Авторегрессионные трансформеры: ImageGPT, VQGAN, DALL·E и Parti

Предыдущая заметка строила модели изображений из цепного правила и маскированных свёрток. Остались два ограничения: свёртка видит фиксированную окрестность, а изображение $256 \times 256$ пикселей — это последовательность почти из 200 000 значений. Трансформеры снимают первое ограничение, потому что причинное self-attention на каждом слое читает всё прошлое, и усугубляют второе, потому что внимание сравнивает все пары позиций. Модели этой заметки отвечают тем, что меняют сам токен: сначала это цвет палитры у маленького изображения (ImageGPT), затем код обученного дискретного автоэнкодера (VQGAN, DALL·E, Parti). Второй ответ превращает генерацию изображений в задачу, которую языковые модели уже решают, а генерацию по тексту — в продолжение подписи.

Обозначения. Последовательность токенов $s = (s_1, \dots, s_N)$ из словаря размера $K$ моделируется как $p(s) = \prod_i p(s_i \mid s_{<i})$. Сетка кодов из $h \times w$ токенов читается построчно, так что $N = hw$. Условие $c$, например подпись, — вторая последовательность токенов.

От свёрток к вниманию#

Декодер трансформера на каждом слое вычисляет внимание от каждой позиции к предшествующим. Причинная маска ставит оценку $-\infty$ всем более поздним позициям перед softmax, так что выход во входной позиции $i$ зависит от входов в позициях $1, \dots, i$, включая её саму. Оставшийся доступ к ответу убирает сдвиг между входом и целью: вход — $(\text{BOS}, s_1, \dots, s_{N-1})$ с начальным токеном BOS, цель — $(s_1, \dots, s_N)$, и выход во входной позиции $i$, который видел $(\text{BOS}, s_1, \dots, s_{i-1})$, предсказывает $s_i$. Для токенов $(7, 3, 9)$ вход — $(\text{BOS}, 7, 3)$, цели — $(7, 3, 9)$, а функция потерь равна $-\log p(7 \mid \text{BOS}) - \log p(3 \mid \text{BOS}, 7) - \log p(9 \mid \text{BOS}, 7, 3)$. Маска не пускает будущее, сдвиг не пускает саму цель; маска A в PixelCNN делала и то и другое сразу, скрывая текущее значение, так что в обеих моделях выход $i$ зависит только от $s_{<i}$. Заметка о GPT из курса NLP выводит следствия, и они переносятся без изменений. Обучение — один параллельный проход с teacher forcing и кросс-энтропией в каждой позиции. Сэмплирование последовательно, по токену за проход. Слепого пятна нет: каждый слой дотягивается до любой предшествующей позиции.

Каждый слой сравнивает $N(N+1)/2$ пар позиций, так что стоимость растёт как квадрат длины последовательности, и это решает, какими могут быть токены изображения.

Рисунок

Сколько токенов в одном изображении

У сэмплирования есть второй квадратичный расход. Если для каждого токена заново прогонять сеть по всему префиксу, то всего сравнивается $\sum_{t=1}^{N} t(t+1)/2 = N(N+1)(N+2)/6$ пар на слой: 179 481 600 при $N = 1024$. Ключи и значения прежних позиций не меняются, когда добавляется токен, поэтому KV-кэш хранит их и вычисляет внимание только для нового запроса: $\sum_t t = N(N+1)/2 = 524\,800$ пар, в $(N+2)/3 = 342$ раза меньше. Цена — память: два вектора на слой и позицию. У Qwen2-0.5B, небольшой открытой языковой модели с 24 слоями и двумя головами ключей и значений размерности 64, при хранении в 16-битных числах это $2 \cdot 24 \cdot 2 \cdot 64 \cdot 2 = 12\,288$ байт на токен, 12 МиБ на 1024 токена и 384 МиБ на весь контекст из 32 768 токенов. Без grouped-query attention, с 14 головами ключей и значений, было бы в семь раз больше. Раздел с кодом проверяет, что кэш не меняет ни одного выхода.

Настройки, превращающие предсказанное распределение в сэмпл, — температура, top-$k$ и top-$p$ — те же, что для текста, и описаны в заметке о GPT. На изображениях они действуют как трюк с усечением у GAN: меняют разнообразие на типичность сэмплов, как показывают результаты VQGAN ниже.

