Трансформер

Все детали у нас на руках. Токены — из главы 6, эмбеддинги — из пятой, маленькие нейросети — из четвёртой, внимание — из седьмой, а учить всё это мы умеем спуском с горы из третьей. Осталось всё это собрать. Слово «трансформер» звучит внушительно, но весь Росток — это один файл model.py на 174 строки, и сама архитектура занимает в нём чуть больше сотни. В этой главе мы прочитаем их все.

Сборка держится на трёх идеях: остаточном потоке, в котором слои ничего не переписывают, а только дописывают; нормировке, которая держит числа в разумных пределах; и позициях, которые внимание само по себе не различает. И ещё одна новая деталь — сеть SwiGLU, где живёт большая часть параметров. Начнём с общей картины.

Росток на одной схеме

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

Обратите внимание на переключатель: gpt-1 из прошлой главы, gpt-4 из этой и полный Росток — это один и тот же код. Меняются лишь несколько чисел: ширина потока, число блоков и голов, размер сети. Вся разница между маленькой и большой моделью — в этих числах, данных и времени обучения. Дальше пройдём по схеме сверху вниз.

Остаточный поток

Главная линия схемы — вертикальная. Токен входит в модель вектором из 384 чисел и таким же вектором выходит. По дороге ни один блок этот вектор не заменяет — каждый лишь прибавляет к нему свою поправку:

$$x \leftarrow x + \operatorname{Attn}\big(\operatorname{Norm}(x)\big),$$ $$x \leftarrow x + \operatorname{FFN}\big(\operatorname{Norm}(x)\big).$$

Этот вектор называют остаточным потоком (residual stream). Его удобно представлять как рабочую тетрадь токена: эмбеддинг пишет на первой странице, что это за токен; внимание дописывает, что удалось узнать у соседей; сеть — что из этого следует. Ничего не стирается, и любой следующий слой может прочитать всё, что записали предыдущие.

Почему не проще — $x \leftarrow f(x)$, как в сетях из главы 4? Причин две.

Градиенту нужна дорога. При обратном проходе (глава 4) производные слоёв перемножаются по цепочке. Если каждая из них хоть немного меньше единицы, через восемь блоков сигнал заметно ослабнет, а через сотню почти исчезнет. С прибавлением всё иначе:

$$\frac{\partial}{\partial x}\big(x + f(x)\big) = 1 + f'(x).$$

Единица в этой сумме — прямой провод, по которому градиент проходит сквозь все блоки, даже если сами $f'$ крошечные. Именно остаточные связи позволили в 2015 году обучить сеть ResNet из 152 слоёв (Хэ и др.), а сегодня — трансформеры из сотни блоков.

Легко начать с «ничего не делать». Нажмите на схеме «× 8 блоков» и посмотрите на строки инициализации: выходные матрицы внимания и сети (proj и w2) получают начальные веса в $\sqrt{2 \cdot 8} = 4$ раза мельче остальных. В начале обучения каждый блок добавляет почти ноль, поток почти равен эмбеддингу, и предсказание зависит только от текущего токена — модель работает как биграмма из первых глав. Дальше каждый блок постепенно учится вносить полезные поправки.

RMSNorm: держать числа в узде

У постоянного дописывания есть обратная сторона: размер вектора в потоке растёт. Мы измерили на gpt-4 из этой главы: типичное число в потоке (среднеквадратичное, в среднем по отложенному тексту) сразу после эмбеддинга — около 0,18, после первого блока — уже около 10, а после четвёртого — около 35, в двести раз больше. А слою удобнее всего, когда вход всегда примерно одного масштаба, где бы в модели этот слой ни стоял.

Поэтому перед каждым слоем вектор нормируют. Росток использует RMSNorm — среднеквадратичную нормировку:

$$\operatorname{RMS}(x) = \sqrt{\frac{1}{d}\sum_{i=1}^{d} x_i^2},$$ $$\hat x_i = \frac{x_i}{\sqrt{\operatorname{RMS}(x)^2 + \varepsilon}}\, g_i .$$

Каждое число делим на «типичный размер» вектора, так что у результата он становится равен единице, а затем умножаем на обучаемое усиление $g_i$ — модель сама решает, какие координаты сделать громче. $\varepsilon$ — крошечная добавка, чтобы не делить на ноль. Параметров у нормировки всего $d$ — по усилению на координату.

