Как модель выбирает слово

Обучение позади. Росток прочитал 330 миллионов токенов и теперь для любого начала умеет назвать вероятность каждого из своих 8 192 токенов. Но список вероятностей — ещё не текст. Кто-то должен решить, какой токен всё-таки будет написан, и, как ни странно, решает это вовсе не нейросеть. Решает короткий кусок кода, который работает уже после модели: весов у него нет, и он ничему не учится. Его называют сэмплером, а сам процесс — выборкой.

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

Пульт выборки

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

Ещё не тронув ни одной ручки, можно заметить две вещи. Во-первых, у распределения есть голова и очень длинный хвост. Горстка токенов забирает почти всю вероятность, а тысячи остальных делят крошки: на нижнем графике линия падает на много порядков от первого места до 8 192-го. Во-вторых, крошки набегают: под столбиками виджет складывает всё, что не попало в первую дюжину. После предсказуемого начала хвост тоненький, после открытого бывает увесистым — и тогда заметная доля бросков улетает куда-то в него.

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

Температура: одна ручка для смелости

Хочется одну ручку, которая плавно ведёт от «всегда бери фаворита» к «можно что угодно» и при этом не переставляет предпочтения модели: фаворит остаётся фаворитом, меняется только то, насколько сильно его предпочитают.

Вспомним, откуда берутся вероятности (глава 3): сеть выдаёт для каждого токена оценку — логит $z_i$, — а softmax превращает оценки в вероятности. Температура $T$ просто делит все логиты перед softmax:

$$p_i(T) = \frac{e^{z_i/T}}{\sum_{j} e^{z_j/T}}$$

Чтобы понять, что это даёт, сравним два токена. Знаменатели сокращаются:

$$\frac{p_i(T)}{p_j(T)} = \frac{e^{z_i/T}}{e^{z_j/T}} = e^{(z_i - z_j)/T}.$$

При $T = 1$ то же отношение равно $e^{z_i - z_j}$. Если возвести его в степень $1/T$, получится ровно строчка выше, то есть

$$\frac{p_i(T)}{p_j(T)} = \left(\frac{p_i(1)}{p_j(1)}\right)^{1/T}.$$

В этом весь секрет температуры: каждое отношение вероятностей возводится в степень $1/T$. Пусть Росток считает « girl» втрое вероятнее, чем « boy». При $T = 0{,}5$ степень равна 2, и « girl» становится вероятнее уже в $3^2 = 9$ раз. При $T = 2$ степень равна ½, и разрыв сжимается до $\sqrt{3} \approx 1{,}7$. Порядок мест не меняется никогда — меняются только зазоры между соседями.

Теперь две крайности. При $T \to 0$ каждое отношение возводится в огромную степень, и фаворит забирает всю вероятность себе: это жадный выбор, всегда верхний токен. При $T \to \infty$ каждое $z_i / T$ стремится к нулю, каждое $e^{z_i/T}$ — к $e^0 = 1$, и все 8 192 токена становятся равновероятными: чистый шум. Между ними — $T = 1$, ровно то распределение, которому модель научилась.

Как измерить, насколько распределение размазано? Инструмент у нас есть с главы 2 — энтропия $H = -\sum_i p_i \ln p_i$, и её легко себе представить. У честного кубика с $n$ гранями энтропия равна $\ln n$, так что распределение с энтропией $H$ так же неопределённо, как честный кубик с $e^H$ гранями. На пульте это число подписано «эффективных вариантов». Поставьте температуру 0,3 — и часто остаётся всего один-два настоящих варианта; поставьте 1,5 — и у кубика вырастают десятки, а то и сотни граней.

Вот одно и то же начало, дописанное четыре раза при четырёх температурах с одним и тем же зерном случайности:

Посмотрите на числа под колонками. Они показывают, насколько каждый текст удивляет самого Ростка, если дать ему перечитать написанное при обычной температуре: среднее $-\ln p$ на токен, та самая величина, которую он учился уменьшать. Холодный текст куда предсказуемее настоящих историй: люди не всегда говорят самое вероятное. Горячий текст ошарашивает даже ту модель, что его написала. Около $T = 1$ тексты примерно так же неожиданны, как истории, на которых учился Росток (на настоящих отложенных историях его ошибка — 1,55), — а на практике сэмплер часто держат чуть холоднее, жертвуя долей живости ради меньшего числа ляпов.

При T = 1 Росток считает « girl» в 4 раза вероятнее, чем « boy». Во сколько раз при T = 0,5?

