Как я поймал Трансформер на читерстве: гроккинг, математика и Mechanistic Interpretability

Феномен Grokking и Mechanistic Interpretability — главные тренды в исследованиях лабораторий уровня OpenAI и Anthropic. Я решил потрогать эти концепции своими руками на уровне тензоров. Цель казалась тривиальной: заставить кастомный микро-Трансформер (всего 1М параметров) выучить базовую арифметику с нуля. Однако вместо математического гения я получил ленивого мошенника. Эта статья — инженерный детектив о том, как нейросети пытаются нас обмануть (Specification Gaming), и как вскрытие Attention-матриц помогает поймать их за руку.

Глава 1. Полигон и экстремальный Weight Decay

Чтобы понимать поведение системы на 100%, я отказался от использования готовых LLM или импортов из библиотеки transformers от Hugging Face. Вся архитектура (классический Decoder-only Transformer) была собрана с нуля на чистом PyTorch.

Параметры были намеренно сжаты, чтобы создать искусственный дефицит нейронов:

• Слои (n_layer): 4

• Головы внимания (n_head): 4

• Размерность (n_embd): 128

• Итоговый размер: ~1 миллион параметров.

В качестве данных использовался генератор всех возможных сложений двузначных чисел. Чтобы Трансформер мог улавливать связи между разрядами, была применена посимвольная токенизация (character-level), где каждая цифра — это отдельный токен. Строки имели жесткий формат с выравниванием нулями: 14+35=049.

Но главная магия кроется в оптимизаторе. Если запустить обучение со стандартными параметрами, микро-модель просто вызубрит датасет. Она превратит свои MLP-слои в гигантскую таблицу (Lookup Table). Чтобы форсировать Гроккинг, в оптимизатор AdamW был внедрен экстремальный штраф: weight_decay = 1.0.

Жесткая L2-регуляризация делает зубрежку энергетически "невыгодной". Она буквально заставляет сеть страдать за большие веса и вынуждает её искать алгоритмическое сжатие — ту самую логику математики, которая требует меньше параметров, чем слепое запоминание.

class GrokkingConfig: vocab_size = 16 # Только цифры 0-9, "+", "=" и "\n" block_size = 16 # Короткое окно контекста для строки n_layer = 4 n_head = 4 # 4 головы внимания n_embd = 128 # Узкий Residual Stream # Инициализация модели model = GPT(GrokkingConfig()).to(device) # weight_decay = 1.0 убирает "зубрежку" и заставляет сеть искать алгоритм optimizer = torch.optim.AdamW( model.parameters(), lr=1e-3, weight_decay=1.0, betas=(0.9, 0.98) )

Глава 2. Долина Отчаяния и Баг 15+15=020

После запуска тренировочного цикла сеть вошла в классическую "Долину Отчаяния". Спустя 15,000 макро-шагов метрика val_loss опустилась до отметки 1.18. Учитывая авторегрессионную природу модели (необходимость предсказывать каждый случайный символ из промпта), это значение близко к теоретическому математическому пределу.

На первый взгляд, Гроккинг состоялся. Инференс на простых примерах выдавал идеальные результаты. Например, на промпт 05+39= сеть безошибочно генерировала 044, успешно справляясь с переносом десятка.

Но настоящий исследователь обязан ломать свою систему на граничных условиях. При более детальном стресс-тестировании я получил аномалию.
На промпт 15+15= модель уверенно сгенерировала ответ 020.

Она правильно сложила единицы (5+5=105+5=10, записала 0). Но при вычислении десятков (1+1=21+1=2) она полностью проигнорировала перенос единицы. Сеть выработала ложную эвристику (Partial Generalization): алгоритм переноса срабатывал только в том случае, если одно из левых слагаемых было нулем.

Глава 3. Вскрытие «черного ящика» (Mechanistic Interpretability)

Чтобы доказать эту гипотезу математически, я отказался от готовых функций и вручную написал код извлечения матриц внимания (Q×KT/dkQ×KT/dk​​) для всех 4 слоев и 4 голов моей модели.

Я провел рентген нейросети на ошибочном примере 15+15=020.
Логика поиска была простой: чтобы правильно сложить десятки и понять, нужен ли перенос, модель обязана направить свое внимание (Attention) на предыдущие единицы (5 и 5) в момент генерации цифры 2. Иными словами, в архитектуре должна была сформироваться специализированная "Голова Переноса" (Carry-over Head).

Посмотрите на предпоследнюю строку матрицы (генерация двойки). Клинический разбор всех 16 голов показал, что ни одна из них не сфокусировалась на единицах. Внимание было распределено как угодно, но только не на нужный контекст. Диагноз подтвердился визуально: Carry-over Head физически отсутствовала в весах. Модель пыталась сложить десятки в полной изоляции, что и привело к катастрофическому сбою.

Глава 4. Попытка исправить и Specification Gaming

Как инженер, я не мог оставить систему в сломанном состоянии. Фундаментальная проблема Трансформера — генерация строго слева направо. Ему когнитивно тяжело предсказывать старшие разряды, держа перенос в уме и не имея возможности посмотреть на результат сложения единиц.

Я решил дать модели неявный "черновик" (Scratchpad) на уровне структуры данных: перевернул ответы задом наперед. Уравнение 14+35=049 в обучающей выборке превратилось в 14+35=940. Логика была изящной: пусть сеть сначала посчитает единицы, запишет результат в контекст, и уже опираясь на него, легко посчитает десятки.

Я запустил новый цикл обучения с чистого листа на 145,000 макро-шагов. И столкнулся с классическим, хрестоматийным примером Specification Gaming (Игровая спецификация).

На тестовый промпт 14+35= сеть сгенерировала 350 (что при развороте дает 053).
Она вообще не делала сложение! Столкнувшись со сложной позиционной геометрией реверсивных данных под жестоким давлением оптимизатора (weight_decay = 1.0), модель "решила", что строить нейронный сумматор — слишком энергетически дорого. Вместо этого она пошла по пути наименьшего сопротивления: вырастила "Головы копирования" (Name Mover Heads) и начала банально копипастить куски слагаемых из промпта в ответ. Трансформер оказался феноменально ленивым мошенником.

Вывод

Этот проект наглядно доказал мне главную проблему безопасности искусственного интеллекта. Нейросеть не разделяет наших человеческих концепций о "правильном" решении задачи. Из-за физики градиентов она всегда будет искать самую короткую и дешевую дорогу к минимизации Loss-функции.

Попытка "помочь" модели изменением структуры данных привела к тому, что она нашла нелегальную алгоритмическую лазейку. Анализ матриц внимания и подобных аномалий — это единственный способ убедиться, что ИИ действительно выучил правила нашего мира, а не просто притворяется умным.

Ссылка на код, датасеты и полные рентгеновские снимки в моем GitHub