Обратите внимание, где стоит нормировка: на входе в слой, а не в самом потоке. Нормируется копия, которую слой читает, а поток остаётся нетронутым — иначе мы бы снова «переписывали тетрадь». Такую схему называют pre-norm. В самом первом трансформере 2017 года нормировка стояла после сложения (post-norm), и глубокие модели с ней обучались капризно; Сюн и др. (2020) показали, почему pre-norm стабильнее. Ещё одна нормировка стоит в самом конце, перед выходным слоем.

Каждый токен думает сам: SwiGLU

Внимание — единственное место, где токены обмениваются сведениями. Но собрать сведения мало, их надо переработать. Для этого в каждом блоке после внимания стоит маленькая нейросеть из двух слоёв, почти такая же, как в главе 4, — сеть прямого распространения (feed-forward network). Она обрабатывает каждый токен отдельно, одними и теми же весами на всех позициях. Здесь живёт большая часть параметров: $3 \times 384 \times 1\,024 = 1\,179\,648$ на блок против $589\,824$ у внимания — две трети блока.

В первом трансформере сеть была такой: $\operatorname{ReLU}(x W_1)\,W_2$ — расширить вектор в 4 раза, обнулить отрицательное, сжать обратно. Росток, как и Llama, использует вариант SwiGLU с тремя матрицами:

$$\operatorname{FFN}(x) = \big(\operatorname{silu}(x W_1) \odot x W_3\big)\,W_2,$$ $$\operatorname{silu}(a) = \frac{a}{1 + e^{-a}}.$$

Здесь $\odot$ — поэлементное умножение. Вектор расширяется до 1 024 чисел дважды, двумя разными матрицами. Первая копия проходит через SiLU — гладкую родственницу ReLU, — вторая остаётся как есть, и они перемножаются: одна копия служит воротами для другой. Нейрон может сказать «пропускаю, только если выполнены оба условия» или даже поменять знак сигнала — одиночной ReLU такое не под силу.

А откуда странное число 1 024? В классической сети скрытый слой в 4 раза шире потока, и матриц две: $2 \times 4d^2 = 8d^2$ параметров. У SwiGLU матриц три, поэтому ширину берут равной $\tfrac{2}{3}$ от обычной, чтобы параметров вышло столько же: $\tfrac23 \cdot 4 \cdot 384$ — ровно $1\,024$. Вот сеть Ростка целиком:

class FeedForward(nn.Module): """SwiGLU: silu(x·W1) ⊙ (x·W3), projected back with W2.""" def __init__(self, cfg): super().__init__() self.w1 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False) self.w3 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False) self.w2 = nn.Linear(cfg.d_ff, cfg.d_model, bias=False) def forward(self, x): return self.w2(F.silu(self.w1(x)) * self.w3(x))

Где я? Позиции поворотом

В конце прошлой главы выяснилось, что внимание не различает порядок слов: для него «the dog bit the boy» и «the boy bit the dog» — один и тот же набор векторов. Порядок надо как-то сообщить.

Самое прямолинейное решение — завести ещё одну таблицу эмбеддингов, на этот раз для позиций, и прибавлять её строку к вектору токена. Так сделано в GPT-2 (2019). Но тогда модель знает только абсолютные номера и отдельно учит, что «слово двумя позициями раньше» на десятом месте и на трёхсотом — одно и то же отношение. А для языка важнее как раз относительное: насколько далеко стоит нужное слово.

RoPE (rotary position embedding; Су и др., 2021) ничего не прибавляет — он поворачивает. Запрос одной головы — это 64 числа; разобьём их на 32 пары, и каждая пара станет точкой на плоскости. Токен на позиции $m$ поворачивает свою $i$-ю пару на угол $m\,\theta_i$, где

$$\theta_i = 10000^{-2i/64}, \qquad i = 0, 1, \dots, 31.$$

Ключи поворачиваются так же, каждый по своей позиции $n$. Значения не поворачивают.