В 2 раза В 4 раза: температура сохраняет отношения В 8 раз В 16 раз

Отношение возводится в степень 1/T = 2: 4² = 16. Низкая температура раздвигает все зазоры, высокая — сближает.

Жадный выбор ходит по кругу

Если модель знает, какой токен самый вероятный, почему бы всегда не брать его? Жадный выбор кажется самым надёжным. Глава 0 уже намекала на подвох; теперь рассмотрим его как следует. Ниже обе колонки пишут жадно, так что результат полностью определяется началом. Красным отмечен каждый токен, который повторяет фразу из четырёх токенов, уже встречавшуюся в тексте.

Почему модель, прочитавшая сотни миллионов слов, попадает в петли? Потому что повторы — настоящая закономерность языка. Имена возвращаются, припевы возвращаются, « the» возвращается каждые несколько слов, а у трансформера есть головы внимания, чья работа — найти, где текущий токен уже встречался, и предсказать то, что шло за ним в прошлый раз (их называют индукционными головами; подробно их изучили Олссон и коллеги в 2022 году). Стоит фразе прозвучать дважды, и этот механизм копирования делает третий раз вероятнее, а третий — четвёртый. Жадному выбору некуда свернуть, и текст застревает в петле, как игла в царапине на пластинке. Случайную выборку из петли вытаскивает везение: рано или поздно кубик выбросит что-то другое.

Зацикливается не каждое начало, и чем лучше обучена модель, тем реже это случается. Попробуйте третье начало: жадный Росток пишет целую историю про Лили и её большой красный мяч и нигде не застревает. Но утка и мышка показывают, что совсем склонность не исчезает, — это одна из причин, по которым чат-боты не выбирают жадно.

Штраф за повторы

Грубое, но действенное лекарство появилось вместе с моделью CTRL (Кескар и соавторы, 2019): перед выбором ослабить каждый токен, который уже есть в недавнем тексте. Наш движок смотрит на 64 последних токена и для каждого найденного там токена $i$ меняет логит с коэффициентом $r > 1$:

$$z_i \leftarrow \begin{cases} z_i / r, & z_i > 0, \\ z_i \cdot r, & z_i \le 0. \end{cases}$$

Обе ветки понижают оценку: положительный логит сжимается к нулю, отрицательный уходит ещё ниже. Подвигайте ползунок в виджете выше: небольшого штрафа обычно хватает, чтобы разорвать петлю. А потом поставьте 2 и посмотрите на обратную сторону. Штраф не отличает скучную петлю от нужного повтора: текст начинает избегать « the», точек и даже имени героя — и становится странным. Вот почему чат на этом сайте использует только мягкий штраф, $r = 1{,}1$.

Отрезаем хвост: top-k и top-p

Если жадный выбор слишком негибок, почему не тянуть из собственного распределения модели — $T = 1$, и больше ничего? Ведь именно ему она и училась. Беда в хвосте. Каждый токен хвоста по отдельности маловероятен, но их тысячи, а история длиной в сотни токенов. Пусть на каждом шаге хвосту принадлежит всего 3%. Шанс ни разу не попасть в него за 200 токенов равен $0{,}97^{200} \approx 0{,}002$. Иначе говоря, в случайно вытянутой истории почти наверняка найдётся несколько токенов из хвоста.

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

top-k

Первый рецепт (Фань, Льюис и Дофен, 2018): оставить $k$ самых вероятных токенов, остальным поставить ноль и растянуть выживших, чтобы в сумме снова получилась единица. При $k = 40$ хвоста больше нет. Но одно и то же $k$ не может подойти всем распределениям сразу. После «Once upon a» разумный токен по сути один — « time», а top-40 оставит рядом с ним ещё 39 лишних. После «For breakfast she ate» годятся десятки слов, а top-5 выбросит хорошие.

top-p, или ядро

Ари Хольцман и его соавторы (2019, статья вышла на ICLR 2020) предложили резать не по числу токенов, а по массе вероятности. Упорядочим токены от самого вероятного к наименее вероятному, $p_{(1)} \ge p_{(2)} \ge \dots$, и оставим самую короткую голову, чьи вероятности в сумме дают хотя бы $p$:

$$k^{*} = \min\Big\{k : \sum_{i=1}^{k} p_{(i)} \ge p\Big\}, \qquad p'_{(i)} = \frac{p_{(i)}}{\sum_{j=1}^{k^{*}} p_{(j)}} \ \text{ при } i \le k^{*}.$$

