Хотите наладить разметку данных с помощью LLM или сделать чат-бота с определенной личностью? Хорошей идеей будет показать модели несколько примеров, то есть сделать in-context или few-shot learning. Такой подход позволяет с минимальными затратами, без дополнительного обучения влиять на качество работы языковой модели. Впервые этот эффект был обнаружен при применении GPT-2, которая могла отвечать на вопросы, суммаризировать и переводить тексты без явного дообучения на эти задачи. Это событие положило начало современной эпохе языковых моделей.

Меня зовут Михаил Коновалов, я работаю ML-инженером в Okko. В этой статье я хочу рассмотреть несколько публикаций, посвященных интерпретации механизма Attention в трансформерах.

Рассмотрим пример, как на практике выглядит few-shot learning. На изображении выше модели сначала показали вопрос и ответ, по образцу которых она дает ответ на следующий вопрос. Контекстные примеры «подсказывают» модели, какую задачу нужно решать, а также позволяют задать формат ответа, что может быть важно для автоматических запросов в LLM. Информация из контекста также может служить справочником фактов для LLM, помогая ей отвечать на вопросы по узкоспециализированной теме.

Как именно работают механизмы, которые позволяют LLM так эффективно взаимодействовать с контекстом? Принято рассматривать нейронные сети как black-box, не разбираясь, что на самом деле происходит во время их инференса. Однако можно немного заглянуть «в мозги» нейронным сетям и попытаться понять, за что отвечают те или иные группы параметров модели. В последние годы выходит много статей, посвященных попыткам понять, что происходит внутри языковой модели, при решении задач на основе примеров из контекста. В данном разборе я хочу рассмотреть несколько таких работ.

Чтобы глубже понять эти вещи, стоит вспомнить базовую архитектуру трансформеров, являющихся де-факто стандартом в современном NLP. Трансформерные модели состоят из блоков, каждый из которых включает два типа слоев: Attention и Feed-forward. Обычно их роль интерпретируется так: с помощью feed-forward модель «понимает», что значит каждый отдельный токен, а Attention отвечает за анализ связей между разными токенами. Хотя ключевые элементы механизма in-context learning обусловлены attention слоями, понимание работы feed-forward слоев дает важную интуицию, которая будет полезна при изучении Attention. Поэтому начнем с краткого обзора двух публикаций, посвященных feed-forward слоям.

Transformer Feed-Forward Layers Are Key-Value Memories

Статья посвящена тому, как модель распознает и обрабатывает значения слов и паттернов в тексте. Авторы приходят к выводу, что память трансформера организована как key-value хранилище, которое связывает текстовые паттерны и активации слоев трансформера. Это key-value хранилище реализуется в виде матриц feed-forward слоев: каждый столбец первой матрицы авторы рассматривают как «ключ», связанный с определенным паттерном текста, а вторая матрица, в свою очередь, хранит  соответствующие «значения».

Во время инференса вычисляется скалярное произведение между эмбеддингами токенов и ключами. Высокое значение этого произведения показывает сильную связь между токеном и ключом, что позволяет модели распознавать паттерны текста. В таблице представлены примеры таких ключей, которые связаны как с отдельными словами, так и с семантическими паттернами.

Например, скалярное произведение соответствующего столбца в feed-forward слое с эмбеддингами токенов, связанных с телешоу, возвращает высокое значение (последняя строчка в таблице выше). Это означает, что данный столбец отвечает за распознавание моделью тематик, связанных с телевидением.

Таким образом, feed-forward слои играют роль key-value хранилища и запоминают текстовые паттерны и значения слов. Это позволяет модели в ходе инференса распознавать эти паттерны в тексте и соответствующим образом изменять эмбеддинги слов, обогащая их новой информацией.

The Dual Form of Neural Networks Revisited

Если в первой статье обсуждалась связь между паттернами текста и параметрами feed-forward слоя, то авторы второй статьи идут дальше и устанавливают связь между параметрами и примерами из обучающего датасета. Они показывают, что конкретные параметры модели запоминают информацию из определенного подмножества обучающих данных. В таблице представлен пример: один из столбцов матрицы оказался связан с обучающими текстами, содержащими слова principal, primary и main, при этом схожее по структуре, но иное по смыслу предложение со словом first не приводит к активации этой группы параметров. Это означает, что в процессе обучения модели на каждую группу параметров влияет определенная часть обучающего датасета. Некоторые примеры из него оказывают большее влияние на определенные параметры, благодаря чему модель запоминает информацию из этих примеров и может использовать ее во время инференса.

Далее авторы рассматривают процесс обучения feed-forward слоев. Обновление параметров слоя методом градиентного спуска может быть записано следующим образом:

где

• — начальные параметры;

• — вектор градиента для i‑го примера;

• — входные данные.

Инференс нового примера можно выразить так:

Это выражение показывает, как модель использует информацию из обучающей выборки для генерации предсказания. Авторы сравнивают это с механизмом Unnormalised Dot Attention, который отличается от классического Attention отсутствием softmax и нормировки:

где и — ключи и значения из контекста.

Эта формула очень похожа на предыдущую: обучающие данные играют роль ключей, а градиенты – роль значений. Авторы вводят концепцию двойственной формы (The Dual Form of linear layers) линейного слоя, показывая, что обучение feed-forward слоя можно рассматривать как механизм Attention, применяемый к обучающим данным:

где

• — матрица градиентов;

• — обучающие данные.

Итак, feed-forward слой обладает двойственностью: он обучается градиентным спуском, но в то же время его обучение — это Attention на обучающих примерах. Это очень полезная аналогия, которая показывает, что во время обучения модель «фокусирует» внимание на тех примерах, которые будут наиболее полезны для предсказания.