Пиксели как токены: ImageGPT#

ImageGPT (Chen et al., 2020) применяет GPT-2 к изображениям с минимумом изменений. Последовательность $224 \times 224 \times 3$ сделала бы логиты внимания в десятки тысяч раз больше, чем в языковых моделях, поэтому изображение сначала уменьшают до $32 \times 32$ или $48 \times 48$, а три канала каждого пикселя заменяют ближайшим из 512 цветов, найденных k-means: получается последовательность из $32^2 = 1024$ или $48^2 = 2304$ токенов над словарём из 512. Палитра — простейший токенизатор, VQ-VAE с тождественными энкодером и декодером, как замечает статья. iGPT-L, модель, к которой относятся результаты на CIFAR-10 ниже, — это GPT-2 с 48 слоями и шириной 1536, 1.4 миллиарда параметров; в работе обучена и более крупная iGPT-XL с 6.8 миллиарда.

Статью интересует прежде всего представление, а не сэмплирование. Модель, обученная только предсказывать следующий пиксель, выучивает признаки, по которым классы CIFAR-10 линейно разделяются с точностью 96.3%, а после дообучения — 99.0%. Статья показывает и достраивание половинок изображений, но в этих малых размерах: пиксельные токены не масштабируются до разрешений моделей ниже.

Обученные токены: двухэтапный рецепт#

Выход уже был в VQ-VAE: сначала обучить дискретный автоэнкодер, затем авторегрессионную модель над его кодами. При коэффициенте понижения $f$ сетка содержит $(H/f) \times (W/f)$ позиций, и трансформер не видит ни одного пикселя.

Рисунок

Два этапа: токенизатор изображений и трансформер над токенами

Два этапа делят работу так же, как в латентной диффузии, появившейся позже в той же группе, что и VQGAN: токенизатор отвечает за текстуры и мелкие детали, априорная модель — за композицию. Поэтому токенизатор должен сильно сжимать и всё же восстанавливать резкие изображения — то, чего не делает обычный VQ-VAE с квадратичной функцией потерь.

VQGAN#

VQGAN (Esser, Rombach and Ommer, 2021) сохраняет векторный квантизатор, слагаемые кодовой книги и commitment из VQ-VAE и меняет функцию потерь реконструкции. Квадратичная ошибка заменяется перцептивной, как в LPIPS, и добавляется дискриминатор по патчам, который оценивает небольшие фрагменты реконструкции как настоящие или нет. Состязательное слагаемое получает адаптивный вес

$$\lambda = \frac{\lVert \nabla_{G_L} \mathcal L_{\text{rec}} \rVert}{\lVert \nabla_{G_L} \mathcal L_{\text{GAN}} \rVert + \delta},$$

отношение норм градиентов двух функций потерь по последнему слою $G_L$ декодера, чтобы ни одно слагаемое не подавляло другое. При $f = 16$ изображение $256 \times 256$ становится сеткой $16 \times 16$, и статья обычно берёт кодовую книгу из $K = 1024$ кодов: 256 токенов по 10 бит, то есть 2560 бит против 1 572 864 бит исходных пикселей. Второй этап — трансформер GPT-2 medium с 307 миллионами параметров, и последовательность из 256 токенов была самой длинной, что помещалась на GPU с 12 ГБ. Условия добавляются в начало как токены: метка класса или коды карты сегментации либо карты глубины, полученные вторым VQGAN. Изображения мегапиксельного размера генерируются со скользящим окном внимания: каждый токен предсказывается по окрестности $16 \times 16$ кодов.

Токенизатор ограничивает то, чего может достичь второй этап, потому что каждый сэмпл — выход декодера. На ImageNet реконструкции через сетку $16 \times 16$ имеют FID 7.94 относительно валидационных изображений при 1024 кодах и 4.98 при 16 384; статья считает это число оценкой лучшего FID, достижимого вторым этапом. Это эмпирический ориентир, а не теорема: он измеряет одно распределение выходов декодера — реконструкции, — а априорная модель может распределить массу по выходам декодера иначе.

Настройки сэмплирования на изображениях#

