Перейти к содержимому
9/30Глава 9 из 30

Attention и блок transformer: вывод из среднего

Начинаем с самого дешёвого резюме контекста — среднего, измеряем провал и выводим формулу attention из исправления.

На этой странице

Вы приходите сюда с tokenizer из главы 7, таблицей embedding из главы 8 и соответствующей им целью: по уже данным tokens назначить вероятность следующему.

Не хватает середины. Чтобы предсказать token tt, модели нужен один вектор, который суммирует всё до него, а ничто из уже построенного вами такого вектора не даёт. Embedding token t1t-1 — не он: это bigram-модель, и она не может знать, что предложение началось с вопроса. Конкатенация всех предыдущих embeddings тоже не подходит: их число меняется на каждом шаге, а фиксированная матрица весов не может принимать вход переменной длины.

Итак: один вектор фиксированного размера, суммирующий переменное число векторов. В этом вся задача, и attention — это то, что получается, если решить её самым ленивым возможным способом, а затем исправить две вещи, которые ломаются.

Ответ, который уже был у области, и почему мы строим не его

Ссылка на раздел: Ответ, который уже был у области, и почему мы строим не его

С 1997 примерно до 2017 года таким резюме было recurrent-состояние: хранить вектор h\mathbf{h} и обновлять его на каждом token, ht=f(ht1,xt)\mathbf{h}_t = f(\mathbf{h}_{t-1}, \mathbf{x}_t). Фиксированный размер, переменный вход, ровно нужная форма.

Оно ломалось тремя способами, и архитектура этой главы отвечает на все три. Backpropagation через TT шагов перемножает TT якобианов, поэтому gradient исчезает или взрывается — болезнь, которую глава 5 измеряла внутри одного узла tanh\tanh. LSTM1 была спроектирована ровно против этого и расширила рабочий диапазон с десятков шагов до сотен, но не изменила того факта, что информация от token 5 доходит до token 500, только пережив 495 последовательных обновлений. Весь источник должен был поместиться в один вектор: в sequence-to-sequence-переводе2 encoder сжимает вход в своё финальное состояние. Bahdanau, Cho и Bengio назвали это узкое место и исправили его в 2014 году, за три года до transformer, позволив decoder брать взвешенную сумму всех состояний encoder с весами, которые он вычислял сам.3 Всё ниже — та же идея, применённая последовательностью к самой себе, но без recurrence. И обновление по конструкции последовательное: ht\mathbf{h}_t нуждается в ht1\mathbf{h}_{t-1}, и GPU с десятью тысячами ядер ничего не может с этим сделать. Победившая архитектура не очевидно умнее; это та, чей дорогой шаг является умножением матриц.

Другой классический inductive bias, convolution — провести один маленький фильтр по всему входу, чтобы признак, найденный где угодно, находился везде, — здесь тоже не строится; он почти идеально подходит для изображений и делегируется курсу по vision. Ни recurrence, ни convolution после этой страницы не возвращаются, поэтому ни одна из них не получает главы: глава 1 обещала, что пропуски будут объявлены, а не спрятаны.

Самая очевидная функция, которая принимает переменное число векторов и возвращает один вектор, — это среднее:

ct=1ti=1txi\mathbf{c}_t = \frac{1}{t}\sum_{i=1}^{t} \mathbf{x}_i

Любое число входов, фиксированный размер выхода, дифференцируемо, бесплатно. Таблица embedding плюс это среднее плюс линейный слой к словарю — полноценная языковая модель в пятнадцать строк. Она ещё и ужасна, а то, как именно она ужасна, и есть весь вывод.

Корпус ниже — один мегабайт Shakespeare, 1 115 394 символа, пропущенных через byte-level BPE tokenizer того типа, который строился в главе 7, со словарём 1024: 459 760 tokens по 2,43 символа каждый, разбиение 90/10. Каждая модель имеет ширину 128, видит 128 tokens и обучается 3000 шагов AdamW при 10310^{-3} с batch 64. Перплексия считается на отложенной части.4

модельпараметрывалидационная перплексия
только текущий token, вообще без контекста263,16859.71
плюс равномерное среднее всего перед ним263,168248.07
плюс обучаемые position embeddings279,552245.93
равномерное среднее добавляется к token, а не заменяет его263,16860.45

Прочитайте вторую строку дважды. Усреднение контекста не помогает чуть-чуть; оно делает модель в четыре раза хуже, чем полное игнорирование контекста. Причины две, и обе доказуемы, а не эмпиричны.

Среднее не видит порядок. Сложение коммутативно, поэтому перемешивание окна оставляет резюме неизменным — не приблизительно:

order.pyPYTHON
A = torch.tril(torch.ones(T, T))
A = A / A.sum(1, keepdim=True)          # rows of the averaging matrix
y = x[torch.randperm(T)]                # the same tokens, shuffled
print((A[-1] @ x - A[-1] @ y).abs().max().item())
TEXT
2.9802322387695312e-08