Эту голову авторы назвали ядром (nucleus). Оно подстраивается само: когда модель уверена, первый токен в одиночку добирает до $p$, и ядро сжимается до одного токена; когда правдоподобных продолжений много, оно разрастается до сотен. Проверьте на пульте: поставьте top-p = 0,9, переключайте готовые начала и смотрите, как скачет число выживших токенов. Затем переключите график на «накопленную сумму»: top-p — это горизонтальная линия на высоте $p$, и где её пересечёт кривая, там и граница, а top-k была бы вертикальной чертой на фиксированном месте.

Вот все три ручки в нескольких строчках numpy: выдуманное распределение из восьми слов и по десять тысяч бросков на каждую настройку. Запустите прямо здесь:

import numpy as np logits = np.array([3.0, 2.2, 1.5, 0.3, -0.5, -1.0, -2.0, -4.0]) words = ["girl", "boy", "dog", "cat", "bird", "king", "spoon", "the"] def next_token_probs(logits, temperature=1.0, top_k=0, top_p=1.0): z = logits / max(temperature, 1e-5) order = np.argsort(-z) # most likely first if top_k: z = z.copy() z[order[top_k:]] = -np.inf # everything below the k-th is gone p = np.exp(z - z.max()) p /= p.sum() if top_p < 1.0: sorted_p = p[order] drop = np.cumsum(sorted_p) - sorted_p > top_p p[order[drop]] = 0 # keep the head up to top_p p /= p.sum() return p settings = [(1.0, 0, 1.0), (0.5, 0, 1.0), (2.0, 0, 1.0), (1.0, 3, 1.0), (1.0, 0, 0.8)] rng = np.random.default_rng(0) for T, k, top_p in settings: p = next_token_probs(logits, T, k, top_p) draws = rng.choice(len(p), size=10_000, p=p) counts = np.bincount(draws, minlength=len(p)) shown = " ".join(f"{w} {c}" for w, c in zip(words, counts) if c) print(f"T={T} top_k={k} top_p={top_p}: {shown}")

Посмотрите на строку с top_p = 0.8: выжили только «girl» и «boy», потому что у «girl» одной около 55%, а вдвоём они переваливают за 80%. Строчка drop = np.cumsum(sorted_p) - sorted_p > top_p — изящный трюк: токен выбрасывается, только если токены перед ним уже набрали больше $p$, поэтому токен, который перешагнул черту, остаётся.

А вот настоящий цикл генерации из model.py Ростка. Порядок важен: сначала температура, потом top-k, потом softmax, потом top-p и, наконец, взвешенный бросок через torch.multinomial.

@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

Обратите внимание на одну расточительную строчку: self(idx[:, -self.cfg.context:]) прогоняет всю модель по всему тексту заново ради каждого нового токена. Коротко, правильно и для скрипта в терминале вполне годится. Движок в браузере устроен экономнее — скоро увидим как.

top-p = 0,9. Росток на 95% уверен, что следующий токен — « time». Из скольких токенов может выбирать сэмплер?

Из одного: « time» Из 90% словаря, это около 7 400 токенов Из первой десятки, какими бы ни были вероятности

Первый токен в одиночку уже набирает больше 0,9, так что ядро состоит из одного токена и этот шаг фактически жадный. Когда модель уверена, top-p позволяет ей быть уверенной; когда колеблется — ядро расширяется.

Откуда берётся случайность и где она кончается

«Вытянуть токен» — это тот же взвешенный кубик, что и в главе 1: берём случайное число $u$ от 0 до 1, идём по оставленным токенам, складывая их вероятности, и останавливаемся на первом, где сумма перевалила за $u$. Токену с вероятностью 0,3 принадлежит отрезок длиной 0,3, и попадают в него в 30% случаев.

Только компьютеры не бросают кубиков. Случайные числа выдаёт генератор псевдослучайных чисел: небольшая формула, которая превращает одно число — зерно (seed) — в длинную последовательность, похожую на случайную. В движке Ростка стоит 32-битный генератор такого рода. То же зерно, та же модель, те же ручки — и вы получите тот же текст, токен в токен. На этом держится виджет с четырьмя температурами: закрепите зерно, нажмите кнопку ещё раз — и каждая колонка повторится в точности; разрешите зерну меняться — и каждый запуск будет новым. Зерно делает эксперименты воспроизводимыми, а сравнения честными. А кнопка «сгенерировать заново» в чат-боте даёт другой ответ по простой причине: она берёт новое зерно.