Модель VQGAN с обусловливанием на класс ImageNet (кодовая книга из 16 384) показывает, как много значат настройки сэмплирования. Её FID при простом top-$k$ сэмплировании с $k = 973$ равен 29.20; при $k = 250$ — 15.98, а top-$p = 0.88$ поверх $k = 973$ даёт 15.78. Если оставить только сэмплы, которые классификатор ResNet-101 уверенно относит к их классу, лучшие 5%, FID падает до 5.88 — ниже 7.53 у BigGAN. Каждый из этих шагов убирает маловероятные сэмплы и меняет полноту на точность: FID улучшается потому, что оставшиеся сэмплы типичнее для своего класса; данные модель охватывает не лучше.

Текст в изображение: DALL·E#

DALL·E (Ramesh et al., 2021) применяет рецепт в большом масштабе. Токенизатор — дискретный VAE, обученный с релаксацией Gumbel-softmax, который переводит изображение $256 \times 256$ в сетку $32 \times 32$ над 8192 кодами: последовательность в 192 раза короче пиксельной. Подпись кодируется не более чем 256 BPE-токенами из словаря на 16 384, и текст вместе с 1024 токенами изображения склеивается в одну последовательность длиной до 1280 токенов. Весь поток моделирует один трансформер-декодер с 12 миллиардами параметров и 64 слоями разреженного внимания, обученный на 250 миллионах пар «изображение — подпись». Генерация по тексту — это продолжение подписи: подать токены текста и вытянуть 1024 токена изображения.

Сэмплы вытягиваются пачками и ранжируются заново: контрастная модель изображений и текстов (CLIP) оценивает каждого из 512 кандидатов по соответствию подписи, и остаются лучшие. В оценке людьми на подписях MS-COCO против DF-GAN, GAN, обученной на этом датасете, большинство из пяти оценщиков сочло сэмпл DALL·E более реалистичным в 90.0% случаев и лучше соответствующим подписи в 93.3%, хотя DALL·E никогда не обучалась на подписях MS-COCO.

Parti: энкодер–декодер в большом масштабе#

Parti (Yu et al., 2022) рассматривает генерацию по тексту как перевод. Энкодер читает подпись — не более 128 токенов из словаря SentencePiece на 16 000, — а декодер генерирует 1024 токена изображения ViT-VQGAN, то есть VQGAN, у которого энкодер и декодер — визуальные трансформеры, с 8192 кодами. Сеть сверхразрешения примерно с 15 миллионами параметров поднимает декодированные изображения $256 \times 256$ до $1024 \times 1024$. Статья обучает четыре размера — 350 миллионов, 750 миллионов, 3 миллиарда и 20 миллиардов параметров — и находит, что качество и соответствие тексту продолжают расти с размером. Модель на 20 миллиардов достигает FID 7.23 на MS-COCO без обучения на нём и 3.22 после дообучения.

Сэмплы улучшают два инструмента из других семейств. Classifier-free guidance работает на логитах так же, как на score: модель обучают, иногда выбрасывая подпись, а при сэмплировании безусловное и условное предсказания комбинируются с весом больше единицы. Контрастное переранжирование оставляет лучший из 16 сэмплов на запрос вместо 512 у DALL·E. Статья о Parti также ввела PartiPrompts — набор из более чем 1600 запросов, описанный в заметке об оценке людьми.

Модели рядом#

Модель	Токены изображения	Словарь	Трансформер	Параметры
ImageGPT (iGPT-L)	$32 \times 32$ пикселя палитры уменьшенного изображения	512 цветов	декодер	1.4 млрд
VQGAN	$16 \times 16$ кодов для $256 \times 256$	1024 (16 384 на ImageNet)	декодер	307 млн
DALL·E	$32 \times 32$ кода после не более 256 токенов текста	8192	декодер, разреженное внимание	12 млрд
Parti	$32 \times 32$ кода, текст в отдельном энкодере	8192	энкодер–декодер	20 млрд

Авторегрессионные трансформеры и диффузия#

У двух семейств больше общего, чем кажется. Оба работают в латентном пространстве отдельно обученного автоэнкодера, оба обусловливаются на текст через внимание, и оба управляются classifier-free guidance. Различаются они тем, как строится сэмпл. Авторегрессионная модель фиксирует по одному токену — 1024 прохода сети для сетки $32 \times 32$ — и не может исправить токен, когда он уже вытянут. Диффузионная модель уточняет весь латент сразу за 20–50 шагов. Взамен авторегрессионная модель даёт точное правдоподобие токенов и разделяет с языковыми моделями архитектуру, код обучения и поведение при масштабировании.