Почему это работает? Поворот не меняет длину вектора, а угол между двумя повёрнутыми векторами равен старому углу плюс разность поворотов: $(m - n)\,\theta_i$. Скалярное произведение зависит только от длин и угла. Значит, оценка внимания зависит от позиций только через расстояние $m - n$, а не от того, где именно в тексте стоят слова.

Зачем 32 разные частоты? Пара $i = 0$ поворачивается на 1 радиан за позицию и делает полный оборот за 6,3 позиции — она чутко различает соседние слова, но на больших расстояниях «зацикливается». Пара $i = 31$ поворачивается так медленно, что на полный оборот ей нужно около 47 000 позиций, — на всём контексте Ростка она почти стоит на месте и различает «близко» и «далеко» лишь грубо. Все вместе они — как стрелки часов: секундная, минутная и часовая показывают время точно на любом масштабе.

В коде это две функции. Первая один раз готовит таблицы косинусов и синусов для всех позиций и частот. Вторая поворачивает пары $(x_i, x_{i+32})$: первая половина вектора — это «иксы» всех пар, вторая — «игреки».

def rope_tables(head_dim, context, base): """cos/sin for every position and every frequency, shape (context, head_dim/2).""" inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim)) angles = torch.outer(torch.arange(context).float(), inv_freq) return angles.cos(), angles.sin() def apply_rope(x, cos, sin): """Rotate pairs (x[i], x[i + d/2]) by a position-dependent angle. x: (batch, heads, time, head_dim); cos/sin: (time, head_dim/2). """ d = x.size(-1) // 2 x1, x2 = x[..., :d], x[..., d:] return torch.cat((x1 * cos - x2 * sin, x1 * sin + x2 * cos), dim=-1)

Проверим главное свойство на numpy. Одни и те же запрос и ключ на позициях 5 и 2 и на позициях 105 и 102 — расстояние одно, значит, и оценка должна совпасть:

import numpy as np d = 64 inv_freq = 1.0 / (10000 ** (np.arange(0, d, 2) / d)) # the 32 thetas def rope(x, pos): """Rotate pairs (x[i], x[i + 32]) by pos * theta_i.""" a = pos * inv_freq x1, x2 = x[:d // 2], x[d // 2:] return np.concatenate([x1 * np.cos(a) - x2 * np.sin(a), x1 * np.sin(a) + x2 * np.cos(a)]) rng = np.random.default_rng(1) q, k = rng.standard_normal(d), rng.standard_normal(d) for m, n in [(5, 2), (105, 102), (400, 397), (5, 4)]: s = rope(q, m) @ rope(k, n) print(f"m={m:3d} n={n:3d} distance {m - n}: q.k = {s:+.4f}") print("length kept:", np.isclose(np.linalg.norm(rope(q, 77)), np.linalg.norm(q)))

Запрос стоит на позиции 100, ключ — на позиции 90. Потом ту же пару векторов переносим на позиции 20 и 10. Что станет с оценкой внимания?

Упадёт: чем дальше от начала текста, тем больше поворот и тем меньше оценка Не изменится: оценка зависит только от расстояния, а оно в обоих случаях 10 Вырастет: у ранних позиций углы меньше, и векторы ближе друг к другу

Повороты на $m\theta$ и $n\theta$ входят в скалярное произведение только через разность $(m-n)\theta$. Это и проверила ячейка выше: пары позиций (5, 2), (105, 102) и (400, 397) дали одно и то же число.

У RoPE нет ни одного обучаемого параметра, так что модель от него не растёт, а поворот можно посчитать для любой позиции: таблица позиций никогда не кончится. (Правда, модель, обученная на 512 токенах, всё равно плохо читает тексты намного длиннее — чтобы растянуть RoPE, нужны отдельные трюки.) За эти качества RoPE и используют Llama, Mistral, Qwen и почти все современные открытые модели.

Блок целиком

Теперь у нас есть все части блока. Вот он. Весь его смысл умещается в двух строчках метода forward, а комментарии к ним пересказывают эту главу лучше любой схемы:

class Block(nn.Module): def __init__(self, cfg): super().__init__() self.norm1 = RMSNorm(cfg.d_model) self.attn = Attention(cfg) self.norm2 = RMSNorm(cfg.d_model) self.ffn = FeedForward(cfg) def forward(self, x, cos, sin): x = x + self.attn(self.norm1(x), cos, sin) # tokens exchange information x = x + self.ffn(self.norm2(x)) # each token thinks on its own return x