Почему четыре колонки часто начинаются с одинаковых слов

Колонки, которые тянут жребий, берут случайные числа из одного потока (жадной колонке случайность не нужна вовсе). На первом шаге фавориту часто принадлежит большой отрезок, так что одно и то же $u$ попадает в него сразу в нескольких колонках. Стоит одной колонке выбрать другой токен — и тексты расходятся, дальше их уже ничто не связывает. Чем ниже температура, тем шире отрезок фаворита, поэтому холодные колонки дольше всех идут вместе.

А как генерация заканчивается? Одним из трёх способов. Модель может выдать служебный токен: при обучении каждая история заканчивалась меткой <|endoftext|>, базовая модель научилась ставить её в конце истории, и наш движок на ней останавливается. Мы можем назначить свои стоп-токены: в следующей главе чат-модель будет заканчивать каждую реплику меткой <|end|>. Или генерация упирается в ограничение длины, которое мы задали сами. Есть и жёсткий предел внутри модели: Росток видит не больше 512 токенов. Когда текст становится длиннее, наш движок оставляет последние 256 токенов, перечитывает их и продолжает, забыв начало.

Почему это быстро: кэш ключей и значений

Ещё раз посмотрим на generate: чтобы выбрать сто первый токен, она прогоняет модель по всем ста токенам; для сто второго — по ста одному, и так далее. Почти вся эта работа повторяется. Можно ли без неё обойтись?

Вспомним внимание (глава 7). Каждый токен в каждом блоке выдаёт запрос, ключ и значение. Из-за причинной маски ключ и значение токена $j$ зависят только от токенов $1 \dots j$. Новый токен в конце ничего не меняет в прошлом: все старые ключи и значения остаются в точности такими же. Значит, их можно посчитать один раз и сохранить. Это хранилище и называют кэшем ключей и значений, KV-кэшем. Для каждого нового токена модель считает только его собственные запрос, ключ и значение, кладёт ключ и значение в кэш, а запросом смотрит на всё сохранённое.

Сколько это экономит? Пусть в начале $P$ токенов, а сгенерировать нужно ещё $n$. Без кэша шаг, который выдаёт токен номер $t+1$, прогоняет через модель все $t$ токенов, и всего выходит

$$\sum_{t=P}^{P+n-1} t \;=\; nP + \frac{n(n-1)}{2}\ \text{ проходов токена},$$

а это растёт как $n^2$. С кэшем каждый токен проходит через модель один раз: $P + n$ проходов. Для начала из 14 токенов и 30 новых это 855 проходов против 44. Проход токена через Ростка стоит около 17,3 миллиона умножений со сложением (каждый вес используется один раз) плюс внимание по сохранённым токенам. Вот гонка — по-настоящему, прямо в вашем браузере:

Кэш не бесплатен: он занимает память. Для каждого токена Росток хранит ключ и значение по 384 числа в каждом из 8 блоков: $2 \times 8 \times 384 = 6\,144$ числа на токен. Полный контекст из 512 токенов — это 3,1 миллиона чисел, 12,6 МБ во float32, около трёх четвертей от 17,5 МБ, которые занимают сами веса. У больших моделей с длинным контекстом кэш занимает гигабайты, и есть целое направление исследований о том, как его ужать.

Как уместиться в браузер: восьмибитные веса и SIMD

У Ростка 17,31 миллиона параметров. В обычных 32-битных числах с плавающей точкой это 69,2 МБ — тяжеловато для веб-страницы. Файл, который загрузил ваш браузер, весит 17,5 МБ: каждая матрица весов хранится восьмибитными целыми числами с одним множителем на строку. Это называется квантованием.

Рецепт для одной строки матрицы: найти в ней наибольшее по модулю число и сопоставить ему 127. Знаковый байт вмещает целые от −128 до 127, и мы берём симметричную часть этого диапазона.

$$s_r = \frac{\max_c |W_{rc}|}{127}, \qquad q_{rc} = \operatorname{round}\!\left(\frac{W_{rc}}{s_r}\right) \in [-127, 127], \qquad W_{rc} \approx q_{rc}\, s_r.$$

Округление сдвигает каждый вес не больше чем на полшага, $s_r / 2$. Почему множитель на каждую строку, а не один на всю матрицу? Потому что веса бывают очень разными: один большой вес — выброс — растянул бы общий множитель, и все остальные числа матрицы округлялись бы грубо. С множителем на строку выброс портит только свою строку. Вот функция из export.py, которая упаковала Ростка для браузера:

