ColBERT, MUVERA и скучная идея, которая оказалась выгоднее обеих

Есть особый жанр разочарования, знакомый каждому, кто занимался поиском всерьёз: читаешь статью, где новый метод обходит старый на несколько пунктов, воспроизводишь всё у себя и видишь совсем другую картину. С ColBERT и MUVERA у меня вышло именно так: на обычных текстах обещанного прироста почти не было, зато по дороге выяснилось, где эти методы действительно работают, а где нет.

Добраться до них пришлось почти случайно. Мне попалась статья о MUVERA, реализация которой уже появилась в экосистеме Qdrant. С Qdrant я работаю давно и как раз готовил по нему практический курс, поэтому решил разобраться и с этой частью поиска, тем более что рассказывать о механизме, который сам ни разу не запускал, было бы странно. Первым делом я просто проверил всё на практике.

Дальше началась история о том, почему результаты оказались совсем не такими, как я ожидал. По дороге не обошлось без собственных ошибок, и каждая из них, как ни странно, постепенно приближала к ответу. А сам ответ оказался хоть и скучнее, чем хотелось бы, но от этого, пожалуй, только полезнее.

Красивая идея и её цена

В ColBERT хочется верить, и понятно почему. Обычный плотный поиск обходится с документом грубо: прогоняет через модель и сворачивает в один вектор. Если в документе десять абзацев, а к запросу относится один, вектор описывает документ «в среднем», и нужный абзац в этом среднем растворяется. ColBERT [1] предлагает не сворачивать. Каждый токен документа получает собственный вектор, а при поиске для каждого токена запроса находится самый похожий токен документа. Такой подход называют поздним взаимодействием, late interaction. Усреднения нет — значит, растворяться нечему.

Беда в цене. Плотный вектор документа занимает около четырёх килобайт, матрица ColBERT на документ средней длины — около ста пятидесяти. На миллионе документов это разница между четырьмя гигабайтами и ста сорока. В качестве компромисса предложили MUVERA [2]: она сжимает матрицу в один длинный вектор фиксированного размера, по нему быстро отбирает кандидатов, а точный подсчёт делается только для них. Качество от одной технологии, цена от другой — выглядит как идеальная пара.

Для проверки у меня было четыре размеченных набора: русская и английская Википедия из MIRACL [3], кодексы РФ [4] и научные аннотации SciFact [5], около трёх тысяч запросов. Соперником был обычный плотный поиск на модели multilingual‑e5-large. Чтобы меня не обвинили в сравнении гиганта с малышом, первым делом пришлось сверить размеры: e5-large и jina‑colbert‑v2 — это 560 и 559 миллионов параметров, почти близнецы.

Метрика, которая чуть не соврала

Первый результат был простым и обидным: ColBERT не находил больше документов. Ни на одном наборе recall@10, то есть доля запросов, для которых нужный документ попадал в первую десятку, заметно не вырос. Вывод «не окупается» напрашивался сам, и до отчёта ему оставался один шаг.

Меня остановило одно число. На кодексах плотный поиск и так находил нужный документ в первой десятке в 93.7% случаев. Когда базовая линия стоит так высоко, выигрывать негде, и отсутствие разницы говорит уже не о равенстве методов, а о том, что метрика перестала что‑либо различать. Пришлось пересчитать всё по recall@1, то есть по доле запросов, где нужный документ оказался на первом месте. Здесь плотный поиск угадывал лишь в 70.5% случаев, и прибавка ColBERT выросла с 0.3 пункта до статистически значимых 4.7. Эффект, который едва не остался незамеченным, оказался в пятнадцать раз больше, чем выглядел.

Та же метрика прятала и обратное. На русской Википедии, пересчитанной по первым позициям, ColBERT выдачу значимо ухудшил. Итог по четырём наборам вышел пёстрым: один домен выиграл, один проиграл, два почти не заметили разницы. Технология, которая по замыслу должна помогать везде, помогала в одном случае из четырёх. Оставалось понять почему.

Две причины вместо одной

Первая подсказка нашлась в анонсе самой модели: Jina пишет, что их ColBERT значительно превосходит BM25 — подсчёт совпадающих слов без всяких нейросетей — на всех языках MIRACL [6]. Почти все публичные сравнения late interaction устроены так же: соперником выступает BM25 или плотная модель первого поколения, а не e5, которую обучали примерно на миллиарде пар текстов. К тому же ColBERTv2, на который чаще всего ссылаются, — это не только позднее взаимодействие, но и особая схема обучения: дистилляция из кросс‑энкодера с трудными негативами [7]. А обычная одновекторная модель, дистиллированная из ColBERT‑учителя, забирает бóльшую часть его выигрыша [8]. Часть заслуг, которые приписывали архитектуре, принадлежала способу обучения, а его современные плотные модели давно впитали. Late interaction был отличным способом выжать качество из слабых энкодеров, но энкодеры перестали быть слабыми.

