UltraBERT умещает 32K токенов в один вектор. Конец разбиения на чанки?

Источник: Qdrant•

UltraBERT умещает 32K токенов в один вектор. Конец разбиения на чанки?

Каждый пайплайн извлечения данных начинается с одного вопроса: какого размера должен быть чанк? Слишком маленький чанк теряет контекст. Слишком большой — модель эмбеддингов обрезает. UltraBERT-Embed-150M от лаборатории DMIS Lab сдвигает эту границу до 32K токенов. Это плотная модель эмбеддингов…

Каждый пайплайн извлечения данных начинается с одного вопроса: какого размера должен быть чанк? Слишком маленький чанк теряет контекст. Слишком большой — модель эмбеддингов обрезает.

Модель UltraBERT-Embed-150M от лаборатории DMIS Lab сдвигает эту границу до 32K токенов. Это плотная модель эмбеддингов со 151M параметров, созданная на базе UltraBERT, англоязычного энкодера, обученного на входах длиной до 32K токенов. Большинство документов помещаются за один проход, и на выходе вы получаете один 768-мерный вектор.

Так неужели разбиение на чанки умерло? Мы протестировали его в Qdrant: один вектор на документ против чанков по 512 токенов, их комбинацию и оба варианта в сжатом виде.

Как мы проводили тестирование

Мы использовали пять задач из бенчмарка длинных документов LongEmbed. Длина документов в нем составляет в среднем от 7 600 до 13 800 токенов. Каждая задача предполагает свой тип вопросов:

  • SummScreenFD: запрос представляет собой краткое содержание эпизода сериала, а правильным документом является его стенограмма.
  • 2WikiMultihopQA: многошаговые вопросы по фрагментам Википедии.
  • QMSum: запрос суммирует одну из частей стенограммы совещания.
  • Passkey: каждый документ содержит скрытый ключевой пароль среди «воды», а запрос требует этот пароль найти.
  • Needle: каждый документ представляет собой эссе с одним добавленным нерелевантным предложением, и запрос спрашивает об этом предложении.

Для тестирования в масштабе мы также использовали MLDR: 200 000 статей англоязычной Википедии и 800 вопросов, каждый из которых относится к одной статье.

Мы создали эмбеддинги для каждого документа двумя способами:

  • Весь документ целиком: один вектор для полного текста размером до 32K токенов.
  • Чанки по 512 токенов: один вектор на чанк. Оценка документа определяется по его лучшему чанку.

Полные результаты приведены в конце.

Чанки по-прежнему находят детали

Вектор целого документа действительно считывает все 32K токенов. В задаче Passkey он ставил правильный документ на первое место в 99% случаев для текстов любой длины и в 100% случаев для 32K токенов. Но в большинстве других задач чанки набрали больше баллов.

Весь документ или чанки UltraBERT-Embed-150M. nDCG@10 или recall@1 для Passkey и Needle. Больше — лучше. Оценка SummScreen2WikiQMSumPasskeyNeedleMLDR0.00.20.40.60.81.0Весь документЧанки по 512 токенов

Чанки побеждали, когда ответ содержался в нескольких фрагментах: в 2WikiMultihopQA, QMSum и MLDR. Вектор целого документа выиграл только в SummScreenFD, где запрос описывает весь эпизод целиком.

Задача Needle показывает, почему чанки побеждают. Один вектор не может вместить все детали длинного документа. При длине 256 токенов вектор целого документа ставил правильное эссе на первое место в 84% случаев. При длине 32K токенов этот показатель упал всего до 30%. Для той же длины чанки по 512 токенов показали результат в 56%.

Мы также попробовали использовать оба подхода. В Qdrant одна точка может содержать два именованных вектора: вектор целого документа и векторные чанки в виде мультивектора. Гибридный запрос производит поиск по обоим и объединяет результаты. Это немного помогло в SummScreenFD и Needle, улучшив результат по сравнению с одними чанками на 0.02 и 0.03 соответственно. В 2WikiMultihopQA, QMSum и MLDR этот подход сработал хуже, чем одни чанки, сильнее всего отстав на MLDR.

Векторные чанки для 200 000 статей примерно за 140 МБ

Чанки означают большее количество векторов. Статьи из MLDR в среднем разбиваются на 7 чанков каждая, что в сумме дает 1,43 миллиона векторов. В формате float32 это около 4 ГБ.

1-битный TurboQuant уменьшает размер каждого вектора с 3 072 байт до 104. Это в 30 раз меньше и составляет около 140 МБ для всех 200 000 статей. Мы сохранили каждый чанк как отдельную точку Qdrant, а сгруппированный поиск возвращает каждую статью ровно один раз.

По сравнению с точным поиском, сжатые векторы сохранили 95% качества, а с повторным ранжированием (rescoring) — 98%. Несжатые векторы в том же индексе сохранили 99%. Повторное ранжирование имеет свою цену. В нашей установке оригинальные векторы находились на диске. Чтение добавляло от 4 до 11 мс на запрос или более 100 мс, пока дисковый кэш был холодным.

Векторные чанки: в 30 раз меньше На одну статью. Меньше — лучше. 01020Оригинальные векторы несжатые21.4 КБ1-битный TurboQuant без повторного ранжирования0.724 КБ1-битный TurboQuant с повторным ранжированием0.724 КБКБ на статью

Вектор целого документа сжимается хуже, чем векторные чанки. На MLDR он сохранил 89% качества при 1 бите. При 2 битах он показал такой же результат, как и несжатый вектор в том же индексе, поэтому лучше сразу использовать 2 бита или оставить включенным повторное ранжирование.

Так умерло ли разбиение на чанки?

Пока еще нет. В этих тестах чанки по-прежнему лучше находили детали. Мы бы оставили их в качестве отправной точки для поиска деталей, даже когда документ помещается целиком.

UltraBERT дает нам возможность экспериментировать с поиском в разных масштабах. Мы можем осуществлять поиск по целым документам, чтобы найти нужную тему, а затем искать по их чанкам в поисках ответа. Это стоит протестировать, тем более что объединение обоих видов поиска в данном случае не дало стабильного улучшения.

Векторы целых документов по-прежнему обходятся дороже. Создание эмбеддингов для стенограмм SummScreenFD целиком заняло примерно в 3 раза больше времени, чем создание эмбеддингов для чанков по 512 токенов на нашей видеокарте NVIDIA L4. Мы использовали SDPA от PyTorch; рекомендуемый FlashAttention 2, возможно, сократит этот разрыв, но мы еще не замеряли это.

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

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

Полные результаты

Точный поиск. DBSF и RRF — это два способа объединения результатов поиска по целому документу и по чанкам в гибридном запросе. RRF ×1.5 присваивает поиску по чанкам вес 1.5 — это значение мы подобрали для данных задач.

200 000 статей англоязычной Википедии и 800 запросов, каждый из которых относится к одной статье, поиск выполнен с помощью HNSW. При повторном ранжировании используется двукратное избыточное сэмплирование. Размеры учитывают только те векторы, которые ищутся в оперативном памяти. 1 КБ равен 1 024 байтам.

Для сравнения с длинными документами: короткие отрывки, по одному вектору на каждый, на пяти наборах данных BEIR (SciFact, NFCorpus, ArguAna, SCIDOCS и FiQA) с использованием HNSW. Сохраненная оценка приведена относительно точного поиска во float32.

Спасибо лаборатории DMIS Lab за UltraBERT-Embed и за то, что дали нам веский повод переосмыслить подход с разбиением на чанки.

О чём эта статья

Ещё в разделе «Разработка ПО»

Все →

Ещё от Qdrant