Atlas: совместное преобучение ретривера и ридера превосходит LLM с 540 млрд параметров, используя лишь 11 млрд

Atlas — это продолжение работы Изакара и Грава над их собственной статьей о Fusion-in-Decoder. В Atlas FiD расширена до полностью совместно обучаемой системы, где ретривер и ридер обучаются вместе с нуля. Я изучаю эту работу сейчас, потому что она завершает архитектурную преемственность от оригинальной статьи о RAG через FiD к совместно обучаемому поиску — это именно то пространство решений, в котором приходится ориентироваться любой QA-системе для бухгалтерских журналов.

Статья

В работе «Atlas: Few-shot Learning with Retrieval Augmented Language Models» (Izacard et al., JMLR 2023) ставится вопрос: могут ли модели с дополненной выборкой сравниться с гигантскими LLM в знаниеемких задачах при обучении по нескольким примерам (few-shot). Основной вклад — это тщательно преобученная система с дополненным поиском, которая совместно обучает плотный ретривер на базе Contriever вместе с ридером Fusion-in-Decoder на базе T5. Ключевая идея в том, что именно совместное преобучение, а не архитектура, определяет производительность при работе со знаниями в режиме few-shot. Система извлекает топ-20 документов, кодирует каждый независимо в энкодере, а затем объединяет их в перекрестном внимании (cross-attention) декодера — та же конструкция FiD из статьи авторов 2021 года.

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

• Atlas-11B достигает точности 42,4% на Natural Questions всего с 64 обучающими примерами, превосходя PaLM (540 млрд параметров) примерно на 3 пункта, при этом используя в 50 раз меньше параметров.
• На TriviaQA (64 примера) Atlas-11B достигает 74,5% на отфильтрованном наборе и 84,7% на неотфильтрованном скрытом тесте, что показывает, как компонент поиска эффективно компенсирует ограниченный контроль за выполнением задач.
• Оцениваются четыре цели обучения ретривера: дистилляция внимания (ADist), EMDR2 (трактовка извлеченных документов как латентных переменных), дистилляция перплексии (PDist) и LOOP (leave-one-out). Различия в производительности между ними невелики; PDist выбрана из-за вычислительной эффективности.
• Совместное преобучение на неразмеченном тексте является решающим фактором: все конфигурации с преобучением поиска значительно превосходят базовую модель, где поиск настраивался только на этапе fine-tuning.
• Индекс документов можно обновлять после обучения без переучивания модели, что архитектурно важно для динамических баз знаний. Временное несоответствие индексов заметно снижает производительность.
• На MMLU (5 примеров) Atlas-11B достигает 47,9%, превышая заявленные 43,9% у GPT-3, несмотря на то, что параметров у него примерно в 16 раз меньше.

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

Главное утверждение — что поиск обеспечивает высокую производительность в задачах на знание при малом количестве примеров и лишь части параметров — подтверждается убедительно. Результат в 42,4% на NQ с 64 примерами поразителен, а сравнение с PaLM справедливо, так как на тот момент PaLM был эталоном масштабируемости.

Однако у меня есть три замечания. Во-первых, точность поиска не идеальна даже после совместного обучения: независимый анализ показывает, что Contriever упускает хотя бы одно эталонное утверждение примерно в 85% случаев и достигает около 47% точности поиска для QA. Совместное обучение улучшает поиск по сравнению с несовместными моделями, но ридер проделывает огромную работу, компенсируя несовершенство выборки — громкие цифры few-shot отражают потолок всей системы, а не качество компонента поиска. Во-вторых, инфраструктурные затраты реальны: обновление индексов документов во время преобучения добавляет примерно 30% вычислительных накладных расходов, а полный индекс Wikipedia+CommonCrawl требует 587 ГБ в формате fp16. Это приемлемо для исследовательских целей, но является серьезным операционным ограничением для внедрения в продакшн. В-третьих, проблема утечки данных признается, но не решается: 2,8% вопросов MMLU встречаются дословно в корпусе CCNet, использованном для преобучения, что завышает результаты MMLU на неизвестную величину.

Существует также более тонкое архитектурное ограничение, которое в статье не рассматривается в полной мере: FiD кодирует каждый извлеченный фрагмент независимо перед объединением. Это помогает параллелизму, но означает, что в энкодере отсутствует внимание между фрагментами. Длинные цепочки рассуждений, требующие соединения информации из разных фрагментов, должны полностью обрабатываться в декодере — а при 20 извлеченных фрагментах на перекрестное внимание декодера ложится огромная нагрузка.

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

Для QA-систем бухгалтерских журналов Beancount наиболее актуальным вкладом Atlas является эмпирическая демонстрация того, что совместное обучение ретривера и ридера оправдывает себя в условиях малого количества примеров, а также честный анализ того, когда оно не работает. Агент Beancount, запрашивающий историю транзакций за несколько лет, сталкивается именно с проблемой динамического индекса: новые записи появляются ежедневно, и индекс месячной давности будет давать неверные ответы. Atlas показывает, что индекс можно заменить «на горячую» без переучивания, что обнадеживает с точки зрения архитектуры.

Однако цифры точности поиска отрезвляют. Если Contriever упускает нужную запись в журнале в 53% попыток поиска даже после совместного обучения на общих текстах, то финансовому агенту, работающему с журналами Beancount — с их специфическими названиями товаров, иерархиями счетов и директивами bean — потребуется либо адаптивное обучение ретривера под домен, либо поиск, дополненный методами структурированных запросов (точное соответствие счетов, фильтрация по датам). Одной лишь выборки в стиле RAG, даже при совместном обучении, будет недостаточно для высокоточных операций с бухгалтерскими книгами.

Сравнение с PaLM также проясняет архитектурный компромисс: поиск позволяет сжать знания в меньшее количество параметров, снижая стоимость инференса. Для такого продукта, как Beancount.io, где стоимость инференса имеет значение при масштабировании, философия дизайна Atlas привлекательна. Но стоимость индекса в 587 ГБ перекладывает нагрузку на инфраструктуру хранения и поиска — это иного рода операционное ограничение, которое не отображается в бенчмарках.

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

• REALM: Retrieval-Augmented Language Model Pre-Training (Guu et al., arXiv:2002.08909, ICML 2020) — более ранняя платформа совместного преобучения ретривера и ридера, которую расширяет Atlas; необходима для понимания того, что именно Atlas улучшил, а что оставил без изменений.
• RA-DIT: Retrieval-Augmented Dual Instruction Tuning (Lin et al., arXiv:2310.01352, ICLR 2024) — достигает производительности, сопоставимой с Atlas, используя настройку инструкций (instruction tuning), а не совместное преобучение с нуля; предполагает, что разрыв между совместным и независимым обучением можно сократить без затрат на инфраструктуру.
• RETRO: Improving Language Models by Retrieving from Trillions of Tokens (Borgeaud et al., arXiv:2112.04426, ICML 2022) — подход DeepMind к поиску во время преобучения на другом масштабе; дополняет общую картину подходов к преобучению с дополненной выборкой перед выбором архитектуры для QA бухгалтерских журналов.