Вторая причина прозаичнее. На английском SciFact мультиязычная jina‑colbert‑v2 проиграла английской answerai‑colbert‑small, которая меньше её в семнадцать раз: recall@1 у маленькой модели — 64.7%, у большой — 58.3%. И только маленькая модель значимо улучшила плотный поиск. Похоже, модель, рассчитанная на десятки языков, на отдельном языке уступает модели, обученной только на нём. А почти все мои замеры сделаны как раз мультиязычной jina‑colbert‑v2, так что ColBERT в них, скорее всего, выглядит слабее, чем выглядел бы с моделью под конкретный язык.

Гипотеза, которая была слишком хороша

С MUVERA вышло поучительнее, потому что здесь у меня было объяснение, выглядевшее неизбежным, — и оно оказалось неверным.

MUVERA в моих экспериментах показывала результат ниже, чем связка с обычным плотным поиском. Здесь важно понимать устройство двухэтапного поиска: сначала быстрый поиск по всему корпусу отбирает кандидатов — это делает либо обычный плотный поиск, либо MUVERA, — а уже потом ColBERT подробно пересчитывает только их. Поэтому я проверял, что происходит при разном количестве кандидатов.

На кодексах РФ, где 3786 документов, первый этап отбирал 100 кандидатов. Когда их отбирал плотный поиск, связка с ColBERT давала MRR 0.823 (MRR — метрика того, насколько высоко в выдаче стоит правильный ответ), а когда MUVERA — 0.737. Для сравнения, обычный плотный поиск без второго этапа давал 0.792.

Увеличение списка кандидатов улучшало результат MUVERA: при 500 кандидатах MRR составлял 0.800, при 1000 кандидатах 0.811. Однако даже при таком увеличении числа кандидатов результат оставался ниже, чем у связки, где кандидатов отбирал плотный поиск.

Объяснение, как мне показалось, нашлось в формуле построения сжатого вектора. MUVERA раскладывает векторы токенов документа по корзинам — делит пространство случайными гиперплоскостями — и внутри каждой корзины их усредняет. Усредняет — то есть делает ровно то, от чего ColBERT был призван избавить. При стандартных настройках корзин тридцать две, на кодексах в каждую попадало в среднем шестнадцать токенов, и нужное слово, казалось, тонуло среди соседей.

Если виновато усреднение, лечение очевидно: сделать корзины мельче, чтобы в каждую попадало меньше токенов. Число корзин в MUVERA задаётся настройкой, и его можно менять, не трогая общий размер вектора. Я прошёл от восьми корзин до ста двадцати восьми. При восьми в каждую попадало под сорок токенов, при ста двадцати восьми — около шести, то есть усреднение слабело в шесть раз. А качество не двигалось: 0.744 на восьми корзинах, 0.750 на шестнадцати, 0.740 на тридцати двух, 0.730 на шестидесяти четырёх. Упало оно только на ста двадцати восьми — и не потому, что усреднение вернулось, а потому, что корзин стало так много, что большинство из них пустовало. Красивая гипотеза не пережила первой же проверки.

Осталось объяснение попроще. Сжатый вектор MUVERA собран из случайных проекций: в нём нет ничего обученного, он не знает, что такое релевантность. Плотный вектор, наоборот, целиком продукт обучения на задаче поиска. Проверил я это в самой жёсткой форме: выставил против сорокакилобайтного вектора MUVERA маленькую модель e5-small, чей вектор весит полтора килобайта. На кодексах маленькая обученная модель отбирала кандидатов лучше при любом размере пула, и тем заметнее, чем пул меньше: при 25 кандидатах recall@1 после пересчёта ColBERT был 74.7% против 61.8% у MUVERA, а при 500 разница сократилась до 1.4 пункта. Вектор в двадцать шесть раз меньше выигрывал, потому что был обучен.

Справедливости ради, авторы MUVERA ничего подобного не обещали: они сравнивали её с поиском по токенным векторам самого ColBERT, а не с отдельной плотной моделью, и вопрос практика — ставить ли MUVERA вместо уже имеющегося плотного поиска — у них просто не задан. Зато у неё есть ловушка: качество зависит от доли корпуса, а не от числа кандидатов. Сто кандидатов — это пять процентов корпуса в две тысячи документов, но лишь одна десятая процента в корпусе на сто двадцать пять тысяч, где нужный документ у меня попадал в первую десятку лишь в 14% случаев. И ничто вокруг не сигналит о поломке.

Скучная идея

