Transformer Circuits: «A Mathematical Framework for Transformer Circuits» — конспект

Статья открыла серию Transformer Circuits Thread — журнала Anthropic о механистической интерпретируемости. Её цель скромная и смелая одновременно: начать «обратную разработку» трансформеров так, как программист восстанавливает исходный код по бинарнику. Авторы берут самые простые модели — трансформеры из 0, 1 и 2 слоёв только с вниманием (без MLP) — и показывают, что их можно понять почти полностью, если записать трансформер в непривычной, но математически эквивалентной форме.

О лицензии и формате конспекта

На сайте transformer-circuits.pub нет открытой лицензии для текстов, поэтому здесь пересказ своими словами с формулами и нашими числовыми примерами, без копирования рисунков. Оригинал интерактивный (графики голов внимания можно «трогать»), и к нему есть серия видеолекций авторов. Конспект — по тексту статьи на transformer-circuits.pub.

Коротко

• Остаточный поток — это канал связи. Каждый слой читает из него линейной проекцией и пишет обратно, прибавляя свой результат. Сам поток ничего не вычисляет, и он полностью линеен.
• Головы внимания независимы и аддитивны. Привычное «склеить выходы голов и умножить на $W_O$» — то же самое, что сумма вкладов отдельных голов.
• У каждой головы две почти независимые части: QK-схема решает, откуда брать информацию (паттерн внимания), а OV-схема — что переносить и как это влияет на выход. Ключи, запросы и значения — лишь промежуточные продукты низкоранговых матриц $W_Q^\top W_K$ и $W_O W_V$.
• Модель раскладывается в сумму путей. Если заморозить паттерны внимания, модель с одним вниманием становится линейной функцией токенов, и её можно читать прямо из весов.
• 0 слоёв — биграммы; 1 слой — биграммы + skip-триграммы «A… B → C»; 2 слоя — композиция голов и индукционные головы [A][B] … [A] → [B], общий алгоритм обучения в контексте.

Как читать вместе с курсом

• До статьи нужно понимать устройство блока трансформера и внимание: уроки 4.4.1 и 4.5.1 (там же — остаточное соединение).
• Вместе с уроком 9.3.1 «Заглядываем внутрь трансформера»: урок показывает остаточный поток, logit lens и индукционные головы на практике, статья — теорию, почему всё так устроено.
• Что пропустить при первом чтении. Тензорную запись через произведение Кронекера (раздел «Attention Heads as Information Movement» можно понять по нашему пересказу ниже), анализ собственных значений и «Term Importance Analysis». Вернитесь к ним, когда будете читать статью об индукционных головах.
• Видео. У статьи есть серия из 7 коротких видео авторов (ссылки в разделах оригинала): 0 слоёв, 1 слой — теория и результаты, 2 слоя — теория, важность слагаемых и результаты.

Что упрощено в моделях

Авторы честно предупреждают, что изучают «игрушечные» трансформеры:

• только внимание, без MLP — сильное упрощение: в обычном трансформере MLP занимают около двух третей параметров, а с ними у авторов «гораздо меньше успехов»;
• без смещений (bias) — их можно «вшить» в веса, добавив всегда равную единице координату;
• без нормализации слоёв (LayerNorm) — с точностью до масштаба её можно влить в соседние веса;
• декодер-трансформер вроде GPT (не encoder-decoder из оригинальной статьи о трансформере).

Остаточный поток как канал связи

Трансформер — это эмбеддинг токена, затем серия остаточных блоков (внимание, потом MLP), затем «разэмбеддинг» в логиты. Ключевая мысль статьи — смотреть на остаточный поток (residual stream) как на общую шину, через которую общаются все слои:

$$x_{\text{после слоя}} = x_{\text{до слоя}} + \text{слой}(x_{\text{до слоя}}) .$$

Словами: каждый слой не заменяет вектор, а добавляет к нему свой вклад. Поэтому вектор в потоке — просто сумма исходного эмбеддинга и выходов всех предыдущих слоёв.

Следствия, которые выводят авторы:

1. Нет «привилегированного базиса». Поток полностью линеен: если повернуть его (и согласованно повернуть все матрицы, которые с ним работают), поведение модели не изменится. Значит, отдельные координаты потока сами по себе ничего не значат.
2. «Виртуальные веса». Раз всё линейно, можно перемножить выходную матрицу одного слоя и входную матрицу другого, даже далёкого, и получить «виртуальную» связь между ними: насколько поздний слой читает то, что записал ранний.
3. Подпространства и «пропускная способность». Поток многомерен (сотни измерений в маленьких моделях, десятки тысяч в больших), а голова внимания работает в маленьком подпространстве (64 или 128 измерений). Слои могут писать в разные подпространства и не мешать друг другу; записанное там остаётся, пока кто-то его не сотрёт.
4. Узкое горлышко. «Вычислительных» измерений (нейронов MLP, выходов голов) намного больше, чем измерений потока: один слой MLP обычно в 4 раза шире потока. Пример из статьи: на 25-м слое 50-слойной модели до потока «стучатся» в 100 раз больше нейронов, чем у него измерений, — информация обязана передаваться «в суперпозиции». Такие активации-«горлышки» трудно интерпретировать напрямую — эта мысль позже выросла в статью о суперпозиции.
5. «Уборщики». Есть намёки, что некоторые нейроны и головы стирают информацию из потока: читают направление и пишут его со знаком минус.