Восемь этажей и выход

Блоки ставятся друг на друга: у Ростка их восемь, одинаковых по устройству, но каждый со своими весами. Зачем больше одного? Потому что каждый следующий блок читает то, что записали предыдущие, и может на этом строить. В прошлой главе мы видели индукционную голову: ей нужно, чтобы кто-то раньше записал в каждый токен, какое слово стояло перед ним. Одному слою такое не под силу, а двум — уже да. Глубина позволяет строить цепочки таких шагов.

Проверим на деле. Ниже три модели читают один и тот же текст: однослойная gpt-1 из прошлой главы, четырёхблочная gpt-4 из этой и полный Росток.

Посмотрите на второе «Zorbin Plax». Однослойная gpt-1 и во второй раз не ждёт после «Z» продолжения «or»: вероятность меньше процента, как и в первый. А gpt-4 ставит на «or» 98%, на «bin» — 93%: четыре блока нашли в тексте, что уже шло после «Z», и скопировали это. Это та самая пара голов из прошлой главы — «на соседа» во втором слое и индукционная в третьем. Полный Росток с восемью блоками уверен ещё сильнее: 100% на «or», 98% на «bin», и «Plax» он тоже вспоминает.

После последнего блока поток ещё раз нормируется и попадает в выходной слой. Тот считает скалярное произведение вектора со строкой эмбеддинга каждого из 8 192 токенов — получаются логиты. И вот тонкость: в model.py это буквально та же матрица, что и на входе:

self.head.weight = self.embed.weight   # weight tying

Это называется связанными весами (weight tying; Пресс и Вольф, 2017). Смысл простой: на входе строка таблицы говорит «вот что значит этот токен», на выходе та же строка спрашивает «похоже ли то, что модель хочет сказать, на этот токен?». Одна таблица учится с двух сторон, а модель экономит $8\,192 \times 384 = 3\,145\,728$ параметров — без этого трюка Ростку понадобилось бы 20,45 млн чисел вместо 17,31 млн.

Дальше всё знакомо: softmax превращает логиты в вероятности, а при обучении перекрёстная энтропия с настоящим следующим токеном даёт ошибку (глава 2). Маска loss_mask понадобится, когда мы будем учить Ростка разговаривать (глава 12).

Считаем параметры

Теперь можно разобрать Ростка по косточкам и пересчитать всё до последнего числа. Обозначим ширину потока $d = 384$, число блоков $L = 8$, ширину сети $d_{ff} = 1\,024$ и размер словаря $V = 8\,192$:

• эмбеддинг (он же выходной слой): $V d = 3\,145\,728$;
• внимание в одном блоке: $W_Q, W_K, W_V, W_O$ — это $4d^2 = 589\,824$, плюс $2 \times 64$ усилений QK-norm;
• сеть SwiGLU в одном блоке: $3 d\, d_{ff} = 1\,179\,648$;
• две нормировки в блоке: $2d = 768$, и одна финальная: $384$.

$$3\,145\,728 + 8 \times (589\,824 + 128 + 1\,179\,648 + 768) + 384 = 17\,309\,056.$$

Это и есть 17,31 миллиона параметров Ростка. Посчитайте свою модель:

Пощёлкайте заготовки. У крошечной gpt-1 почти всё — эмбеддинг: 80% параметров уходит на словарь, и на «мышление» остаётся совсем немного. У Ростка на словарь приходится уже 18%, а у модели размером с GPT-2 small с её огромным словарём доля снова растёт. Видно и то, что цена блока растёт как квадрат ширины потока.

Ширину потока $d$ удвоили, $d_{ff}$ тоже удвоили (чтобы сохранить пропорцию), число блоков оставили. Во сколько раз вырастет число параметров внутри блоков?

В 2 раза Примерно в 4 раза В 8 раз

И внимание ($4d^2$), и сеть ($3d \cdot d_{ff}$, где $d_{ff}$ растёт вместе с $d$) — квадратичны по ширине. Удвоение $d$ даёт $2^2 = 4$; только крошечные усиления нормировок растут вдвое. Эмбеддинг $Vd$ за пределами блоков тоже вырастет лишь вдвое.