Самый важный эксперимент оказался последним. Всё это время ColBERT у меня соревновался с плотным поиском, где на документ приходится один вектор, — а это нечестно по отношению к плотному поиску: ColBERT получает в тридцать семь раз больше места в индексе. Что если отдать плотному поиску тот же бюджет? У него есть давний ответ — нарезать документ на куски и хранить вектор для каждого.

Проверить это удобнее всего было на русской Википедии из MIRACL: её я ещё раньше проиндексировал в двух видах — 7341 статья целиком и те же статьи, нарезанные на 125 тысяч фрагментов, в среднем по семнадцать на статью. Чтобы сравнение было честным, я оценивал все режимы по статьям: если поиск находил фрагмент, засчитывалась статья, из которой он взят, а повторные фрагменты той же статьи в выдаче пропускались. Так каждый режим отвечал на один и тот же вопрос: в какой статье ответ.

Нарезка стоила в семнадцать раз больше места и принесла почти пять пунктов. ColBERT стоил сверху ещё в двадцать один раз больше — и не добавил ничего: 85.3% против 85.3%, до десятой доли. Оговорюсь честно: этот замер сделан мультиязычной моделью, и со специализированной картина могла бы сложиться иначе.

Кажется, в этой картинке спрятано объяснение всей истории. Дополнительное место в индексе можно потратить двумя способами. Можно увеличить покрытие — дать документу больше независимых точек, в которых он может совпасть с запросом. Это и делает нарезка: вместо одного усреднённого вектора у статьи появляется семнадцать, и нужный абзац больше не растворяется в остальных. А можно увеличить разрешение — описать ту же самую точку подробнее, сотней векторов вместо одного. Это делает ColBERT. Первое работает и стоит дёшево. Второе, во всяком случае против современной плотной модели, почти не работает. MUVERA в эту схему не укладывается вовсе. Её вектор занимает сорок килобайт, но собран из случайных проекций, а не обучен, и при отборе кандидатов он проигрывал вектору маленькой e5-small, который весит всего полтора килобайта, зато обучен искать. У MUVERA дело не в том, сколько места потрачено, а в том, на что.

Где плотная модель слепнет

Из всего этого напрашивается вывод, что MUVERA нужна лишь там, где хорошей плотной модели нет. Но это рассуждение, а не замер, и мне хотелось найти такое место на деле. Трудного текста для этого мало — трудного текста e5 видела сколько угодно. Нужны были данные, которых она никогда не видела, и выбор пал на байткод Python — то, во что интерпретатор компилирует функцию перед исполнением.

Я взял пять тысяч функций из стандартной библиотеки и скомпилировал каждую отдельно. Запросом служила та же функция после небольшой правки: я переименовал локальные переменные, сдвинул числа и вставил в начало лишнюю строку. Найти нужно было исходную. Функции я хранил сырыми байтами — парами шестнадцатеричных чисел вроде «7c 00 64 01», по‑настоящему чужими для текстовой модели.

На сырых байтах e5 ставила нужную функцию первой в 6% случаев. Это почти слепота: в одном из примеров правильный ответ оказался у неё на 220-м месте, а первой шла функция из модуля sqlite3, не имеющая с ним ничего общего. ColBERT угадывал в 97% случаев: сравнивая токены по одному, он находил совпадающие куски, которые одновекторная модель размазала в среднем. Здесь наконец нашлась и ниша MUVERA. Когда сто кандидатов для ColBERT отбирала e5, выходило 48% — нужная функция просто не попадала в сотню. Когда их отбирала MUVERA — 80%, а при пятистах кандидатах, десятой части корпуса, — 93%. На втором корпусе, функциях из сторонних библиотек вроде torch и transformers, картина повторилась: 4% у e5 против 99.7% у ColBERT.

Радоваться, правда, пришлось недолго. BM25 по тройкам соседних инструкций — подсчёт совпадений без единой нейросети — давал 94–96%. Весь выигрыш ColBERT на по‑настоящему чужих данных свёлся к нескольким пунктам.

Выходит, ниша у late interaction есть, но она уже, чем обещают анонсы: данные, которых плотная модель не видела вовсе и где смысл держится на совпадении мелких кусков. Похожее видели авторы BEIR [9] ещё в 2021 году: на биомедицинском BioASQ плотные модели того времени проседали ниже BM25, а ColBERT держался на его уровне. Какие ещё данные в неё попадают, я пока не проверял. Скорее всего, похожая ситуация возникает с данными, которые по своей структуре напоминают байткод: длинные последовательности символов без привычных слов. Это могут быть, например, последовательности генов, химические формулы или артикулы товаров. Но и здесь я бы сначала попробовал простой BM25 по n‑граммам, то есть по коротким последовательностям соседних символов. Для байткода это, например, могут быть тройки инструкций.