Итоги#

• Причинный трансформер — это PixelCNN без слепого пятна: каждый слой видит всё прошлое ценой стоимости, квадратичной по длине последовательности.
• KV-кэш превращает кубическую стоимость наивного сэмплирования в квадратичную ценой памяти: два вектора на слой и токен.
• Пиксели дают слишком длинные последовательности. ImageGPT сокращает их палитрой и маленьким изображением; двухэтапный рецепт заменяет их кодами дискретного автоэнкодера.
• VQGAN делает токенизатор достаточно хорошим при сильном сжатии благодаря перцептивной и состязательной функциям потерь; его FID реконструкций — практический ориентир для того, чего может достичь априорная модель, а не доказанная граница.
• DALL·E и Parti обусловливаются на текст, помещая его в последовательность или в энкодер, и улучшают сэмплы переранжированием и classifier-free guidance; top-$k$, top-$p$ и отбраковка меняют разнообразие на типичность сэмплов.

Ноутбук#

Ноутбука об авторегрессионных трансформерах для изображений в курсе нет. Его ноутбук о LLM, Transformers-LangChain.ipynb в репозитории HSE-LAMBDA Deep Generative Models, запускает предобученные трансформеры на тексте и звуке с Hugging Face Transformers, а затем авторегрессионную чат-модель с LangChain и LangGraph. Номера ячеек считают все ячейки, включая текстовые, начиная с 1. Выходы ячеек сохранены, и числа ниже взяты из них. Идеи, стоящие за ячейками, разобраны в заметках курса NLP, на которые ссылаются пункты списка.

