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

Росток закончил обучение: он прочитал 330 миллионов токенов и теперь оценивает возможные продолжения текста. Допустим, после «The little» он отвёл одну вероятность токену « girl», другую — « boy», третью — « cat». Какой из них окажется на странице? Это решает алгоритм декодирования. Можно всегда брать лидера, а можно выбирать случайно с учётом вероятностей. Второй способ называется выборкой, или сэмплированием; код, который его выполняет, — сэмплером.

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

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

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

Ещё не тронув ни одной ручки, можно заметить две вещи. Во-первых, у распределения есть голова и очень длинный хвост. Горстка токенов забирает почти всю вероятность, а тысячи остальных делят крошки: на нижнем графике линия падает на много порядков от первого места до 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^+$ вероятность сосредоточивается на токене с максимальным логитом; если максимум делят несколько токенов, предельное распределение остаётся равномерным между ними. На практике «температура 0» — отдельный режим: взять один токен с максимальной оценкой, без деления на ноль. Это жадное декодирование. При $T \to \infty$ все конечные логиты после деления стремятся к нулю, и без дополнительных фильтров распределение становится равномерным по словарю. При $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 уже намекала на подвох; теперь рассмотрим его как следует. Ниже обе колонки пишут жадно, так что результат полностью определяется началом. Красным отмечен каждый токен, который повторяет фразу из четырёх токенов, уже встречавшуюся в тексте.

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

Жадный выбор не обязательно приводит к петле. Попробуйте третье начало: Росток может закончить историю про Лили и красный мяч без заметного зацикливания. Утка и мышка дают другие примеры поведения. Выбор между жадным декодированием и выборкой зависит от задачи: для рассказа часто хочется разнообразия, а в некоторых задачах полезнее воспроизводимый ответ. Универсального запрета на жадный режим нет.

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

Грубое, но действенное лекарство появилось вместе с моделью 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) — строит длинную последовательность чисел. При одинаковом начале текста, весах, настройках и вычислениях то же зерно воспроизводит те же выборы токенов. Так устроено сравнение четырёх температур: сохраните зерно и повторите запуск, затем смените его. Для других движков и устройств одного seed может быть недостаточно: на результат влияют также порядок вычислений и численные погрешности. «Сгенерировать заново» обычно запускает новый случайный выбор, но не обещает обязательно другой ответ.

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

Колонки, которые тянут жребий, берут случайные числа из одного потока (жадной колонке случайность не нужна вовсе). На первом шаге фавориту часто принадлежит большой отрезок, так что одно и то же $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$ токенов достаточно $P + n - 1$ проходов позиций; если обработать и последний выданный токен, получится $P+n$. Наш виджет считает именно второй вариант: для начала из 14 токенов и 30 новых — 44 вместо 855. При этом внимание нового токена всё ещё обращается ко всему сохранённому контексту: выигрыш не делает всю работу постоянной по длине текста. Проверьте скорость прямо в браузере:

Кэш не бесплатен: он занимает память. Для каждого токена Росток хранит ключ и значение по 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.$$

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

А насколько сильно можно сжать? Проверьте сами. Виджет ниже загружает веса Ростка в эту страницу и заново округляет все 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, но теперь вы управляете тем, как она выбирает. Убавьте температуру — получится спокойная, предсказуемая история, прибавьте — дикая; и помните, что творится на краях шкалы. Истории Росток пишет неплохо, но для роли собеседника пока не дообучен: задайте ему вопрос, и он просто продолжит писать, будто ваш вопрос — строчка из рассказа. В следующей главе научим его разговаривать.