Шум floating-point на переупорядоченной сумме: два резюме — один и тот же вектор. Модель, чей единственный взгляд на контекст — среднее, не может отличить the dog bit the man от the man bit the dog. Третья строка доказывает, что это не исправляется добавлением позиций ко входам: обучаемый position embedding на каждом token перед усреднением дал 2,14 пункта из 188. Позиции попадают в сумму, а сумма их забывает.

И среднее топит настоящее. В позиции 100 текущий token — одна сотая резюме. У этого есть дешёвое исправление, которое у вас уже есть: оставить token и добавить к нему резюме — residual connection из главы 6, и четвёртая строка показывает, что она делает. Когда разбавление исправлено, равномерное среднее не даёт вообще ничего: 60,45 против baseline 59,71. Каждый token там есть, с одинаковым весом, а одинаковое взвешивание эквивалентно отсутствию информации.

Проблема не в усреднении. Она в весах.

Среднее — это умножение матриц, а маска — это softmax

Ссылка на раздел: Среднее — это умножение матриц, а маска — это softmax

Усреднение по растущему префиксу выглядит как цикл. На самом деле это одно умножение на нижнетреугольную матрицу, строки которой суммируются в единицу, — и также, в точности, softmax:

mechanics.pyPYTHON
loop = torch.stack([x[:t + 1].mean(0) for t in range(T)])   # the obvious version

A = torch.tril(torch.ones(T, T))
A = A / A.sum(1, keepdim=True)
mat = A @ x                                                  # the same thing

S = torch.zeros(T, T).masked_fill(torch.tril(torch.ones(T, T)) == 0, float("-inf"))
soft = F.softmax(S, dim=-1) @ x                              # and the same thing again
TEXT
loop vs matmul   max |diff| = 5.960464477539063e-08
loop vs softmax  max |diff| = 5.960464477539063e-08

the averaging matrix A (rows sum to 1, upper triangle is zero):
  1.000 0.000 0.000 0.000 0.000 0.000
  0.500 0.500 0.000 0.000 0.000 0.000
  0.333 0.333 0.333 0.000 0.000 0.000
  0.250 0.250 0.250 0.250 0.000 0.000
  0.200 0.200 0.200 0.200 0.200 0.000
  0.167 0.167 0.167 0.167 0.167 0.167

Три именованных компонента transformer теперь на экране. Треугольник — это causal mask, вынужденная целью: если бы позиция tt могла видеть позицию t+1t{+}1, ответ был бы во входе — утечка, которую глава 6 велела проверять, только теперь внутри архитектуры. softmax — это способ реализации маски: установка запрещённых элементов в -\infty отправляет их ровно в ноль и нормализует оставшееся, так что маскирование и нормализация становятся одной операцией. (Используйте -\infty, а не -1e9: это значение, которое маска означает, оно переживает приведение к float16 как -\infty и избавляет вас от решения, достаточно ли велика выбранная константа для диапазона, в котором вы случайно оказались, — а это floating-point-рамка из главы 2, задающая вопрос, на который вам не нужно отвечать.) А scores — свободный параметр. Равномерное среднее получается, когда каждый разрешённый score — одно и то же число; поставьте туда любые числа, и softmax превратит их в допустимые веса.

Остаток этой главы — один вопрос: откуда берутся эти числа?

Они не могут быть обычными параметрами. Обучаемая матрица T×TT \times T была бы одинаковой для каждого предложения — она могла бы кодировать «смотри на четыре tokens назад», но никогда «смотри на существительное, к которому относится это местоимение». Вес, связывающий позицию tt с позицией ii, должен зависеть от того, что находится в обеих позициях, потому что релевантность — это отношение, а не свойство: слово it не является релевантным само по себе, оно релевантно чему-то.

Самая дешёвая функция двух векторов, возвращающая число, — dot product из главы 1. Оцените позицию ii для позиции tt как xtxi\mathbf{x}_t \cdot \mathbf{x}_i, и механизм заработает — плохо, двумя способами, которые вынуждают всё остальное. Dot product вектора с самим собой — это квадрат его нормы, поэтому каждый token mostly attend к самому себе. И отношение было бы симметричным: если it сильно attends к animal, то animal сильно attends к it, что неверно для языка, где прилагательному существительное нужно куда больше, чем существительному прилагательное.

Поэтому дайте каждому token две роли, как два обучаемых линейных отображения от него: что эта позиция ищет, qt=Wqxt\mathbf{q}_t = W_q\mathbf{x}_t, query; и что она предлагает, чтобы её нашли, ki=Wkxi\mathbf{k}_i = W_k\mathbf{x}_i, key. Score qtki\mathbf{q}_t \cdot \mathbf{k}_i — и симметрия исчезает, потому что WqWkW_q \neq W_k: token может рекламировать одно, а искать другое.

