Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks

Разбор статьи

Проблема и мотивация

Предварительно обученные языковые модели хранят значительные объемы фактологических знаний внутри своих параметров (параметрическая память). Однако у них есть существенные недостатки: они с трудом расширяют или обновляют свои знания без повторного обучения, не могут легко предоставить источники или интерпретируемость своих выводов и склонны к галлюцинациям.

Хотя гибридные параметрические/непараметрические модели, такие как REALM и ORQA, объединяли дифференцируемый поиск с маскированным языковым моделированием, они были ограничены экстрактивным вопросно-ответным поиском. Стандартные архитектуры sequence-to-sequence (seq2seq), такие как BART и T5, оставались преимущественно «закрытыми книгами» (closed-book). Данная работа устраняет этот пробел, предлагая Retrieval-Augmented Generation (RAG) — общий фреймворк тонкой настройки, который объединяет предварительно обученный плотный нейросетевой ретривер с авторегрессионным генератором seq2seq как для генерации в открытой предметной области, так и для классификации.

Ключевые идеи

• Гибридная параметрическая и непараметрическая память: непараметрическая память состоит из плотного векторного индекса по 21 млн 100-словных фрагментов Wikipedia, доступ к которому осуществляется через Maximum Inner Product Search (MIPS). Параметрическая память представляет собой предварительно обученный генератор BART.
• Сквозная тонкая настройка (End-to-End Fine-Tuning): кодировщик запросов и генератор обучаются совместно, рассматривая извлеченные фрагменты как скрытые переменные и минимизируя отрицательное маргинальное логарифмическое правдоподобие целевых токенов без явного контроля со стороны разметки ретривера.
• Две формулировки кондиционирования:
  • RAG-Sequence: обусловливает всю генерируемую последовательность одним и тем же извлеченным документом.
  • RAG-Token: позволяет каждому сгенерированному токену обращать внимание и маргинализировать по различным извлеченным фрагментам.
• Горячая замена знаний о мире: внешний индекс можно заменить во время инференса (например, заменив индекс Wikipedia 2016 года на индекс 2018 года), чтобы обновить знания о мире без изменения весов модели.

Метод

Для входной последовательности $x$ ретривер $p_{\eta }(z|x)$ формирует распределение по документам $z$, а генератор $p_{\theta }(y_i|x,z,y_{1:i-1})$ генерирует целевые токены, обусловленные как $x$, так и $z$.

Ретривер использует Dense Passage Retriever (DPR): $$p_{\eta }(z|x)\propto \exp \left(\mathbf{d}(z{)}^{\top }\mathbf{q}(x)\right)$$ где $\mathbf{d}(z)={\text{BERT}}_d(z)$ и $\mathbf{q}(x)={\text{BERT}}_q(x)$. Top-$k$ кандидатов извлекаются с помощью FAISS MIPS. Индекс документов и кодировщик документов ${\text{BERT}}_d$ остаются фиксированными, в то время как ${\text{BERT}}_q$ и параметры генератора BART $\theta$ подвергаются тонкой настройке.

RAG-Sequence

Модель рассматривает извлеченный документ как единую скрытую переменную для всей последовательности $y$: $$p_{\text{RAG-Sequence}}(y|x)\approx \sum _{z\in \text{top-}k(p(\cdot |x))}{p}_{\eta }(z|x)\prod _i^N{p}_{\theta }(y_i|x,z,y_{1:i-1})$$ Декодирование можно оценивать с помощью Thorough Decoding (запуск beam search для каждого документа и вычисление прямых проходов для гипотез-кандидатов по всем документам) или Fast Decoding (аппроксимация отсутствующих вероятностей луча нулем).

RAG-Token

Модель выполняет маргинализацию по документам на каждом шаге генерации токена, что позволяет смешивать контент из нескольких источников на уровне токенов: $$p_{\text{RAG-Token}}(y|x)\approx \prod _i^N\sum _{z\in \text{top-}k(p(\cdot |x))}{p}_{\eta }(z|x)p_{\theta }(y_i|x,z,y_{1:i-1})$$ Это дает стандартную авторегрессионную вероятность перехода, которая напрямую интегрируется в beam search.

Результаты

• Open-Domain QA: RAG устанавливает новые SOTA-результаты на открытых бенчмарках по метрике точного совпадения (exact match, EM):
  • Natural Questions: 44.5 (RAG-Seq) против 41.5 (DPR) и 36.6 (T5-11B+SSM).
  • TriviaQA (стандартный открытый сплит): 56.8 (RAG-Seq); на тестовом сплите TQA-Wiki: 68.0 (RAG-Seq) против 60.5 (T5-11B+SSM).
  • WebQuestions: 45.5 (RAG-Token) против 41.1 (DPR).
  • CuratedTrec: 52.2 (RAG-Seq) против 50.6 (DPR).
  • Примечательно, что на Natural Questions RAG достиг 11.8% EM в случаях, когда правильный ответ полностью отсутствовал в извлеченном контексте (где экстрактивные модели получают 0%).
• Абстрактивная генерация: на Open MS-MARCO RAG-Sequence достигает 40.8 Rouge-L и 44.2 Bleu-1, превосходя BART (38.2 Rouge-L, 41.6 Bleu-1).
• Генерация вопросов Jeopardy: эксперты-оценщики признали RAG-Token более фактологичным, чем BART, в 42.7% случаев, тогда как BART оказался лучше только в 7.1%. Эксперты также сочли RAG более специфичным (37.4% против 16.8%).
• Верификация фактов: на задаче FEVER (3 класса) RAG достигает точности меток 72.5% (с отставанием всего в 4.3% от конвейерных систем, использующих явный контроль по золотым фактам); на 2 классах RAG достигает точности 89.5%.
• Обновление знаний через горячую замену: в темпоральном тесте с 82 мировыми лидерами в период с 2016 по 2018 год RAG показал точность 70% для лидеров 2016 года с индексом 2016 года и 68% для лидеров 2018 года с индексом 2018 года, при этом результаты резко упали (до 12% и 4%) при несовпадении индексов.

Ограничения и открытые вопросы

• Коллапс поиска (Retrieval Collapse): для задач, требующих более слабого явного фактологического обусловливания (таких как генерация историй с открытым финалом), ретривер может вырождаться в выбор статичных, неизменных документов, пока генератор учится их игнорировать.
• Фиксированный индекс документов: кодировщик документов остается статичным во время тонкой настройки задачи из-за высокой вычислительной стоимости постоянной переиндексации 21 млн векторов, что ограничивает адаптацию ретривера.
• Обоснование источников и смещения: непараметрическая память остается уязвимой к дезинформации и предвзятостям, заложенным в базовом корпусе знаний (Wikipedia).
• Совместное предварительное обучение: компоненты ретривера и генератора обучались предварительно по отдельности до тонкой настройки на прикладных задачах, что оставляет их совместное предварительное обучение с нуля открытым направлением для исследований.