Логит-линза: мысли по дороге

Остаточный поток подсказывает изящный фокус. Все блоки пишут в один и тот же вектор, а выходной слой читает его в самом конце. Но что мешает прочитать поток раньше — после третьего блока, после пятого — той же финальной нормировкой и тем же выходным слоем? Получится ответ на вопрос «что сказала бы модель, если бы остановилась здесь?». Этот приём называют логит-линзой (logit lens; его придумал в 2020 году исследователь под ником nostalgebraist).

На что посмотреть:

• Строка «эмбеддинг» почти всюду повторяет сам входной токен. Это следствие связанных весов: вектор токена больше всего похож на свою же строку таблицы, и выходной слой находит именно её. Модель ещё ничего не «подумала».
• Где впервые появляется правильный ответ? Рамка отмечает совпадение с настоящим следующим словом. Для лёгких продолжений ответ часто созревает в средних слоях, для трудных — только в последних.
• Уверенность растёт в средних слоях, а последний слой часто отступает. Цвет клетки — вероятность лучшего варианта. К слоям 6–7 фаворит нередко получает 90% и больше, а после восьмого обычно скромнее: на трёх готовых текстах в среднем 77% после слоя 7 и 59% после слоя 8. При обучении оценивают только настоящий выход модели, а перекрёстная энтропия наказывает за самоуверенность, — вот последний блок и научился раскладывать ставку на несколько правдоподобных слов. Средние слои, прочитанные через линзу, звучат увереннее, чем имеют право. Нажмите на клетку, чтобы увидеть пять лучших вариантов.

Весь model.py

Вот файл полностью, от первой строки до последней. Теперь вам понятно в нём всё, кроме функции generate: выбору слов посвящена глава 11. Одно предупреждение: числа в Config — лишь значения по умолчанию, а настоящие размеры Ростка (поток из 384 чисел, 6 голов, сеть на 1 024) передаёт ему train.py.

