NeoMME: Легкий мультимодальный энкодер для поиска по изображениям

Команда Hcompany представила NeoMME — семейство мультимодальных энкодеров (260M и 800M параметров), спроектированных с нуля для эффективного поиска по визуальным документам. В отличие от популярных VLM-подходов, NeoMME не использует отдельный предобученный визуальный башню (vision tower) или причинный декодер. Вместо этого единый двунаправленный Transformer обрабатывает и текстовые токены, и сырые патчи изображений, обучаясь на задаче маскированной дискретной диффузии.

Модель демонстрирует выдающуюся эффективность: на GPU NVIDIA L40S при разрешении 2048×2048 пикселей модель NeoMME-260M кодирует около 51 страницы в секунду, что в два раза быстрее ColModernVBERT. При этом она достигает состояния Парето по соотношению размера модели и качества поиска (nDCG@10) на датасете ViDoRe v3.

01Архитектура: Один трансформер для всего

Ключевое отличие NeoMME от адаптаций ModernBERT или SigLIP2 заключается в нативной мультимодальности. Текст и изображения проходят через один и тот же вычислительный путь, что упрощает параллелизацию и обслуживание (serving).

• Входные данные: Текст использует факторизованные эмбеддинги, изображения разбиваются на непересекающиеся патчи 32×32 и проецируются через MLP. Оба типа данных попадают в один Transformer Encoder.
• Контекст: Длина контекста составляет 16,384 токенов, что позволяет обрабатывать изображения высокого разрешения (до двух кадров 4K UHD) без потери деталей.
• Механизмы внимания: Используются grouped-query attention (GQA), нормализация query-key, gated attention, 2D RoPE и squared-ReLU MLP. Слои чередуются: большинство используют скользящее окно (sliding window), каждые шестой и последний слои — глобальное внимание.
• Токенизатор: BPE-токенизатор с словарем 131k токенов, обученный на мультиязычном тексте, коде, математике и транскриптах изображений.

02Обучение: Маскированная диффузия

Модель обучалась с нуля на 524 миллиардах упакованных токенов. Специфика обучения заключается в объективе маскированной текстовой денормализации:

1. Для текстовых примеров маска применяется с вероятностью от 0 до 1.
2. Для мультимодальных примеров (текст + изображение) маска текста составляет 0.3–1.0. Патчи изображения остаются видимыми, а модель должна восстановить замаскированный текст, опираясь на визуальный контекст.

Высокий уровень маскировки (high masking) лишает модель возможности использовать только языковые подсказки, заставляя её учиться привязывать текстовые концепции к визуальным патчам. Для оптимизации использовался оптимизатор NorMuon, что позволило достичь хороших результатов при относительно небольшом бюджете токенов (290B текстовых токенов против 2T у ModernBERT).

03NeoMME-Retriever: Плотные и поздние взаимодействия

Для задач визуального поиска документов (Visual Document Retrieval) базовая модель дообучается с добавлением двух голов (heads):

• Dense Head: Среднее пулирование скрытых состояний дает один нормализованный вектор. Идеально для быстрого поиска по ANN-индексам.
• Late-Interaction Head: Проецирует каждый токен запроса и патч изображения в 128-мерный вектор. Это сохраняет локальные совпадения (как в ColBERT), позволяя точно ранжировать документы по layout, таблицам и шрифтам, которые теряются при OCR.

Один forward pass возвращает оба типа эмбеддингов. Рекомендуется использовать late-interaction для высокой точности (с библиотеками типа NextPlaid) или dense-эмбеддинги для предварительной фильтрации больших корпусов с последующим reranking.

04Бенчмарки и производительность

Модели NeoMME-Retriever показывают выдающиеся результаты на датасетах ViDoRe, занимая лидирующие позиции по соотношению «качество/размер».

NeoMME-260M превосходит ColModernVBERT почти вдвое по метрике nDCG@10 при схожем размере. NeoMME-800M приближается к ColQwen2.5 (3.75B), используя в 4.5 раза меньше параметров.

05Установка и запуск

Модели доступны в Hugging Face Transformers под лицензией Apache 2.0. Для локального запуска на GPU с поддержкой CUDA (например, NVIDIA L40S, A100, RTX 4090) используйте следующий код:

Установка зависимостей:

pip install transformers torch accelerate

Пример загрузки и инференса для NeoMME-Retriever:

from transformers import AutoModel, AutoTokenizer

# Загрузка модели и токенизатора
model_name = "Hcompany/NeoMME-Retriever-800M"
model = AutoModel.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Пример текста и изображения (предполагается наличие пути к файлу)
# В реальном сценарии изображение нужно предварительно обработать
# в соответствии с требованиями токенизатора (патчи 32x32)
inputs = tokenizer(
    text="What is the main chart about?",
    images=["path/to/document_page.png"],
    return_tensors="pt"
)

# Forward pass
with torch.no_grad():
    outputs = model(**inputs)
    # outputs содержит dense и late-interaction embeddings
    dense_embedding = outputs.dense_embeddings
    late_interaction_embeddings = outputs.late_interaction_embeddings

06Кому подойдёт / что запустится

• Visual RAG: Идеально для систем, где нужно искать информацию по сканам документов, чертежам или фотографиям без предварительного OCR. Сохраняет layout и визуальные подсказки.
• Low-latency Search: Модель 260M может работать на менее мощных GPU (даже с 8-12 ГБ VRAM) с высокой пропускной способностью (51 стр/с на L40S). Подходит для сервисов с высокими требованиями к скорости ответа.
• Cost-sensitive Infra: Если вы хотите избежать затрат на запуск больших VLM (3B+ параметров), NeoMME-800M дает сопоставимое качество при меньшем потреблении ресурсов.
• Железо: Требуется GPU с поддержкой PyTorch и CUDA. Для late-interaction поиска с большими корпусами рекомендуется 16+ ГБ VRAM, для dense-поиска достаточно 8 ГБ.

Источник: Hugging Face ↗