• Ячейки 1–9: конвейеры. Ячейка 4 сводит в таблицу одиннадцать задач pipeline(): классификацию текста, генерацию текста и суммаризацию для текста; классификацию изображений, сегментацию изображений и обнаружение объектов для изображений; классификацию звука и распознавание речи для звука; ответы на вопросы по изображению, ответы на вопросы по документу и генерацию подписей к изображениям для смешанных входов. Ячейка 6 создаёт конвейер анализа тональности, не называя модели. Конвейер берёт модель по умолчанию, distilbert/distilbert-base-uncased-finetuned-sst-2-english в ревизии 714eb0f, и предупреждает, что конвейер без зафиксированных модели и ревизии не рекомендуется для промышленного использования, поскольку модель по умолчанию может смениться. Ячейки 7–9 классифицируют одно предложение, затем список из двух, для которого возвращается по одному результату на каждый вход. «We are very happy to show you the 🤗 Transformers library.» получает POSITIVE с оценкой 0.9998; «We hope you don’t hate it.» — NEGATIVE с 0.5309: почти подбрасывание монеты на осторожной фразе, построенной на отрицании негативного слова.
• Ячейки 10–17: распознавание речи. facebook/wav2vec2-base-960h — это wav2vec 2.0, дообученная с выходным слоем CTC (connectionist temporal classification) на 960 часах аудиокниг LibriSpeech с частотой дискретизации 16 кГц; она работает на CPU. Предупреждение о том, что masked_spec_embed инициализирован заново, безвредно: этот вектор заменяет замаскированные кадры только при обучении. Ячейка 13 загружает обучающую часть en-US набора MInDS-14 — устных запросов к интернет-банку, записанных с частотой 8 кГц, — и приводит его звуковой столбец к 16 кГц экстрактора признаков модели, так что каждый клип передискретизируется при чтении. Ячейка 15 распознаёт первые четыре клипа — все они о совместном счёте. Расшифровки набраны заглавными буквами без знаков препинания, потому что в словаре CTC есть только 26 заглавных букв, апостроф и разделитель слов. Телефонная речь далека от начитанных аудиокниг, и ошибки это показывают: «how do I start a joint account» превращается в «HOW DO I FURN A JOINA COUT». Ячейка 17 передаёт генератор Python вместо списка, так что конвейер читает и распознаёт клипы по одному, не держа весь набор данных в памяти; расшифровки те же.
• Ячейки 18–21: модель с Hub. nlptown/bert-base-multilingual-uncased-sentiment, многоязычный BERT, дообученный предсказывать от 1 до 5 звёзд, ставит французскому предложению 5 звёзд с вероятностью 0.73.
• Ячейки 22–31: что конвейер делает внутри. Токенизатор этой модели превращает «We are very happy to show you the 🤗 Transformers library.» в 14 идентификаторов, от 101 ([CLS]) до 102 ([SEP]). Для эмодзи в словаре из 105 879 записей нет куска, и оно становится 100 — идентификатором [UNK], а token_type_ids для одного сегмента все равны 0 (эти идентификаторы объясняет заметка о токенизации). Ячейка 28 токенизирует оба английских предложения одним батчем с padding=True, truncation=True, max_length=512 и return_tensors="pt"; ячейка 31 прогоняет на нём AutoModelForSequenceClassification и применяет softmax к логитам. По звёздам от 1 до 5 первое предложение получает [0.0021, 0.0018, 0.0115, 0.2121, 0.7725]; второе — [0.2084, 0.1826, 0.1969, 0.1755, 0.2365], почти равномерное распределение: та же неуверенность, что показала на нём английская модель тональности. Прямой проход идёт без torch.no_grad(), так что выход несёт граф градиентов, который при инференсе не нужен.
• Ячейки 32–35: дообучение с Trainer, только в текстовых ячейках. DistilBERT (distilbert-base-uncased) на предложениях из рецензий на фильмы Rotten Tomatoes: шаг обучения $2 \cdot 10^{-5}$, размер батча 8, 2 эпохи, оценка каждые 500 шагов (значение по умолчанию, когда выбрана стратегия "steps"), набор данных токенизируется один раз через map(batched=True) без паддинга, а DataCollatorWithPadding дополняет каждый батч до его собственной самой длинной последовательности (зачем — объясняет заметка о дообучении). Ничего не запускается. Перед запуском нужно поменять две детали: текущие версии Transformers принимают только eval_strategy и отвергают evaluation_strategy, а оценка идёт на тестовой части, хотя в наборе данных есть валидационная — именно за ней следят во время обучения.
• Ячейки 36–41: LangChain. Ячейка 38 перечисляет, для чего используют библиотеку: чат-боты с памятью, агенты с инструментами и генерация с извлечением. Ячейка 40 называет два места, где выбирать модель: рейтинг LMArena и модели генерации текста на Hugging Face, отсортированные по популярности (trending). Ячейка 41 загружает Qwen/Qwen2-0.5B-Instruct, модель из примера с KV-кэшем выше, как конвейер генерации текста не более чем на 100 новых токенов и оборачивает его как чат-модель LangChain. Настройки генерации модели сэмплируют с температурой 0.7, top-$k$ 20, top-$p$ 0.8 и штрафом за повторы 1.1, так что ответы меняются от запуска к запуску.
• Ячейки 42–46: своей памяти у модели нет. На вопрос «What’s my name?» без истории она отвечает общими словами об именах; с предыдущими репликами на входе отвечает «Your name is Bob.» Разговор — это префикс, который подаётся заново на каждом ходу, как объясняет заметка об инструментах и агентах.
• Ячейки 47–51: разговор с контрольными точками. StateGraph из LangGraph над MessagesState содержит один узел, который вызывает модель, и ребро к нему от START и компилируется с MemorySaver в роли хранилища контрольных точек (checkpointer). thread_id в конфигурации выбирает разговор; второй invoke передаёт только новый вопрос, и ответ на него строится по сохранённой истории.
• Ячейки 52–64: шаблоны промптов и расширенное состояние. ChatPromptTemplate.from_messages ставит системное сообщение перед MessagesPlaceholder для истории. С инструкцией говорить как пират (ячейки 55–58) модель на 0.5 миллиарда параметров отвечает обычным английским. Ячейки 59–64 расширяют состояние до TypedDict с полем messages: Annotated[Sequence[BaseMessage], add_messages], чей редьюсер (reducer) add_messages дописывает новые сообщения, и обычным полем language. Второй invoke не передаёт язык, сохранённое в контрольной точке значение «Spanish» по-прежнему действует, и ответ — «Tu nombre es Bob.»
• Ячейки 65–70: обрезка. trim_messages с бюджетом 45 токенов, strategy="last" (оставлять самые новые сообщения), include_system=True (оставлять системное сообщение), allow_partial=False (никогда не разрезать сообщение), start_on="human" (начинать оставленную историю с реплики пользователя) и token_counter=chat_model сокращает пример из 11 сообщений до 7. Затем граф обрезает сохранённую историю перед каждым вызовом, а ячейка 70 читает контрольную точку обратно через get_state(config). Обрезка — способ удержать чат в фиксированном контексте, 32 768 токенов для этой модели, и она ограничивает KV-кэш.