Головы внимания: независимые и аддитивные

Обычно слой внимания записывают так: выходы $r^{h_1}, r^{h_2}, \dots$ всех голов склеивают в один длинный вектор и умножают на большую матрицу $W_O^H$. Но если разрезать $W_O^H$ на блоки по головам, получится

$$W_O^H \begin{bmatrix} r^{h_1} \\ r^{h_2} \\ \vdots \end{bmatrix} = \sum_i W_O^{h_i}\, r^{h_i} .$$

Словами: «склеить и умножить» — то же самое, что каждой голове умножить свой выход на свою матрицу и сложить результаты. Склейка удобна для быстрых вычислений, а для понимания удобнее думать о головах как о независимых слагаемых.

Числовой пример (наш). Две головы с выходами $r^{1} = (1, 2)$ и $r^{2} = (3)$, блоки $W_O^{1} = \begin{bmatrix}1 & 0\\0 & 1\end{bmatrix}$, $W_O^{2} = \begin{bmatrix}1\\1\end{bmatrix}$. Склейка: $[W_O^1, W_O^2]\cdot(1, 2, 3)^\top = (1 + 3,\ 2 + 3) = (4, 5)$. По отдельности: $W_O^1 r^1 + W_O^2 r^2 = (1, 2) + (3, 3) = (4, 5)$ — то же самое.

Голова внимания переносит информацию: QK и OV

Главное действие головы — переносить информацию из остаточного потока одного токена в поток другого. Обычное вычисление в три шага (значения $v_j = W_V x_j$, их взвешенная сумма по паттерну внимания $r_i = \sum_j A_{ij} v_j$, выход $W_O r_i$) авторы переписывают так, чтобы стало видно: какие токены участвуют и какая информация переносится — две разные, почти независимые операции.

• Паттерн внимания $A$ решает, откуда куда переносить. Он вычисляется из матрицы $W_{QK} = W_Q^\top W_K$: $A = \operatorname{softmax}(x^\top W_Q^\top W_K\, x)$ (с маской будущих токенов).
• Матрица $W_{OV} = W_O W_V$ решает, что читается из исходного токена и как записывается в целевой.

Наблюдения авторов:

• $A$ — единственная нелинейная часть головы. Если заморозить паттерн внимания, голова становится линейной операцией.
• $W_Q$ и $W_K$ всегда работают вместе, $W_O$ и $W_V$ — тоже. Ключи, запросы и значения «в некотором смысле поверхностны»: это промежуточные результаты при вычислении двух низкоранговых матриц $W_{QK}$ и $W_{OV}$ (ранг не больше $d_{\text{head}}$). Их можно перепараметризовать, не меняя работу головы.
• Произведение двух голов ведёт себя как голова: паттерн $A^{h_2}A^{h_1}$ и матрица $W_{OV}^{h_2}W_{OV}^{h_1}$. Это виртуальные головы внимания.

Если перейти от остаточного потока к токенам (умножить на эмбеддинг $W_E$ и разэмбеддинг $W_U$), получаются две матрицы размера «словарь × словарь»:

Схема	Формула	Что означает элемент
QK-схема (query-key)	$W_E^\top W_{QK}^h W_E$	насколько токен-«запрос» хочет смотреть на токен-«ключ»
OV-схема (output-value)	$W_U W_{OV}^h W_E$	как токен, на который посмотрели, меняет логиты следующего токена

Мысленный эксперимент авторов, объясняющий, почему их можно изучать по отдельности: прогоните модель один раз и запишите все паттерны внимания (это зависит только от QK); прогоните второй раз, подставив записанные паттерны, — и логиты станут линейной функцией токенов (это уже только OV).

Нулевой слой: биграммы

Модель без слоёв: эмбеддинг, сразу разэмбеддинг, $T = W_U W_E$. Информацию между токенами она переносить не может, поэтому лучшее, что ей доступно, — предсказывать следующий токен по текущему, то есть приближать биграммные лог-вероятности. Такое слагаемое $W_U W_E$ («прямой путь») есть в разложении любого трансформера: оно берёт на себя биграммы, которые не объясняются общими правилами, — например, что за «Barack» часто идёт «Obama». Отсюда вывод: не стоит ждать, что разэмбеддинг $W_U$ — обратная матрица к эмбеддингу $W_E$.