Одна вещь всё ещё неправильна. Взвешенная сумма бралась по самим xi\mathbf{x}_i, что заставляет копируемую вещь быть той же вещью, по которой сопоставляют. Сопоставлению нужны признаки, идентифицирующие token; копированию — признаки, полезные ниже по сети. Поэтому выучим третье отображение, vi=Wvxi\mathbf{v}_i = W_v\mathbf{x}_i, value, и будем суммировать их.

Формула теперь — bookkeeping:

Attention(Q,K,V)=softmax ⁣(QKdk+M)V\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}} + M\right)V

где MM — causal mask: ноль на диагонали и ниже, -\infty выше. В коде это тридцать строк, двадцать из которых — shapes:

attention.pyPYTHON
class Head(nn.Module):
    """One head of causal self-attention."""

    def __init__(self, d_model, d_head, block):
        super().__init__()
        self.q = nn.Linear(d_model, d_head, bias=False)      
        self.k = nn.Linear(d_model, d_head, bias=False)      
        self.v = nn.Linear(d_model, d_head, bias=False)      
        self.d_head = d_head
        self.register_buffer("mask", torch.tril(torch.ones(block, block)).bool())

    def forward(self, x):
        T = x.shape[1]
        q, k, v = self.q(x), self.k(x), self.v(x)
        s = q @ k.transpose(-2, -1) / math.sqrt(self.d_head)          
        s = s.masked_fill(~self.mask[:T, :T], float("-inf"))          
        w = F.softmax(s, dim=-1)                                      
        return w @ v                                                  

Score, mask, normalize, mix. Всё остальное — projection.

Деление на квадратный корень и от чего оно защищает

Ссылка на раздел: Деление на квадратный корень и от чего оно защищает

Почти каждое объяснение dk\sqrt{d_k} говорит «чтобы softmax не насыщался», что верно и ничего не объясняет. Аргумент — две строки про дисперсию из главы 2. Если элементы q\mathbf{q} и k\mathbf{k} независимы, имеют среднее ноль и дисперсию один, каждое произведение qjkjq_j k_j имеет дисперсию один, а дисперсии независимых величин складываются:

Var(qk)=j=1dkVar(qjkj)=dk\mathrm{Var}(\mathbf{q}\cdot\mathbf{k}) = \sum_{j=1}^{d_k}\mathrm{Var}(q_j k_j) = d_k

Значит, scores имеют стандартное отклонение dk\sqrt{d_k}. Измерено на двадцати тысячах случайных пар:

TEXT
     d     Var(q.k)         std   sqrt(d)
     4        3.975       1.994     2.000
    16       16.071       4.009     4.000
    64       64.249       8.016     8.000
   256      253.065      15.908    16.000
  1024     1015.562      31.868    32.000

Почему это важно: softmax чувствителен к масштабу так, как линейный слой не чувствителен. Удвоение входа линейного слоя удваивает его выход; умножение scores на десять перед softmax превращает мягкую смесь в жёсткий выбор. Одна строка из 64 scores, с делением и без:

dkd_kнаибольший вес, без деленияэнтропияэффективные tokensнаибольший вес, с делениемэнтропияэффективные tokens
40.2052.94419.00.0813.75842.9
160.4381.6925.40.0753.84946.9
640.4890.8742.40.0853.67339.4
2560.99990.00071.00.1433.54734.7
10241.00000.00001.00.1323.64438.3

«Эффективные tokens» — это экспонента энтропии: по скольким позициям строка реально усредняет. Без деления, при dk=256d_k = 256, только что инициализированная head attends ровно к одному token из 64, выбранному исключительно случайной инициализацией.

Это плохо на forward и ещё хуже на backward, в форме, которую глава 5 уже измеряла на tanh\tanh. Softmax, зафиксировавшийся на одном элементе, почти не имеет производной: диагональ его якобиана — wi(1wi)w_i(1-w_i), ноль на обоих концах. По двум тысячам случайных строк:

dkd_kiwi(1wi)\sum_i w_i(1-w_i) без деленияс делениемнасыщенные строки (наибольший вес выше 0.99)
40.84270.95680.2 % → 0.0 %
640.29400.960917.9 % → 0.0 %
2560.14060.960949.1 % → 0.0 %
10240.06810.961170.4 % → 0.0 %

При dk=1024d_k = 1024 семь строк из десяти заморожены ещё до начала обучения, а head, которая начинает замороженной, не может научиться, куда смотреть. С делением величина ровно держится около 0,96 при любой ширине, и ничто не насыщается.

Теперь часть, которую никто не публикует: меняет ли это итоговую перплексию? Удалим деление и обучим при четырёх ширинах head:

ширина headбез деленияделение на dk\sqrt{d_k}деление на dkd_k
четыре heads, dk=32d_k = 3237.2938.0737.89
одна head, dk=128d_k = 12848.5146.1045.99
одна head, dk=256d_k = 25665.3747.53
одна head, dk=512d_k = 51267.0649.15
одна head, dk=1024d_k = 102476.6959.17

