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

Atlas — це продовження роботи Ізакарда та Грейва після їхньої власної статті про Fusion-in-Decoder, що розширює 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 документів, кодує кожен незалежно в енкодері, а потім об'єднує їх у крос-аттеншн декодера, використовуючи той самий дизайн FiD із роботи авторів 2021 року.

Ключові ідеї

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

Що витримує перевірку часом — а що ні

Головне твердження — що пошук забезпечує результативність у задачах few-shot при набагато меншій кількості параметрів — залишається переконливим. Результат 42,4% на NQ з 64 прикладами вражає, а порівняння з PaLM є справедливим, оскільки PaLM був еталоном масштабу на той час.

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

Існує також більш тонка архітектурна обмеженість, яку стаття не розкриває повністю: FiD кодує кожен отриманий уривок незалежно перед об'єднанням, що допомагає паралелізму, але означає, що енкодер не має крос-аттеншн між уривками. Довгі ланцюжки міркувань (multi-hop), які потребують з'єднання інформації з різних уривків, мають виконувати всю цю роботу в декодері — і при 20 отриманих уривках крос-аттеншн декодера несе велике навантаження.

Чому це важливо для фінансового ШІ

Для системи QA реєстрів Beancount найбільш релевантним внеском Atlas є емпірична демонстрація того, що спільне навчання ретривера та рідера приносить результати в умовах few-shot — та чесний звіт про те, коли це не працює. Агент 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, за допомогою інструктивного донавчання замість спільного попереднього навчання з нуля; припускає, що розрив між спільним та незалежним навчанням можна подолати без витрат на інфраструктуру.
• RETRO: Improving Language Models by Retrieving from Trillions of Tokens (Borgeaud et al., arXiv:2112.04426, ICML 2022) — підхід DeepMind до пошуку під час попереднього навчання на іншому масштабі; доповнює картину підходів до попереднього навчання з доповненим пошуком перед вибором архітектури для QA фінансових реєстрів.