Why Can GPT Learn In-Context?

В предыдущей части мы узнали, что обучение feed-forward слоя градиентным спуском можно представить как механизм Attention. Авторы следующей статьи предлагают взглянуть на эту аналогию с другой стороны: механизм Attention как процесс градиентного спуска. Давайте сначала вспомним, как устроен Attention.

Этот механизм оперирует тремя типами эмбеддингов: запросы, ключи и значения. Запросы (Q) представляет собой эмбеддинги токенов, для которого нужно сделать предсказание, ключи (K) — это эмбеддинги контекстных токенов, а значения (V) — связанные с ними эмбеддинги, которые нужно усреднить для получения обновленного эмбеддинга запроса. Модель вычисляет скалярное произведение запроса и ключей, чтобы определить attention-веса, с которыми суммируются значения. Полученный эмбеддинг используемый для предсказания следующего токена.

Рассмотрим этот процесс в сетапе in-context learning. Пусть — демонстрационные примеры, — пример, для которого нужно сделать предсказание, а – их конкатенация, то есть полный инпут модели. В таблице представлен пример из sentiment-classification датасета, где демонстрационные примеры помогают модели определить тональность текста :

Механизм Linear Attention для этого случая выглядит следующим образом:

где

• — матрицы параметров для ключей и значений;

• — эмбеддинг запроса, для которого делается предсказание.

Из этого выражения видно, что Linear Attention позволяет разделить этот механизм на два независимых слагаемых: одно зависит от токенов инструкции, а другое — от токенов демонстрационных примеров.

Сделаем следующий шаг и обозначим — матрица, соответствующая zero-shot learning, где нет демонстрационных примеров и присутствует только запрос, для которого делается предсказание. Это состояние можно интерпретировать как исходные параметры Attention, аналогичные стартовым параметрам модели перед дообучением на демонстрационных примерах.

Теперь соединим две последние формулы:

где мы обозначили изменение внутреннего состояния Attention под влиянием демонстрационных примеров. Последнее выражение в этой формуле похоже на обновление параметров модели с помощью градиентного шага. Мы видим прямую аналогию между Attention и градиентным спуском: играет роль начальных параметров , а аналогично градиентному шагу . Добавление к авторы называют метаоптимизацией по аналогии с оптимизацией градиентным спуском. В ходе метаоптимизации возникают метаградиенты – так авторы обозначают values в механизме attention, которые действуют аналогично градиентам в обычной оптимизации.

В итоге мы имеем альтернативное понимание роли Attention в in-context learning — оптимизация внутреннего состояния модели методом градиентного спуска. Из токенов контекста получаются метаградиенты, под их влиянием состояние модели меняется и она генерирует более правильное предсказание для токена запроса. Важно помнить, что это изменение виртуально и существует только в момент инференса и вычисления Attention. После прохождения данных через модель «памяти» о них не сохраняется и параметры модели остаются прежними. Это в каком-то смысле реализует механизм краткосрочной памяти, которая позволяет модели эффективно работать с информацией из контекста диалога.

Для проверки корректности предложенной аналогии авторы проводят ряд экспериментов. Они используют две предобученных модели GPT размерами 1.3B и 2.7B параметров, а также 6 классификационных датасетов. Для каждого датасета были выбраны 32 примера, которые применяются в следующих сценариях:

1. ZSL: модель делает предсказание без демонстрационных примеров

2. ICL: к ZSL в контекст добавляются демонстрационные примеры

3. FT: модель проходит одну эпоху градиентного спуска на демонстрационных примерах, после чего делается ZSL

В таблице ниже показаны accuracy для 6 датасетов. Во всех случаях FT улучшает качество предсказания по сравнению с ZSL. Однако использование тех же примеров в ICL сценарии дает еще более высокую accuracy. Вполне ожидаемый результат, что и добавление примеров в контекст, и дообучение на них улучшают качество классификации.

Далее представлена метрика Rec2FTP, которая вычисляется следующим образом:

Она показывает долю точек, где ICL и FT одновременно дают более правильные предсказания, чем ZSL. В среднем в 85% точек ICL и FT дают одинаковые предсказания, что указывает на схожее поведение модели в сетапе дообучения и in-context learning.

Следующая таблица показывает метрику SimAOU:

где

• — векторное представление последнего токена из последнего слоя модели;

• — изменение векторного представления под действием in‑context примеров;

• — изменение векторного представления в ходе файнтюна.

Эта метрика показывает, насколько схожи изменения скрытых представлений ZSL под влиянием метаградиентов в ICL и обычных градиентов в FT. Результаты показывают, что FT изменяет эмбеддинги модели в том же направлении, что и ICL, что означает сонаправленность градиентов и метаградиентов в этих двух механизмах. Для сравнения SimAOU также вычисляется между ICL и случайным изменением эмбеддингов, и в этом случае близость изменений околонулевая.

Таким образом, авторы предложили аналогию между механизмом Attention и градиентным спуском, а также ввели понятия метаградиента и метаоптимизации. Такую аналогию они подкрепили экспериментами, в которой показали, что in-context learning и дообучение на демонстрационных примерах показывают схожее поведение и похожим образом влияют на внутреннее состояние модели.

Предложенная аналогия позволяет глубже понять, как работает Attention в современных NLP моделях. Всякий раз, когда мы показываем LLM несколько примеров выполнения задания, механизм Attention реализует подобие градиентного спуска, на лету обновляет внутреннее состояние модели и делает генерацию более точной. Этот процесс представляет собой некоторое подобие кратковременной памяти и позволяет языковым моделям эффективно работать с длинным контекстом.