Первые две строки взяты из бюджета 3000 шагов выше; последние три — более короткий прогон: 1500 шагов, batch 32, одна head, без нормализации перед projections, при идентичных настройках обеих вариантов.

При dk=32d_k = 32 деление ничего не стоит, и прогон без него совсем немного впереди. Это не лицензия его убрать, потому что при 256 оно стоит 18 пунктов перплексии, а при 1024 — 17. Механизм виден в самих scores:

dkd_kstd scores при initпосле 1500 шагов, без деленияпосле 1500 шагов, с делениемнасыщенные строки, без деленияс делением
25610.49121.672.1391.9 %0.8 %
51215.13836.852.6698.7 %1.3 %
102421.155147.463.4499.9 %16.5 %

Head без деления не восстанавливается. Она убегает: стандартное отклонение её scores растёт с 21 при инициализации до 5147, entropy attention падает до нуля, и 99,9 % строк кладут больше 0,99 своего веса на один token. Как только head стала жёстким selector, её gradient почти нулевой, и ничто не тянет её назад, поэтому коллапс стабилен. Head с делением после того же обучения сидит на стандартном отклонении scores 3,44 — это мягкая смесь, которую всё ещё можно изменить.

Vaswani et al. говорят ровно это и не больше: они подозревают, что произведения «становятся большими по модулю при больших значениях dkd_k», и делят.5 Слово большими несёт нагрузку, а таблицы показывают, где начинается «большое»: ничего при 32, всё к 256.

Больше одного мнения и две трети, о которых никто не говорит

Ссылка на раздел: Больше одного мнения и две трети, о которых никто не говорит

Одна head — это одна строка softmax на позицию, то есть один ответ на вопрос «что здесь релевантно». Предсказание слова после the в the animal that crossed the wet street требует одновременно синтаксического слота, подлежащего и предыдущего token, а одно распределение вероятностей не может быть сосредоточено в трёх местах. Поэтому запускаем несколько heads параллельно, каждая шириной dmodel/hd_{\text{model}}/h, конкатенируем и смешиваем ещё одной матрицей WoW_o: вы разделили ширину, а не добавили к ней.

Attention также делает ровно одну вещь — переносит информацию между позициями. Каждая операция в коде выше линейна вдоль оси признаков, а глава 5 доказала, чем является стек линейных отображений. Поэтому каждый block также содержит небольшой MLP, применяемый к каждой позиции независимо, расширяющий ширину в четыре раза и возвращающийся обратно, с GELU посередине. Разделение труда стоит запомнить: attention смешивает между позициями, feed-forward network вычисляет внутри позиции.

Полная лестница, где каждая строка добавляет один элемент к строке над ней:

модельпараметрывалидационная перплексия
равномерное среднее, добавляется279,55260.45
одна attention head, заменяет token328,70455.47
одна attention head, добавляется328,70446.10
четыре heads вместо одной345,21643.21
плюс feed-forward network476,92839.87
плюс LayerNorm — полный block477,69638.07

Обучаемые веса бьют равномерные на 14 пунктов перплексии — весь аргумент этой главы в одной строке. Четыре heads дают ещё 3 за дополнительные 16 512 параметров. И та же head приносит на 9 пунктов больше, когда добавляется, чем когда заменяет: attention приносит информацию внутрь, но не решает, чем является позиция.

Теперь — где на самом деле находятся параметры, что удивляет людей, видевших только диаграмму:

ширинаheadsattentionfeed-forwardвсего на block
128465,664 (33.2 %)131,712 (66.6 %)197,888
768122,360,064 (33.3 %)4,722,432 (66.6 %)7,085,568
40963267,112,960 (33.3 %)134,238,208 (66.7 %)201,367,552

Две трети каждого transformer block — это feed-forward network, в любом масштабе, потому что у attention четыре матрицы d×dd \times d, а у MLP — эквивалент восьми. Что бы модель ни знала, большая часть параметров, которые это держат, находится в per-position MLP.

Residuals и LayerNorm, унаследованные из главы 6

Ссылка на раздел: Residuals и LayerNorm, унаследованные из главы 6

LayerNorm был построен и измерен в главе 6, и эта глава использует его в том виде, в каком он там остался; residual connections были там названы и ablated, и здесь они встроены. Строки «добавляется, а не заменяет» выше — это residual connections, стоящие 188 пунктов перплексии для среднего и 9 для одной head. LayerNorm7 нормализует каждый пример по его признакам, а глава 6 дала причины, почему здесь выжил именно он, а не BatchNorm: нет зависимости от batch, нет running statistics, одинаков в обучении и inference, равнодушен к длине последовательности — всё это становится требованием, когда вы генерируете по одному token за раз для одного пользователя, к чему в итоге приходит глава 13. Он стоит 768 параметров и покупает 1,8 пункта перплексии.

block.pyPYTHON
class Block(nn.Module):
    def forward(self, x):
        x = x + self.att(self.ln1(x))     
        x = x + self.ff(self.ln2(x))      
        return x

