Код: https://github.com/fanat503/Induction-Heads-Tinystories
Введение
Исследования Induction Heads и SAE очень важны, чтобы разобраться в современных моделях и понять, что же в них происходит. Цель эксперимента — проверить формирование Induction Heads на датасете TinyStories. Почему был выбран именно такой датасет? Потому что он очень простой и понятный, что позволяет изолировать влияние сложности данных на формирование механизмов.
Архитектура и методология
Моделью была GPT-2 Small, реализованная с нуля на PyTorch, на 124М параметров. Сам датасет состоял из 473М токенов (~5.5 эпох, 20,000 шагов), что позволило прогнать модель несколько раз на каждом токене. Метрики проверялись на функциях compute_induction_score и compute_previous_token_score, которые были также написаны самостоятельно. Архитектура SAE состояла из Encoder, который принимал 768 измерений и переносил их в 8192 измерений, а также Decoder, который переносил всё обратно в 768 измерений. l1_coeff равнялся 3.9, количество эпох обучения SAE — 25.
Результаты
3.1 Previous Token Heads
Previous Token Heads чётко формировались в начальных слоях (0–4). Наибольший score показала голова L0H3: 0.20. Стабильная динамика начала вырисовываться уже с 14,000 шагов (из 20,000).
3.2 Induction Heads
Induction Heads показали совсем другие результаты. Максимальный score, который показала голова L6H4, был равен 0.05, что объясняется характеристиками датасета. Пик результатов пришёлся на 15,000–16,000 шагов, а дальше пошёл медленный плавный спад, который почти не отразился на итоговом результате.
Начальная гипотеза эксперимента была такова, что Induction Heads могут слабо сформироваться на таком простом датасете — что позже и подтвердилось с помощью метрик, описанных выше.
3.3 Sparse Autoencoder
SAE обучен на активациях 6-го слоя, так как именно ему принадлежит голова L6H4, показавшая наивысший Induction Score. Active features находились в диапазоне 20–35 фичей на каждый токен, что позволяет хорошо их интерпретировать. Dead features составили приблизительно 3% (~228/8192) — неплохой результат, достигнутый за счёт L1-регуляризации, которая минимизировала количество фичей с ненулевой активацией.
Были найдены 2 главные фичи, которые сначала казались Induction-фичами, но при более длительной проверке на синтетических и реальных данных выяснилось следующее: F#8107 отвечает за синтаксически предсказуемый контекст, а F#635 — за общий контекстный паттерн. В итоге чистых Induction-фичей найдено не было, что подтвердило низкий score на тестах.
Обсуждение
Induction Heads оказались такими слабыми на TinyStories, так как простой датасет не требует глубоких механизмов In-context Learning. Из-за этого модели было "выгодно" сформировать качественные Previous Token Heads и не развивать Induction-механизм. Гипотеза на следующий эксперимент: Induction Heads будут лучше себя проявлять на более сложных датасетах, например OpenWebText, так как там модели для достижения более низкого loss придётся освоить данный алгоритм.
Следующий эксперимент будет посвящён обучению моделей разных размеров на TPU на двух датасетах (TinyStories и OpenWebText), чтобы определить, верна ли эта гипотеза.
Заключение
Previous Token Heads формируются стабильно на любых данных, так как они необходимы языковым моделям для предсказания следующего токена и минимизации loss. Induction Heads зависят от статистики повторов в датасете: если их нет, модели не имеет смысла выделять ресурсы на этот механизм. SAE при этом отлично помогает в визуализации и верификации результатов.
Код: https://github.com/fanat503/Induction-Heads-Tinystories