Две проблемы:

• Сохранённые ответы содержат промпт. Каждое сообщение модели начинается со всего входа в её шаблоне чата (<|im_start|>system …) и лишь затем даёт ответ, потому что конвейер генерации текста по умолчанию возвращает промпт вместе с продолжением. Когда LangGraph сохраняет такие сообщения (ячейки 47–51) и отправляет их обратно на следующем ходу, новый промпт содержит старый, вложенный в реплику ассистента. Короткие диалоги всё равно получают правильные ответы, но контекст растёт с каждым ходом. Параметр return_full_text=False у конвейера оставляет только сгенерированный текст.
• Обрезка считает не те токены. С token_counter=chat_model LangChain считает токены своим запасным токенизатором GPT-2, а не токенизатором Qwen. Из-за повторённого промпта единственное сообщение ассистента, сохранённое ячейкой 69, уже занимает не меньше 66 токенов GPT-2 — больше бюджета в 45 токенов, — так что на втором ходу граф с обрезкой оставил бы только самый новый вопрос. Кроме того, ячейка 68 обрезает историю до того, как шаблон добавляет системное сообщение, поэтому include_system=True в графе ничего не делает.

Ноутбук с заданием курса по LLM — ассистентом с извлечением по аннотациям медицинских статей — описан в разделах «Ноутбук» заметок о семантическом поиске и о RAG.

Код#

Небольшой трансформер-декодер с причинным вниманием, KV-кэшем и фильтрами сэмплирования на PyTorch. Проверка в конце сравнивает логиты с кэшем и без него и подтверждает, что изменение более позднего токена не меняет более ранних предсказаний:

import torch
import torch.nn as nn
import torch.nn.functional as F

class CausalSelfAttention(nn.Module):
    def __init__(self, d, heads):
        super().__init__()
        self.h, self.qkv, self.out = heads, nn.Linear(d, 3 * d), nn.Linear(d, d)

    def forward(self, x, cache=None):                 # x: [B, T, d], только новые позиции
        B, T, d = x.shape
        q, k, v = self.qkv(x).view(B, T, 3, self.h, d // self.h).permute(2, 0, 3, 1, 4)
        if cache is not None:                         # ключи и значения прежних позиций
            if 'k' in cache:
                k, v = torch.cat([cache['k'], k], 2), torch.cat([cache['v'], v], 2)
            cache['k'], cache['v'] = k, v
        S = k.shape[2]                                # позиции, видимые на данный момент
        mask = torch.ones(T, S, dtype=torch.bool).tril(S - T)   # новый запрос i видит ключи до своей позиции
        y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        return self.out(y.transpose(1, 2).reshape(B, T, d))

class TinyGPT(nn.Module):
    def __init__(self, vocab=16, d=32, heads=4, layers=2, max_len=64):
        super().__init__()
        self.tok, self.pos = nn.Embedding(vocab, d), nn.Embedding(max_len, d)
        self.blocks = nn.ModuleList(CausalSelfAttention(d, heads) for _ in range(layers))
        self.mlps = nn.ModuleList(nn.Sequential(nn.LayerNorm(d), nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d)) for _ in range(layers))
        self.norms = nn.ModuleList(nn.LayerNorm(d) for _ in range(layers))
        self.head = nn.Linear(d, vocab)

    def forward(self, idx, caches=None, start=0):     # idx: [B, T], токены в позициях start .. start+T-1
        x = self.tok(idx) + self.pos(torch.arange(start, start + idx.shape[1]))
        for l, (attn, mlp, norm) in enumerate(zip(self.blocks, self.mlps, self.norms)):
            x = x + attn(norm(x), None if caches is None else caches[l])
            x = x + mlp(x)
        return self.head(x)                           # логиты: [B, T, vocab]