Посмотрите, где стоит нормализация: на входе каждого sub-layer, при этом residual path от входа к выходу никогда не нормализуется. Это pre-norm. Статья 2017 года делает наоборот, x = LayerNorm(x + Att(x))post-norm, где LayerNorm ставится на сам residual path.

Xiong et al. объяснили разницу через gradient при инициализации, который в post-norm-сети плохо масштабирован с глубиной — поэтому исходному transformer вообще требовался warmup learning rate, чтобы обучаться.8 Двенадцать blocks, 1000 шагов, learning rate 3×1033 \times 10^{-3}:

TEXT
gradient norm per block at initialisation, before any step
  pre-norm    block 1 0.0498 ... block 12 0.0657   ratio last/first  1.32
  post-norm   block 1 0.0977 ... block 12 0.1613   ratio last/first  1.65

  pre-norm,  no warmup          perplexity   37.82
  pre-norm,  200-step warmup    perplexity   37.62
  post-norm, no warmup          perplexity  308.05
  post-norm, 200-step warmup    perplexity   37.88

Post-norm без warmup в восемь раз хуже, а post-norm с warmup в точности совпадает с pre-norm. Warmup здесь не общая хорошая практика; это заплатка для конкретного расположения нормализации, и перенос LayerNorm устраняет необходимость в ней. Поэтому практически каждая модель с 2019 года — pre-norm, а диаграмму 2017 года следует читать как историю, а не как спецификацию.

Удалите position embeddings, и модель всё равно обучается; она просто не может сказать, где что находится, и это симметрия, а не сбой обучения. Ничто в attention score не упоминает сами tt или ii, поэтому перестановка входа переставляет выход: self-attention эквивариантен к перестановкам. Это слепота среднего к порядку в лучшей маскировке — causal mask восстанавливает часть порядка, потому что каждая позиция видит разный префикс, но внутри префикса все порядки одинаковы.

Четыре способа внедрить позицию, обученные на 64-token окнах и оценённые при 64, 128 и 256 — за пределами любой длины, которую они видели:

позицииперплексия при 64при 128при 256
совсем нет48.7952.6357.52
обучаемые absolute embeddings38.63108.47181.94
фиксированные синусоиды42.9695.26152.25
RoPE44.1250.5284.84
ALiBi44.9543.5142.49

Обучаемые absolute embeddings — один вектор на позицию, добавленный к token, — выигрывают на обученной длине, а затем падают с обрыва, потому что позиция 100 никогда не была в batch, и её embedding всё ещё случайный вектор, с которого он начал. Синусоиды, исходный выбор, вычисляются, а не обучаются, из синусов и косинусов с геометрически разнесёнными частотами; статья 2017 года надеялась, что это будет экстраполировать, а таблица говорит, что нет: функция определена в позиции 200, но модель никогда не училась читать её там. RoPE9 ничего не добавляет, а вместо этого вращает query и key на угол, пропорциональный позиции, в двумерных срезах; поскольку одинаковое вращение обеих сторон dot product оставляет его неизменным, score в итоге зависит только от tit - i, так что позиция бесплатно становится относительной и таблица не может закончиться. Оно деградирует, но именно деградирует. ALiBi10 — самый простой и самый странный результат здесь: линейный штраф на score, пропорциональный расстоянию, с разным наклоном для каждой head. Его перплексия улучшается, когда окно растёт за пределы обученной длины, с 44,95 до 42,49, потому что штраф определён на любом расстоянии, и каждая head продолжает делать то, чему её учили.

Урок переживает таблицу: архитектура, которая не может представить что-то, — это другая проблема, чем архитектура, которая никогда не учила этот диапазон, и кусается именно вторая. Это также механизм за каждым объявлением «мы расширили контекст до 128K»: почти всегда это re-scaling rotary encoding, и именно поэтому глава 16 говорит, что предел контекста движется, а не исчезает.

Dropout наследуется так же: он появляется на attention weights после softmax, на выходе каждого sub-layer перед residual addition и на сумме embeddings, делая ровно то, что описывала глава 6. В больших pretraining-прогонах его часто ставят в ноль, потому что модель, которая видит каждый token один раз, не в положении overfit.

Два тензора в слое имеют форму n×nn \times n, где nn — число tokens: scores и weights после softmax. Всё остальное — каждая projection, весь MLP — линейно по nn.

Один attention layer, ширина 512, 8 heads, batch один, float32, на laptop GPU. Две колонки с миллисекундами читайте только по их соотношениям: это wall clock на 8 GB laptop-карте, которая throttles с 1 785 MHz до ниже 300 MHz при нагреве, так что холодный прогон того же кода возвращается в семь-десять раз быстрее, а занятый — ещё медленнее. Колонки мегабайт — это byte counts allocator, они не меняются.