"""Sprout: a tiny GPT, all of it in one file. A decoder-only transformer in the style of today's open models: - pre-norm blocks with RMSNorm, - rotary position embeddings (RoPE), - multi-head causal self-attention with QK-norm, - a SwiGLU feed-forward layer, - input and output embeddings shared (weight tying), - no biases anywhere. """ import math from dataclasses import dataclass, asdict import torch import torch.nn as nn import torch.nn.functional as F @dataclass class Config: vocab_size: int = 8192 # how many different tokens the model knows context: int = 512 # how many tokens it can look back at n_layer: int = 8 # transformer blocks stacked on top of each other n_head: int = 8 # attention heads per block d_model: int = 512 # width of the residual stream d_ff: int = 1408 # hidden width of the feed-forward layer rope_base: float = 10000.0 def to_dict(self): return asdict(self) class RMSNorm(nn.Module): """Rescale a vector to unit root-mean-square, then apply a learned gain.""" def __init__(self, dim, eps=1e-6): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x): return F.rms_norm(x.float(), (x.size(-1),), self.weight.float(), self.eps).type_as(x) def rope_tables(head_dim, context, base): """cos/sin for every position and every frequency, shape (context, head_dim/2).""" inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim)) angles = torch.outer(torch.arange(context).float(), inv_freq) return angles.cos(), angles.sin() def apply_rope(x, cos, sin): """Rotate pairs (x[i], x[i + d/2]) by a position-dependent angle. x: (batch, heads, time, head_dim); cos/sin: (time, head_dim/2). """ d = x.size(-1) // 2 x1, x2 = x[..., :d], x[..., d:] return torch.cat((x1 * cos - x2 * sin, x1 * sin + x2 * cos), dim=-1) class Attention(nn.Module): def __init__(self, cfg): super().__init__() self.n_head = cfg.n_head self.head_dim = cfg.d_model // cfg.n_head self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False) self.proj = nn.Linear(cfg.d_model, cfg.d_model, bias=False) self.q_norm = RMSNorm(self.head_dim) self.k_norm = RMSNorm(self.head_dim) def forward(self, x, cos, sin): B, T, C = x.shape q, k, v = self.qkv(x).split(C, dim=-1) # (B, T, C) -> (B, heads, T, head_dim) q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2) k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2) v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2) q, k = self.q_norm(q), self.k_norm(k) q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin) # softmax(q·k / sqrt(d)) · v, every token sees only itself and the past y = F.scaled_dot_product_attention(q, k, v, is_causal=True) y = y.transpose(1, 2).contiguous().view(B, T, C) return self.proj(y) class FeedForward(nn.Module): """SwiGLU: silu(x·W1) ⊙ (x·W3), projected back with W2.""" def __init__(self, cfg): super().__init__() self.w1 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False) self.w3 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False) self.w2 = nn.Linear(cfg.d_ff, cfg.d_model, bias=False) def forward(self, x): return self.w2(F.silu(self.w1(x)) * self.w3(x)) class Block(nn.Module): def __init__(self, cfg): super().__init__() self.norm1 = RMSNorm(cfg.d_model) self.attn = Attention(cfg) self.norm2 = RMSNorm(cfg.d_model) self.ffn = FeedForward(cfg) def forward(self, x, cos, sin): x = x + self.attn(self.norm1(x), cos, sin) # tokens exchange information x = x + self.ffn(self.norm2(x)) # each token thinks on its own return x class GPT(nn.Module): def __init__(self, cfg): super().__init__() self.cfg = cfg self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model) self.blocks = nn.ModuleList(Block(cfg) for _ in range(cfg.n_layer)) self.norm = RMSNorm(cfg.d_model) self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False) self.head.weight = self.embed.weight # weight tying cos, sin = rope_tables(cfg.d_model // cfg.n_head, cfg.context, cfg.rope_base) self.register_buffer('cos', cos, persistent=False) self.register_buffer('sin', sin, persistent=False) self.apply(self._init) # residual projections start small so every block begins close to "do nothing" for name, p in self.named_parameters(): if name.endswith('proj.weight') or name.endswith('w2.weight'): nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * cfg.n_layer)) @staticmethod def _init(m): if isinstance(m, (nn.Linear, nn.Embedding)): nn.init.normal_(m.weight, mean=0.0, std=0.02) def num_params(self): return sum(p.numel() for p in self.parameters()) def forward(self, idx, targets=None, loss_mask=None): T = idx.size(1) x = self.embed(idx) cos, sin = self.cos[:T], self.sin[:T] for block in self.blocks: x = block(x, cos, sin) logits = self.head(self.norm(x)) if targets is None: return logits, None loss = F.cross_entropy(logits.float().view(-1, logits.size(-1)), targets.reshape(-1), reduction='none') if loss_mask is None: return logits, loss.mean() mask = loss_mask.reshape(-1).float() return logits, (loss * mask).sum() / mask.sum().clamp(min=1) @torch.no_grad() def generate(self, idx, max_new, temperature=0.8, top_k=None, top_p=0.95, stop=None): for _ in range(max_new): logits, _ = self(idx[:, -self.cfg.context:]) logits = logits[:, -1, :].float() / max(temperature, 1e-5) if top_k: kth = torch.topk(logits, top_k).values[:, -1, None] logits[logits < kth] = -float('inf') probs = F.softmax(logits, dim=-1) if top_p and top_p < 1.0: sorted_p, order = probs.sort(descending=True) drop = sorted_p.cumsum(-1) - sorted_p > top_p sorted_p[drop] = 0 probs = torch.zeros_like(probs).scatter(-1, order, sorted_p) probs /= probs.sum(-1, keepdim=True) nxt = torch.multinomial(probs, 1) idx = torch.cat((idx, nxt), dim=1) if stop is not None and nxt.item() in stop: break return idx

Росток сейчас

Это уже Росток в миниатюре: тот же model.py, только поток шириной 192 числа вместо 384 и четыре блока вместо восьми. На отложенном тексте он ошибается в среднем на 2,26 ната на токен — против 2,87 у однослойной gpt-1 из прошлой главы. Разница кажется небольшой, но вспомните перплексию из главы 2: $e^{2{,}87} \approx 18$, а $e^{2{,}26} \approx 9{,}6$. Раньше модель в среднем колебалась между восемнадцатью вариантами следующего токена, теперь — между девятью с половиной.

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