Отдельный кандидат здесь — действительно редкие языки, особенно те, для которых в обучающих корпусах почти нет текста. Например, чукотский: на нём говорят всего несколько тысяч человек, а корпус доступных текстов несопоставимо меньше, чем для русского или английского. Поэтому вполне возможно, что e5 практически не встречала его в исходном виде.

Здесь возникает та же проблема, что и с байткодом. Одновекторная модель пытается представить редкие слова и их сочетания в общем семантическом пространстве, хотя у неё может просто не быть достаточно данных, чтобы научиться их различать. ColBERT в этом смысле устроен иначе: он сохраняет представления отдельных токенов и при поиске может находить совпадения между конкретными фрагментами текста.

Особенно показательный случай возникает, если токенизатор плохо работает с письменностью языка. Тогда разные слова могут распадаться на практически одинаковые последовательности неизвестных или редких токенов. Символьные n‑граммы от этой проблемы не зависят: для них неважно, встречалась ли конкретная буква или слово в обучающем корпусе. Они просто видят последовательность символов. Если же говорить о совсем редких языках, например об иероглифах с наручного компьютера Хищника, то тут всё ещё хуже: в Юникоде их нет, и сначала придётся договориться, как их вообще хранить. Но если договоритесь, моя ставка — на n‑граммы по символам: им всё равно, что это за закорючки, лишь бы совпадали.

Что с этим делать

Я не считаю ColBERT плохой технологией. В моих экспериментах он дал заметный выигрыш в трёх случаях: на английских научных текстах, когда модель была английской, а не мультиязычной; на юридических текстах с точной терминологией, причём даже в мультиязычном варианте; и на байткоде, которого плотная модель до этого не видела. В последнем случае разница была особенно большой. Но этот выигрыш легко приписать самой архитектуре, хотя на практике он сильно зависит от того, на каком языке и на каких данных обучена модель.

Если вам предстоит решить, нужен ли late interaction, я бы начинал с нарезки документов. Это дёшево и часто уже даёт заметный эффект. Затем стоит смотреть качество на первой позиции, а не только на десятой: иначе часть разницы между методами просто теряется в метрике. И только если после этого остаётся проблема, а под рукой есть модель, обученная на нужном языке и домене, имеет смысл пробовать ColBERT. При этом стоит помнить, что его индекс в десятки раз больше плотного и растёт вместе с корпусом.

MUVERA, на мой взгляд, интересна там, где обычная плотная модель слепнет: на данных, сильно непохожих на те, на которых она обучалась. В моих экспериментах это оказался байткод. Но даже там сначала стоит попробовать BM25 по n‑граммам: на байткоде он уступил ColBERT всего несколько пунктов и почти ничего не стоит.

Собственно, так и получилась эта статья: хотелось разобраться с одной новой для меня возможностью Qdrant, а в итоге пришлось копать гораздо глубже и разбираться уже не только с API, но и с тем, как вообще устроен поиск. Такой же практический подход я стараюсь использовать и в своём курсе по Qdrant. Подробнее о нём и о том, что там разбираем, есть в описании профиля.

Источники

1. O. Khattab, M. Zaharia. ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT. SIGIR 2020.

2. L. Dhulipala, M. Hadian, R. Jayaram и др. MUVERA: Multi‑Vector Retrieval via Fixed Dimensional Encodings. NeurIPS 2024.

3. X. Zhang, N. Thakur, O. Ogundepo и др. Making a MIRACL: Multilingual Information Retrieval Across a Continuum of Languages. 2022.

4. Набор данных alekseevpavel04/ru‑law‑retrieval, Hugging Face.

5. D. Wadden, S. Lin, K. Lo и др. Fact or Fiction: Verifying Scientific Claims. EMNLP 2020.

6. Jina AI. Jina ColBERT v2: Multilingual Late Interaction Retriever for Embedding and Reranking. 30.08.2024. Статья о модели: R. Jha, B. Wang, M. Günther и др. Jina‑ColBERT‑v2: A General‑Purpose Multilingual Late Interaction Retriever. 2024.

7. K. Santhanam, O. Khattab, J. Saad‑Falcon и др. ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction. NAACL 2022.

8. S.‑C. Lin, J.‑H. Yang, J. Lin. Distilling Dense Representations for Ranking using Tightly‑Coupled Teachers. 2020.

9. N. Thakur, N. Reimers, A. Rücklé, A. Srivastava, I. Gurevych. BEIR: A Heterogeneous Benchmark for Zero‑shot Evaluation of Information Retrieval Models. NeurIPS 2021, Datasets and Benchmarks.