TEXT
  tokens   ms total    ms x4   ms projections   attn matrix MB    peak MB    MB x4
     128      2.246        -            1.324              0.5       14.6        -
     256      2.855     1.27            2.113              2.0       19.2     1.31
     512      5.761     2.02            3.105              8.0       34.4     1.79
    1024     16.414     2.85            4.008             32.0       89.1     2.59
    2048     51.573     3.14            9.989            128.0      296.1     3.32
    4096    225.432     4.37           20.176            512.0     1100.1     3.72
    8192    832.838     3.69           40.106           2048.0     4300.1     3.91
   16384   OUT OF MEMORY                                 8192.0

fitted exponent (log-log slope, last four rows):  time ~ n^1.91   memory ~ n^1.87

Колонки x4 — отношение к строке выше, и удвоение nn сходится ровно к 4 и для времени, и для памяти: 3,91 на последнем шаге против теоретических 4. Колонка projections — контроль: 4,0 ms при 1024 tokens до 40,1 ms при 8192, фактор десять при факторе восемь. Линейно, как и обещано.

А затем последняя строка. Один attention layer, одна последовательность, без модели вокруг него, исчерпывает память на 8 GB GPU при 16 384 tokens — одна только матрица scores была бы 8 GB: 8 heads × 16 384 × 16 384 × 4 bytes. Не модель; один промежуточный тензор в одном слое.

Это физический факт под тремя последующими главами. Поэтому context window вообще имеет предел, который глава 16 превращает в цену. Поэтому существует FlashAttention, вычисляющий тот же результат tiles, никогда не сохраняя матрицу, — memory optimization прежде, чем speed optimization.11 И это арифметика за ценой длинного prompt, которую глава 24 платит в agent loop, — отдельный вопрос от другого вывода той главы: модель ещё и хуже использует длинный контекст, что она измеряет и отказывается сваливать на эту формулу.

Показать детали

Два варианта, уменьшающих cache, названы здесь и оплачены в главе 13.

Генерация кэширует keys и values уже обработанных tokens — один key и одно value на token, на head на layer. Multi-query attention12 сохраняет hh query projections, но использует одну key projection и одну value projection, общие для всех heads, деля этот cache на hh. Grouped-query attention13 интерполирует: heads группируются, каждая группа делит один key и одно value, так что g=hg = h — ordinary attention, а g=1g = 1 — multi-query. Почти каждая открытая модель с 2023 года использует это с 4 или 8 группами. Ни один вариант не существует ради качества; оба существуют ради размера этого cache, и глава 13 делает арифметику, превращающую это в «какая модель помещается в ваш GPU».

Статья 2017 года описывает encoder-decoder: один stack читает source с unmasked attention, второй каузально генерирует target, а посередине третий вид attention, где queries decoder встречают keys encoder. Это правильно для перевода, где вход и выход — две последовательности.

Победила decoder-only половина — один stack, causal повсюду, вход и выход в одной последовательности, — и причина не в элегантности. «Предскажи следующий token» работает на любом тексте, поэтому training set — это internet, а не parallel corpus, и всё становится этой одной задачей: перевод — документ, содержащий source, затем target; вопрос и ответ — документ; разговор с tool call посередине — документ. Глава 11 о том, как производится последний. Encoders не исчезли: один видит весь вход сразу, что нужно, когда задача — представить текст, а не продолжить его, и именно поэтому retrieval embeddings в главе 19 приходят из encoders, а не из модели, которая ведёт чат.

Когда block определён, размер модели — арифметика. На block, при ширине dd и четырёхкратном расширении: 4d2+4d4d^2 + 4d для Wq,Wk,Wv,WoW_q, W_k, W_v, W_o с biases на всех четырёх, как в GPT-2, — таблица выше опускает bias у трёх из них, отсюда на 2 304 меньше на block при d=768d = 768; 8d2+5d8d^2 + 5d для MLP; 4d4d для двух LayerNorms — 12d2+13d12d^2 + 13d, плюс token table V×dV \times d и, для absolute positions, nctx×dn_{\text{ctx}} \times d. Для формы GPT-2 small — d=768d = 768, 12 blocks, vocabulary 50 257, context 1024, output layer делит weights embedding:

TEXT
  token embeddings     50,257 x 768 = 38,597,376
  position embeddings   1,024 x 768 =    786,432
  one block                             7,087,872
  12 blocks                            85,054,464
  final LayerNorm         2 x 768 =        1,536
  total (weights tied)                124,439,808

Это и есть опубликованный размер этой модели. Формула не приближение; это модель. Заметьте также, что почти треть маленькой модели — embedding table, поэтому размер словаря — архитектурное решение, а не preprocessing-решение: trade-off, который подготовила глава 7.

Перплексия — число о корпусе. Что делает одна head — другой вопрос, и модель, обученная на мегабайте Shakespeare, неправильный инструмент для него: честно о карте attention 500 000-параметрической модели можно сказать, что она в основном неинтерпретируема. Поэтому: язык, где у вопроса есть правильный ответ.

Классическая иллюстрация — the animal did not cross the street because it was too tired, где it — animal, против …because it was too wet, где одно слово переносит референт к street. Это Winograd schemas14 — пары предложений, отличающиеся одним словом, где это слово решает, к чему относится местоимение.