def filter_logits(logits, temperature=1.0, top_k=None, top_p=None):
    logits = logits / temperature
    if top_k is not None:                             # оставить k наибольших логитов
        kth = logits.topk(top_k, dim=-1).values[..., -1:]
        logits = logits.masked_fill(logits < kth, float('-inf'))
    if top_p is not None:                             # оставить наименьший набор с вероятностью >= top_p
        sorted_logits, order = logits.sort(dim=-1, descending=True)
        cum = sorted_logits.softmax(-1).cumsum(-1)
        drop = cum - sorted_logits.softmax(-1) >= top_p   # масса до этого токена уже достигает top_p
        logits = logits.masked_fill(drop.scatter(-1, order, drop), float('-inf'))
    return logits

@torch.no_grad()
def generate(model, prompt, n_new, **kw):             # prompt: [B, T0]
    caches = [{} for _ in model.blocks]
    logits = model(prompt, caches)                    # один параллельный проход по запросу заполняет кэш
    out = prompt
    for _ in range(n_new):
        probs = filter_logits(logits[:, -1], **kw).softmax(-1)
        nxt = torch.multinomial(probs, 1)
        out = torch.cat([out, nxt], 1)
        logits = model(nxt, caches, start=out.shape[1] - 1)   # обрабатывается только новый токен
    return out

torch.manual_seed(0)
model = TinyGPT().eval()
seq = generate(model, torch.tensor([[1, 2, 3]]), 20, temperature=0.7, top_k=5)
with torch.no_grad():
    full = model(seq[:, :-1])                         # пересчитать все позиции заново
    caches = [{} for _ in model.blocks]
    step = torch.cat([model(seq[:, i:i + 1], caches, start=i) for i in range(seq.shape[1] - 1)], 1)
    a = seq[:, :10].clone(); b = a.clone(); b[0, 7] = (b[0, 7] + 1) % 16
    print('max |cached - full| =', (full - step).abs().max().item())
    print('earlier logits unchanged:', torch.allclose(model(a)[:, :7], model(b)[:, :7]))

z = torch.tensor([[2.0, 1.5, 1.0, 0.0, -1.0]])
for kw in [{}, {'temperature': 0.7}, {'top_k': 3}, {'top_p': 0.7}]:
    print(kw, filter_logits(z, **kw).softmax(-1).numpy().round(3))

Запуск печатает наибольшую разницу $4.8 \cdot 10^{-7}$ — ошибку округления 32-битной арифметики — и True. Для логитов $(2, 1.5, 1, 0, -1)$ вероятности равны 0.463, 0.281, 0.170, 0.063 и 0.023; температура 0.7 заостряет их до 0.555, 0.272, 0.133, 0.032 и 0.008; top-$k$ с $k = 3$ оставляет три токена с 0.506, 0.307 и 0.186; а top-$p$ с $p = 0.7$ — два с 0.622 и 0.378, потому что у первого токена меньше 0.7 массы, а у первых двух 0.744.

Дальше#

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

Литература#

• M. Chen, A. Radford, R. Child, J. Wu, H. Jun, D. Luan, I. Sutskever. Generative Pretraining from Pixels. ICML, 2020.
• P. Esser, R. Rombach, B. Ommer. Taming Transformers for High-Resolution Image Synthesis. CVPR, 2021.
• A. Ramesh, M. Pavlov, G. Goh, S. Gray, C. Voss, A. Radford, M. Chen, I. Sutskever. Zero-Shot Text-to-Image Generation. ICML, 2021.
• J. Yu et al. Scaling Autoregressive Models for Content-Rich Text-to-Image Generation. TMLR, 2022.
• J. Yu, X. Li, J. Y. Koh, H. Zhang, R. Pang, J. Qin, A. Ku, Y. Xu, J. Baldridge, Y. Wu. Vector-quantized Image Modeling with Improved VQGAN. ICLR, 2022. ViT-VQGAN.
• A. van den Oord, O. Vinyals, K. Kavukcuoglu. Neural Discrete Representation Learning. NeurIPS, 2017. VQ-VAE.