Один слой: ансамбль биграмм и skip-триграмм

Трюк с разложением на пути

Модель записывается как произведение множителей «по слоям». Если раскрыть скобки, произведение превращается в сумму путей от входа к выходу:

$$\text{логиты} \approx \underbrace{W_U W_E}_{\text{прямой путь}} \;+\; \sum_{h} \underbrace{A^h \otimes W_U W_{OV}^h W_E}_{\text{путь через голову } h} .$$

Словами: логиты — это сумма прямого пути (биграммы) и вкладов каждой головы, где каждая голова смотрит на токены по своему паттерну $A^h$ и сдвигает логиты по своей OV-схеме. Каждый путь можно понимать отдельно.

Skip-триграммы

QK-схема выбирает, на какой прошлый токен-«источник» смотрит текущий токен-«цель», а OV-схема — как этот источник меняет предсказание следующего токена. Вместе три токена образуют skip-триграмму [источник] … [цель] → [выход]. Однослойная модель — это огромная, но понятная таблица таких правил. «Огромная» буквально: при словаре около 50 000 токенов одна развёрнутая OV-матрица содержит около 2,5 млрд элементов; авторы сравнивают модель с «сжатой китайской комнатой» — горой карточек с правилами.

Что нашлось среди крупнейших элементов матриц (примеры из статьи):

