RAPTOR: Recursive Abstractive Processing for Tree-Organized Retrieval
В статье разработана новая архитектура RAG на основе итерационного процесса кластеризации фрагментов текстов БД и их суммаризации. Архитектура позволяет сети отвечать на тематические запросы, требующие суммаризованного контекста всего документа (пример - краткий пересказ книги).
Source: Arxive
Все подходы RAG отлично работают на практике. Тем не менее, у них есть и недостатки. Один из них заключается в том, что большинство существующих методов извлекают только несколько коротких, непрерывных фрагментов текста, что ограничивает их способность представлять и использовать крупномасштабную структуру одного документа или даже нескольких. Это особенно актуально для тематических вопросов, требующих интеграции знаний из нескольких частей текста, таких как понимание целой книги.
Если вкратце перефразировать проблему - чем больше фрагментов текста вы включаете в запрос, тем меньше вам нужен RAG. Ведь по сути вы можете вместе с запросом отправлять в LLM целые книги (контекстные окна сегодняшних моделей позволяют это делать). Например top-K фрагментов текста не смогут дать полного ответа на запрос "Через что прошли дети капитана Гранта, чтобы найти своего отца?".
Также они не смогут полно ответить на сравнительные вопросы типа "Как?": "Как ты думаешь, идти мне на Data Fest или идти на свидание?" из-за большого количества релевантных фрагментов текста. В целом для любой системы RAG найти несколько конкретных частей текста в большом документе является сложной задачей.
Вместо разделения документов на маленькие фрагменты и сохранения их в векторную БД для последующего извлечения, RAPTOR сначала их кластеризует, а после суммаризует каждый кластер с помощью LLM. Он повторяет этот процесс итерационно, пока не остается один, финальный фрагмент текста, в котором содержится вся информация документа.
Каждая нода содержит следующую информацию:
Именно такая архитектура "от общего к частному" отлично работает на любых запросах - краткое содержание книги, сравнение двух книг или даже конкретные факты из обоих документов с сравнением. Все это извлекается в общих чертах с готовой суммаризированной информацией, а если необходимы факты, то можно опуститься на слой ниже и извлечь более детальное summary.
Авторы используют мягкую кластеризацию - то есть узлы могут принадлежать к нескольким кластерам одновременно, не требуя фиксированного числа кластеров. Такая особенность необходима, поскольку отдельные фрагменты текста часто содержат информацию, относящуюся к различным темам, что оправдывает их включение в несколько кластеров.
Алгоритм кластеризации основан на Gaussian Mixture Models (GMMs) - такой подход обеспечивает необходимую гибкость и вероятностную структуру. GMM - это вероятностная модель, которая предполагает, что все точки данных генерируются из смеси конечного числа гауссовских распределений с неизвестными параметрами.
В качестве финального распределения вероятности (к каким кластерам принадлежит вектор текста) авторы используют взвешенные Гауссовы распределения.
Первая формула - условная вероятность принадлежности вектора x (эмбеддинга текста некоторой размерности) к некоторому кластеру k. Остальные параметры - параметры гауссого распределения N. Финальное распределение - взвешенная сумма всех распределений. Параметры π определяют принадлежность каждого распределения к кластеру.
Однако высокая размерность эмбеддингов представляет собой проблему для традиционных GMM, поскольку метрики расстояний могут плохо себя вести при измерении сходства в высокоразмерных пространствах. Поэтому авторы используют Uniform Manifold Approximation and Projection (UMAP) для уменьшения размерности. Варьируя параметр количества соседей в кластерах, они понижают размерность эмбеддингов в 2 этапа:
- Сначала определяются глобальные кластеры
- Затем выполняется локальная кластеризация внутри этих глобальных кластеров
Такой подход позволяет захватить как общие темы, так и конкретные детали.
Оптимальное количество кластеров определяется с помощью Bayesian Information Criterion (BIC) - критерий, который сильно штрафует модель за увеличение количества параметров и вознаградает за уменьшение.
Здесь N - количество фрагментов текста, k - параметры модели а L - максимизированное значение функции правдоподобия модели (likelihood function). Благодаря этому критерию, у нас не возникнет ситуации равенства количества фрагментов текста и кластеров.
Для суммаризации текста авторы использовали gpt-3.5-turbo, однако около 4% summary содержали незначительные галлюцинации. Они не распространялись
на родительские узлы и не оказали заметного влияния на решение задач, связанных с ответами на вопросы.
RAPTOR предоставляет 2 способа запроса - traversal и collapsed tree.
Traversal tree - сначала выбирается top-k наиболее релевантных корневых узлов на основе их косинусного сходства с запросом. На следующем уровне рассматриваются дочерние узлы этих выбранных узлов, и из этого пула снова выбираются top-k узлов на основе их косинусного сходства с вектором запроса. Этот процесс повторяется до тех пор, пока мы не достигнем узлов листа.
Collapsed tree - более простой способ поиска релевантной информации за счет одновременного рассмотрения всех узлов дерева. Вместо того чтобы переходить от слоя к слою, этот метод сглаживает многослойное дерево в один слой, выводя все узлы на один уровень для сравнения.
RAPTOR превосходит базовые показатели каждого из соответствующих методов поиска (SBERT, BM25, DPR) на наборе данных NarrativeQA, используя в качестве языковой модели UnifiedQA-3B.
Также он превосходит поисковики на датасетах QuALITY и QASPER.
Сравнение оценок F-1 на наборе данных QASPER с использованием трех различных языковых моделей (GPT-3, GPT-4, UnifiedQA 3B) и различных методов поиска. «Title + Abstract» отражает производительность, когда для контекста используются только название и аннотация статей.