def quantize_rows(w): scale = np.abs(w).max(axis=1) / 127.0 scale[scale == 0] = 1.0 q = np.clip(np.round(w / scale[:, None]), -127, 127).astype(np.int8) return q, scale.astype(np.float32)

Проверим её на случайной матрице 384 × 384 — такого размера выходной слой внимания у Ростка — с одним нарочно подложенным выбросом и сравним с единым множителем на всю матрицу:

rng = np.random.default_rng(1) W = rng.normal(0, 0.02, size=(384, 384)).astype(np.float32) W[7, 42] = 0.5 # one planted outlier q, s = quantize_rows(W) err_rows = np.abs(q * s[:, None] - W) s_all = np.abs(W).max() / 127.0 # one scale for the whole matrix err_all = np.abs(np.round(W / s_all) * s_all - W) print("bytes:", W.nbytes, "->", q.nbytes + s.nbytes) print(f"average error, a scale per row: {err_rows.mean():.1e}") print(f"average error, one scale for all: {err_all.mean():.1e}") print(f"in the row with the outlier: {err_rows[7].mean():.1e}")

Байтов почти вчетверо меньше, а с множителем на строку ошибка в разы меньше, чем с общим: расплачивается только строка с выбросом. На настоящей модели урон крошечный. Мы измерили его на контрольной точке из обучения Ростка (шаг 3 000 из 10 070): функция потерь на проверочных данных — 1,8388 ната во float32 и 1,8394 после округления до 8 бит с множителем на строку, разница около половины тысячной. На 8 битах не хуже справляется и один общий множитель на матрицу; множитель на строку окупается, когда битов становится мало. На 4 битах округление с множителем на строку стоит 0,06 ната, а с общим множителем на матрицу — полната.

В этом формате спрятан ещё один подарок. Когда строка умножается на входной вектор $x$, множитель не зависит от $c$ и выносится за сумму:

$$\sum_c W_{rc}\, x_c \;\approx\; \sum_c s_r\, q_{rc}\, x_c \;=\; s_r \sum_c q_{rc}\, x_c.$$

Поэтому движок умножает на вход маленькие целые числа и применяет один множитель на строку в самом конце. Эта внутренняя сумма и съедает почти всё время Ростка: 17,3 миллиона умножений со сложением на токен. Считает её крошечная функция на WebAssembly с SIMD (single instruction, multiple data — одна команда, много данных): она загружает сразу 16 весов в 128-битный регистр, расширяет их до чисел с плавающей точкой и умножает по четыре за раз. На Mac с M4 Pro (замер в Node) это ядро делает около 14 миллиардов умножений со сложением в секунду, а Росток пишет примерно 500 токенов в секунду — почти в шесть раз быстрее того же цикла на обычном JavaScript (около 90). Телефон в несколько раз медленнее, но всё равно успевает быстрее, чем вы читаете.

А насколько сильно можно сжать? Проверьте сами. Виджет ниже загружает веса Ростка в эту страницу и заново округляет все 17 миллионов чисел до меньшего числа бит, а потом даёт сжатой модели писать и измеряет её удивление на короткой истории, которую она никогда не видела.

У итогового Ростка шесть, пять и даже четыре бита на этой истории ничего заметного не портят: удивление держится около 1,77 ната на токен, а жадный текст почти не меняется. (На всём проверочном тексте контрольная точка шага 3 000 на 4 битах всё же теряла 0,06 ната — урон настоящий, просто маленький.) Три бита вредят сильно: удивление подскакивает больше чем на полната, и история теряет нить. А на двух битах, где каждый вес превращается в множитель строки, умноженный на −1, 0 или +1, модель разваливается: удивление взлетает примерно до 12 натов, а на контрольной точке шага 3 000 округление до 2 бит подняло функцию потерь на проверочных данных до 10,6. И то и другое хуже, чем равномерно гадать среди всех 8 192 токенов ($\ln 8192 \approx 9{,}0$).

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

Это та же базовая модель, что и в конце главы 10, но теперь вы управляете тем, как она выбирает. Убавьте температуру — получится спокойная, предсказуемая история, прибавьте — дикая; и помните, что творится на краях шкалы. Истории Росток пишет неплохо, но разговаривать пока не умеет: задайте ему вопрос, и он просто продолжит писать, будто ваш вопрос — строчка из рассказа. В следующей главе научим его разговаривать.