• Копирование — примитивное обучение в контексте. Большинство голов однослойных моделей тратят огромную часть ёмкости на копирование: OV-схема повышает вероятность того самого токена, на который посмотрели, а QK-схема смотрит только на токены, которые правдоподобно могут идти следующими. Интересный частный случай связан с токенизацией: « Ralph» с пробелом — один токен, а «Ralph» без пробела (например, в начале абзаца) — два, «R» + «alph». Есть головы, которые, увидев «R», ищут в контексте « Ralph» и предсказывают «alph».
• Код и разметка. После уменьшения отступа в Python растёт вероятность else/elif/except; после open(…, " — режимов rb, wb, r, w; после def … ( — self; в LaTeX после \left … \ — right; в HTML после tbody … < — td; адреса: github … / → issues, blob, pull.
• Язык. keep … in → mind, keep … at → bay; для одной головы с запросом « and»: back and → forth, day and → night, sick and → tired.
• Позиционные головы — небольшая часть голов смотрит в основном на определённые относительные позиции (на текущий или предыдущий токен).

«Баги» skip-триграмм

Голова хранит правило в «факторизованном» виде — часть в QK, часть в OV, — как функцию $f(a, b, c) = f_1(a, b)\,f_2(a, c)$. Поэтому она не может гибко учесть тройное взаимодействие: если голова повышает вероятность «keep… in mind» и «keep… at bay», она неизбежно повышает и «keep… in bay» и «keep… at mind». Авторы называют это ранней демонстрацией того, как интерпретируемость объясняет ошибки модели.

Как найти копирующие головы автоматически

Копирование — это когда токен повышает свою же вероятность, то есть матрица отображает вектор «почти в себя». Удобная (но не идеальная) статистика — собственные значения OV-схемы: у копирующей матрицы они в основном положительные, у случайной — поровну положительных и отрицательных (и много комплексных). В одной из моделей так нашлось, что 10 из 12 голов заметно копируют, — это совпало с ручным просмотром весов.

Понимаем ли мы однослойную модель «полностью»?

Ответ авторов: в том же смысле, в каком понимаем огромную линейную регрессию или базу данных. Алгоритмической загадки больше нет, но удержать модель в голове невозможно — слишком много правил.

Два слоя: композиция и индукционные головы

Три вида композиции

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

Вид	Что читает результат первой головы	Эффект
Q-композиция	$W_Q$ второй головы	меняет паттерн внимания (через запрос)
K-композиция	$W_K$ второй головы	меняет паттерн внимания (через ключ)
V-композиция	$W_V$ второй головы	меняет что переносится — получается «виртуальная голова»

Q- и K-композиция качественно сильнее: они позволяют строить сложные паттерны внимания, которых одна голова не выразит. V-композиция лишь соединяет два переноса в один перенос.

Индукционные головы

В разобранной двухслойной модели почти вся композиция оказалась K-композицией одной головы первого слоя (она смотрит на предыдущий токен) с несколькими головами второго слоя — индукционными головами (induction heads). Большинство остальных голов ни в какой заметной композиции не участвуют и работают как ещё одна порция skip-триграмм.

Что делает индукционная голова: ищет в контексте прошлые появления текущего токена; если нашла — смотрит на токен, который шёл после него, и копирует его; если нет — смотрит на первый (служебный) токен и ничего не делает. Сравнение алгоритмов:

Модель	Правило	Пример
1 слой, копирующая голова	[b] … [a] → [b]	слово встречалось — повысим его вероятность
1 слой, редкий случай токенизации	[ab] … [a] → [b]	« Ralph» … «R» → «alph»
2 слоя, индукционная голова	[a][b] … [a] → [b]	«Mr Dursley … Mr» → «Dursley»

Двухслойный алгоритм сильнее: он помнит, как именно токен использовался раньше, и меньше зависит от выученной статистики. Он работает даже на случайных последовательностях токенов, повторённых несколько раз, — самое жёсткое испытание, потому что никакой «обычной» статистики там нет.

flowchart LR
T1["… Mr"] --> T2["Dursley …"] --> T3["… Mr (сейчас)"]
H1["Голова 1-го слоя: в позицию Dursley пишет «передо мной было Mr»"] -.-> T2
H2["Индукционная голова: запрос «Mr» ищет ключ «передо мной было Mr»"] -.->|"находит позицию Dursley"| T2
T2 -.->|"OV копирует Dursley"| P["Предсказание: Dursley"]

Как работает индукционная голова: ключ сдвинут на один токен (наша схема)

Как это устроено. Центральный трюк: ключ вычисляется из токена, сдвинутого на один назад. Голова первого слоя («смотрит на предыдущий токен») записывает в каждую позицию информацию о том, какой токен стоял перед ней. Индукционная голова через K-композицию использует эту информацию в ключах: запрос «я — Mr» находит ключ «передо мной был Mr» — то есть позицию токена, шедшего после прошлого «Mr». А OV-схема головы копирующая и повышает вероятность этого токена.

Проверка теории. Если теория верна, у индукционной головы должны быть копирующая OV-схема и «сопоставляющая» QK-схема (на слагаемом с K-композицией). По собственным значениям все индукционные головы оказались в крайнем углу «оба положительные» — хотя ничто не требовало, чтобы сильная K-композиция была именно положительной.

Виртуальные головы и важность слагаемых

Чтобы проверить, не важны ли «в сумме» виртуальные головы (V-композиция), авторы придумали алгоритм абляции путей заданной длины (прогоны с замороженным вниманием, где выходы голов подменяются выходами прошлого прогона). Вывод: в маленькой двухслойной модели виртуальные головы почти не важны; главное — прямой путь и отдельные головы, особенно второго слоя. Но в больших моделях виртуальных голов может быть очень много: число пар голов растёт квадратично с числом слоёв, и авторы предполагают, что там они важнее (например, «смотреть на подлежащее предыдущего предложения»).

Исправление в оригинале

В опубликованной версии схема силы Q-, K- и V-композиции между головами содержала ошибку из-за бага в библиотеке для низкоранговых матриц; исправленная схема приведена в разделе комментариев статьи. Выводы о K-композиции и индукционных головах это не меняет.

Что дальше (вывод авторов)

Методы применимы к частям больших моделей: в обычных трансформерах есть схемы, которые держатся в основном на внимании, и большие модели тоже образуют индукционные головы на основе K-композиции с головой «предыдущего токена». Это, по словам авторов, похоже на центральный механизм обучения в контексте у моделей любого размера — тема следующей статьи. Но MLP-слои — две трети параметров — остаются непонятыми; для прогресса нужны интерпретируемые нейроны, а их найти не удавалось. Этот тупик привёл к работам о суперпозиции и словарном обучении.

Что сделать руками

1. Повторите практику урока 9.3.1: найдите индукционные головы в маленькой модели на повторённой случайной последовательности (prefix matching score).
2. Для двухслойной модели только с вниманием из TransformerLens (attn-only-2l) посчитайте матрицу $W_U W_{OV}^h W_E$ для одной головы и найдите 20 самых больших элементов на диагонали — это копирование.
3. Проверьте «баг» skip-триграмм: найдите голову, которая повышает и «in mind», и «at bay» после «keep», и посмотрите, растёт ли вероятность «in bay».
4. Для разминки по линейной алгебре — урок о матричном умножении и ранге на ступени 1: почему ранг $W_O W_V$ не больше $d_{\text{head}}$?

Словарик

EN	RU
mechanistic interpretability	механистическая интерпретируемость
residual stream	остаточный поток
privileged basis	привилегированный базис
virtual weights	виртуальные веса
attention-only transformer	трансформер только с вниманием
QK circuit / OV circuit	QK-схема / OV-схема
path expansion	разложение на пути
direct path	прямой путь
skip-trigram	skip-триграмма (триграмма с пропуском)
Q-/K-/V-composition	Q-/K-/V-композиция
virtual attention head	виртуальная голова внимания
induction head	индукционная голова
copying head	копирующая голова
eigenvalue	собственное значение