Fusion-in-Decoder: как поиск по нескольким отрывкам улучшает генеративные вопросно-ответные системы

Генерация с дополненным поиском (RAG) зависит от того, насколько эффективно генератор может синтезировать данные, распределенные по нескольким документам. В статье Izacard и Grave для EACL 2021 года «Использование поиска отрывков в генеративных моделях для ответов на вопросы в открытой предметной области» предлагается обманчиво простое архитектурное решение — независимое кодирование отрывков и их объединение в декодере — которое значительно превосходит доминировавший тогда фреймворк RAG. Я изучаю её сейчас, потому что принцип проектирования напрямую применим к QA по гроссбухам: прежде чем решать, как извлекать записи в агентах Beancount, стоит понять, какая стратегия слияния (fusion) действительно работает.

О статье

Оригинальный RAG от Льюиса и др. (arXiv:2005.11401) сочетает плотный поисковик (dense retriever) с генератором BART, но заставляет генератор ориентироваться только на один найденный отрывок за раз, усредняя результаты по отрывкам либо для последовательности (RAG-Sequence), либо для токена (RAG-Token). Изакард и Грейв определили это как сдерживающее ограничение: модель, которая видит только один отрывок одновременно, не может легко сопоставить факты, разбросанные по разным документам.

Их решение FiD (Fusion-in-Decoder) элегантно. Каждый найденный отрывок конкатенируется с вопросом, а затем кодируется независимо энкодером T5. Энкодер запускается один раз для каждого отрывка — процесс полностью параллелизуем. Затем декодер выполняет перекрестное внимание (cross-attention) по конкатенации всех представлений отрывков одновременно. Сложность энкодера масштабируется линейно в зависимости от количества отрывков; что критически важно, декодер может обращаться к контексту через границы отрывков на каждом этапе генерации. В статье в качестве основы генератора используются T5-base и T5-large.

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

• FiD-large со 100 найденными отрывками достигает 51,4% точности (exact match) на Natural Questions и 67,6% на TriviaQA open, по сравнению с 47,5% и 56,1% у RAG-Sequence соответственно — прирост составляет примерно 4 и 11 пунктов.
• Производительность на Natural Questions монотонно растет с увеличением количества отрывков: 37,3% при 1 отрывке, 48,8% при 10, 50,8% при 50, 51,4% при 100. Предельная отдача снижается, но никогда не становится отрицательной.
• Показатели TriviaQA улучшаются на 6%, а Natural Questions — на 3,5% при увеличении количества отрывков с 10 до 100. Это доказывает, что декодер действительно агрегирует информацию, а не просто выбирает лучший отрывок.
• Этап кодирования легко параллелизуется: каждая пара (вопрос, отрывок) обрабатывается независимо, поэтому фактическое время выполнения масштабируется сублинейно при наличии соответствующего оборудования.
• FiD-base с 770 млн параметров превосходит T5-11B в режиме «закрытой книги» (44,1% против 36,6% на NQ), демонстрируя, что использование поиска позволяет меньшим моделям показывать результаты значительно выше своего веса.

Что подтверждается, а что — нет

Основной результат надежен и неоднократно подтверждался. Архитектурное решение — независимое кодирование, совместное декодирование — действительно изящно: оно позволяет избежать квадратичного взрыва вычислений self-attention, который возник бы при наивной конкатенации всех отрывков перед энкодером, сохраняя при этом для декодера глобальный контекст всех найденных данных.

Ограничение, которое в статье почти не упоминается, заключается в том, что перекрестное внимание декодера является узким местом во время инференса. Cross-attention должен загружать все пары «ключ-значение» (KV) энкодера для каждого слоя декодера на каждом шаге генерации, и эти тензоры KV растут линейно с количеством отрывков. Исследование 2023 года FiDO (arXiv:2212.08153) показало, что замена многоголового внимания (multi-head attention) на многозапросное (multi-query attention) и прунинг слоев перекрестного внимания обеспечивают 7-кратное ускорение инференса при минимальной потере точности — это подразумевает, что оригинальный декодер FiD избыточен для поставленной задачи.

Также существует проблема калибровки, которую авторы не исследуют: они сообщают о точном совпадении (exact match), что поощряет системы, которые случайно выдают каноническую строку ответа. Для задач синтеза фактов — обобщения выводов по нескольким отрывкам, а не просто извлечения фрагмента — метрика exact match преуменьшает ошибки и преувеличивает уверенность модели. В финансовых сценариях, где неверное число в остальном правильном предложении является критической ошибкой, exact match вообще не подходит.

Почему это важно для финансового ИИ

QA по гроссбухам Beancount по своей природе является задачей поиска по нескольким отрывкам. Вопрос вроде «Сколько я потратил на путешествия в третьем квартале по всем счетам?» требует синтеза десятков записей транзакций с разными датами, счетами и типами валют. Основной вывод FiD — генеративные модели могут агрегировать данные из множества найденных отрывков, и производительность улучшается с увеличением контекста — внушает оптимизм.

Практический вывод для проектирования конкретен: при создании слоя QA для Beancount извлечение большего количества записей-кандидатов (50–100 вместо обычных топ-5) и предоставление генератору совместного доступа к ним, вероятно, будет эффективнее, чем использование переранжирования для выбора одного правильного ответа. Архитектура FiD также четко накладывается на структуру гроссбуха: каждая запись транзакции может быть закодирована независимо (дешево и параллельно), прежде чем декодер синтезирует данные из всех них.

Опасения по поводу стоимости инференса реальны для промышленного использования, но последующая работа над FiDO показывает, что это решается на уровне архитектуры без потери точности. Более насущным ограничением для финансовых агентов является то, что FiD разработан для фактологических ответов с коротким выводом. Анализ гроссбуха часто требует многоэтапной арифметики — сложения сумм, вычисления коэффициентов — а генератор FiD по своей сути не перенаправляет такие задачи интерпретатору. Сочетание слияния в стиле FiD с головой генерации кода в стиле PAL — естественный следующий шаг для обеспечения точности числовых данных.

Что почитать дальше

• FiDO (arXiv:2212.08153, ACL Findings 2023) — многозапросное внимание и прунинг перекрестного внимания позволяют достичь точности FiD при 7-кратном ускорении инференса; необходимо изучить перед внедрением FiD в продакшн.
• REALM: Retrieval-Augmented Language Model Pre-Training (arXiv:2002.08909, ICML 2020) — Гуу и др. показывают, как внедрить поиск на этапе предварительного обучения, а не только при инференсе; дает теоретическую базу, на которой строится FiD.
• Atlas: Few-shot Learning with Retrieval Augmented Language Models (arXiv:2208.03299, JMLR 2023) — расширение FiD от самих Изакарда и др. для условий обучения на малом количестве примеров (few-shot) с совместным обучением поисковика и ридера; наиболее полный синтез этого направления работы.