Они также решаемы читерством, и именно эту часть tutorials пропускают. Если два кандидата — animal и place, tired и wet идентифицируют референт по категории, и модель, которая знает только, какие слова присутствуют, отвечает правильно, ничего не зная о порядке. Измерено на такой версии задачи, с отложенными парами animal/place:

TEXT
uniform causal average           held-out referent accuracy 100.0 %
one transformer block            held-out referent accuracy  91.7 %

Bag of words бьёт transformer. Любая демонстрация, построенная на этом предложении, ничего не доказывает об attention.

Закроем дыру: берём обоих кандидатов из одного набора шестнадцати nouns, каждый из которых может стоять в любом слоте, и делим adjectives по роли, а не категории — четыре делают it тем, кто пересекает (tired, scared, slow, weak), четыре — тем, что пересекают (wet, wide, busy, steep).

TEXT
the {x} did not cross the {y} because it was too {adj} , so the {ref} waited .

Обучаем как обычный next-token predictor, оцениваем одну позицию — слово после so the — и строим held-out set из пар nouns, чей обратный порядок был в обучении, так что всё, что знает, какие два nouns присутствуют, но не знает, какой был первым, должно ответить наоборот.

модельпараметрыheld-outназывает другое noun
только текущий token5,7965.2 %5.2 %
равномерное causal среднее5,79627.9 %50.0 %
одна head learned attention18,08435.4 %64.6 %
четыре heads22,24475.0 %15.6 %
один transformer block55,71692.7 %4.2 %
два transformer blocks105,508100.0 %0.0 %

Случайный выбор среди двух присутствующих nouns — 50 %. Равномерное среднее попадает на 27,9 % и отвечает не тем noun из пары ровно в половине случаев — подпись механизма, который знает, какие слова есть, и ничего об их порядке, как предсказал shuffle test три раздела назад.

Теперь карта: attention в позиции, которая должна назвать референт, усреднённое по четырём heads каждого block, для двух предложений, отличающихся одним словом. Равномерное среднее положило бы 0,067 на каждый из пятнадцати видимых tokens.

TEXT
the animal did not cross the street because it was too tired , so the animal waited .
  blk 1  the:0.00 animal:0.70 did:0.00 not:0.00 cross:0.00 the:0.00 street:0.06
         because:0.00 it:0.00 was:0.00 too:0.00 tired:0.00 ,:0.05 so:0.00 the:0.19
  blk 2  the:0.00 animal:0.00 did:0.00 not:0.00 cross:0.00 the:0.00 street:0.00
         because:0.00 it:0.00 was:0.00 too:0.00 tired:1.00 ,:0.00 so:0.00 the:0.00

the animal did not cross the street because it was too wet , so the street waited .
  blk 1  the:0.00 animal:0.70 did:0.00 not:0.00 cross:0.00 the:0.00 street:0.06
         because:0.00 it:0.00 was:0.00 too:0.00   wet:0.00 ,:0.05 so:0.00 the:0.19
  blk 2  the:0.00 animal:0.00 did:0.00 not:0.00 cross:0.03 the:0.00 street:0.49
         because:0.00 it:0.00 was:0.00 too:0.20   wet:0.03 ,:0.00 so:0.00 the:0.25

Block 1 одинаков в обоих предложениях — 0,70 на первом noun, каким бы ни было adjective. Это не провал, а доказательство: в первом слое query в позиции является функцией собственного token этой позиции и индекса, а the в позиции 14 — один и тот же token в обоих предложениях. Head первого слоя не может зависеть от слова, которое она ещё не fetched. Поэтому block 1 делает единственную полезную доступную вещь и тянет первый noun вперёд.

Block 2 — место, где предложения расходятся, и одна и та же строка по всем восьми adjectives показывает правило, которое модель нашла:

adjectiveblock 2 на animalна streetна adjectiveответ
tired, scared, slow, weak0.0000.0001.000animal
wet, wide, busy, steep0.0000.4910.00–0.03street

Для crosser-adjective второй block тратит весь свой вес на adjective, потому что ответ уже в residual stream — block 1 положил его туда, — и ему нужно только подтверждение. Для crossed-adjective он идёт и fetches другое noun. Это two-hop circuit: одна head переносит кандидата вперёд, head в более позднем слое читает token, который решает, оставить его или нет. Composition across layers — механизм, и именно поэтому один block достиг 92,7 %, а два — 100 %.

Это также форма самого хорошо документированного circuit в реальных моделях. Induction heads — previous-token head, питающая head в следующем слое, которая завершает паттерн [A][B] … [A] → [B], — это то, что interpretability-работа Anthropic находит за большой частью in-context learning, и они формируются в определимый момент pretraining. Эта глава не пытается делать такой анализ: он делегирован, с обеими статьями в ссылках, потому что чтение circuits из реальной модели — исследовательская область, а не раздел.

Наконец, реализация. Тридцать строк выше, с weights, скопированными из собственного PyTorch:

TEXT
ours vs nn.MultiheadAttention           max |diff| = 1.7881393432617188e-07
ours vs F.scaled_dot_product_attention  max |diff| = 1.7881393432617188e-07

1.8×1071.8 \times 10^{-7} на outputs, средняя величина которых 0,159: та же арифметика в другом порядке, с точностью float32.

У вас есть архитектура, из которой построена каждая модель в оставшейся части курса, и она меньше своей репутации: взвешенное среднее с обучаемыми весами, per-position MLP, держащий две трети параметров, две нормализации и два сложения, сложенные в stack.

Чего у вас нет — модели, которая что-то знает, и само по себе stacking этого не исправит. Два blocks на этом корпусе достигают training perplexity 14,49 и validation perplexity 40,57 против 18,77 и 38,07 у одного block: больше capacity, лучше на том, что модель видела, хуже на том, чего не видела, — таблица главы 6, только с transformer внутри. Расстояние между этой моделью и теми, с которыми говорят главы 14–30, не архитектурное. Это тот же block, повторённый больше раз, на несравнимо большем тексте.

Значит, это бухгалтерская задача, и бухгалтерия страннее, чем кажется. Сколько текста, и где его вообще берут? Сколько арифметики, и как оценить её до того, как деньги потрачены? При фиксированном бюджете лучше сделать модель больше или показать ей больше данных — и есть ли правильный ответ, или только мода? Глава 10 отвечает на все три вопроса измерениями и назначает цену самой дешёвой полезной форме вопроса: сколько сегодня стоит обучить модель вроде GPT-2 с нуля?


Три объяснения этого материала лучше этого в своих задачах, и эта глава написана так, чтобы её читали рядом с ними. The Illustrated Transformer Jay Alammar — лучшая когда-либо нарисованная картина потока данных. The Annotated Transformer Harvard NLP — статья 2017 года с исполняемым кодом, вставленным строка за строкой. Let's build GPT: from scratch, in code, spelled out Andrej Karpathy строит ту же модель вживую за два часа, а лестница ablations выше — тот же позвоночник, измеренный на другом корпусе. Для вопроса interpretability, которого эта глава только касается, первоисточники — Elhage et al., A Mathematical Framework for Transformer Circuits (2021) и Olsson et al., In-context Learning and Induction Heads (2022), обе из interpretability group Anthropic.

  1. Hochreiter, S. and Schmidhuber, J. Long Short-Term Memory. Neural Computation 9(8), pp. 1735–1780 (1997).

  2. Sutskever, I., Vinyals, O. and Le, Q. V. Sequence to Sequence Learning with Neural Networks. arXiv:1409.3215 (2014). Encoder-decoder, чей единственный context vector является узким местом.

  3. Bahdanau, D., Cho, K. and Bengio, Y. Neural Machine Translation by Jointly Learning to Align and Translate. arXiv:1409.0473 (2014). Attention за три года до transformer.

  4. Перплексия — экспонента средней cross-entropy на token, из главы 8. Каждое число здесь использует один и тот же tokenizer и одно и то же validation split; только при этом условии две перплексии вообще можно сравнивать.

  5. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł. and Polosukhin, I. Attention Is All You Need. arXiv:1706.03762 (2017). Section 3.2.1 — одно предложение о dk\sqrt{d_k}, измерению которого эта глава посвящает целый раздел.

  6. Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G. and Dean, J. Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. arXiv:1701.06538 (2017).

  7. Ba, J. L., Kiros, J. R. and Hinton, G. E. Layer Normalization. arXiv:1607.06450 (2016). Введён и измерен в главе 6; здесь используется без изменений.

  8. Xiong, R., Yang, Y., He, D., Zheng, K., Zheng, S., Xing, C., Zhang, H., Lan, Y., Wang, L. and Liu, T.-Y. On Layer Normalization in the Transformer Architecture. arXiv:2002.04745 (2020). Gradient-анализ за pre-norm и аргумент, что warmup — симптом.

  9. Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B. and Liu, Y. RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv:2104.09864 (2021).

  10. Press, O., Smith, N. A. and Lewis, M. Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation. arXiv:2108.12409 (2021). Результат экстраполяции, воспроизведённый выше.

  11. Dao, T., Fu, D. Y., Ermon, S., Rudra, A. and Ré, C. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135 (2022).

  12. Shazeer, N. Fast Transformer Decoding: One Write-Head is All You Need. arXiv:1911.02150 (2019).

  13. Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebrón, F. and Sanghai, S. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. arXiv:2305.13245 (2023).

  14. Levesque, H. J., Davis, E. and Morgenstern, L. The Winograd Schema Challenge. KR (2012). Конструкция за предложением animal / street, которое использует каждый tutorial по attention.

Готовы доверить выбор модели LIA?

Создавайте со всеми ИИ-моделями в одном месте — начните бесплатно уже сегодня.