Timur-th/llm-course-hw1

Генератор русских анекдотов: Byte-level BPE + Transformer с нуля

Byte-level BPE токенизатор, декодерный трансформер и цикл обучения, генерация анекдотов на датасете IgorVolochay/russian_jokes.

В репозитории лежат обученный токенизатор (vocabulary.json, merges.json) и лучшая по результатам экспериментов модель (model.safetensors, config.json).

Данные

150 553 анекдота, разбитых на train (135 497) и test (15 056) в пропорции 90/10. Анекдоты длиннее 128 токенов обрезаются.

Токенизатор

Byte-level BPE, как в GPT-2: текст разбивается на слова регулярным выражением, слова переводятся в байты, после чего жадно сливаются самые частые пары. Словарь из 1024 токенов: 256 базовых байтов, 767 слияний и спецтокен [EOS]. Обучен только на train-части. В среднем на анекдот из test приходится 73.3 токена.

Архитектура

Реализованы все варианты из задания, они переключаются через конфиг:

• позиционное кодирование: ALiBi или RoPE (бонус);
• внимание: GQA или Multi-Head Latent Attention в стиле DeepSeek-V2, с низкоранговым сжатием KV в латент и decoupled RoPE (бонус);
• FFN на SwiGLU, pre-norm на RMSNorm, dropout перед residual-соединениями.

Финальная модель (mini + RoPE):

Параметр	Значение
Слои	6
Скрытая размерность	384
Головы внимания / KV-головы	6 / 3 (GQA)
Размерность SwiGLU	1024
Позиционное кодирование	RoPE
Контекст	128 токенов
Dropout	0.1
Параметры	10.5M

Обучение

AdamW (weight decay 0.01), линейный warmup на первых 10% шагов и затем линейное затухание learning rate, gradient clipping 1.0. Лосс — кросс-энтропия по следующему токену, паддинги замаскированы и в attention, и в лоссе.

[EOS] добавляется не только в конец, но и в начало каждого анекдота, где он работает как BOS. Так модель учится начинать новый анекдот после [EOS].

Финальный прогон: lr 6e-4, batch 64, 25 000 шагов (около 12 эпох), GPU T4 в Google Colab.

Отчёт о качестве

Все прогоны используют один сид и один и тот же test-сплит. Метрика — средняя кросс-энтропия на токен (меньше — лучше).

Эксперимент	Размер	Поз. код.	Внимание	lr	Batch	Шаги	Dropout	Параметры	Val loss
mini_long (финальная)	6×384	RoPE	GQA	6e-4	64	25 000	0.1	10.5M	2.410
mini_long_dropout0.2	6×384	RoPE	GQA	6e-4	64	25 000	0.2	10.5M	2.479
mini_bs64	6×384	RoPE	GQA	6e-4	64	10 000	0.1	10.5M	2.542
small	12×768	RoPE	GQA	2e-4	16	10 000	0.1	79.5M	2.659
mini_lr6e-4	6×384	RoPE	GQA	6e-4	16	10 000	0.1	10.5M	2.816
mini_lr1e-3	6×384	RoPE	GQA	1e-3	16	10 000	0.1	10.5M	2.846
mini	6×384	RoPE	GQA	3e-4	16	10 000	0.1	10.5M	2.857
mini_lr1e-4	6×384	RoPE	GQA	1e-4	16	10 000	0.1	10.5M	3.121
nano	3×96	RoPE	MLA	3e-4	16	10 000	0.1	0.53M	3.463
nano	3×96	RoPE	GQA	3e-4	16	10 000	0.1	0.50M	3.477
nano (базовая)	3×96	ALiBi	GQA	3e-4	16	10 000	0.1	0.50M	3.486
nano	3×96	ALiBi	MLA	3e-4	16	10 000	0.1	0.52M	3.488

Выводы:

• Размер модели даёт самый большой прирост: при одинаковом бюджете обучения nano → mini → small снижает лосс с 3.49 до 2.86 и 2.66.
• Объём обучения важнее подбора learning rate. Увеличение батча с 16 до 64 снизило лосс с 2.82 до 2.54, а удлинение обучения до 25k шагов — до 2.41. Признаков переобучения нет: dropout 0.2 сделал результат хуже, значит модель скорее недообучена.
• Learning rate: для mini оптимум около 6e-4; 1e-4 заметно хуже.
• RoPE против ALiBi и MLA против GQA на nano различаются в пределах 0.025. Это прогоны с одним сидом, так что разница сопоставима с шумом. RoPE выбран для дальнейших экспериментов как немного лучший.

Примеры генерации

Финальная модель, top_k=10. Промпт выделен жирным.

Заходит в бар и говорит бармену: - Слушай, дружище, а зачем мне виски на дороге? - А как ты это догадался? - Ну да, да вот и я виски вискарь!

Встречаются два программиста:- Ты слышал, что в России появились новые факсы?- Да, ведь вчера вечером мы с женой встретились.- Ну и что?- Так он же на ней женился!

Встречаются два программиста.- Ну, как твоя новая секретарша? Она такая же какая!

Штирлиц пришел домой и, увидев свой первый срок, спрашивает:- Что ты там делаешь?- Гду?- Гиви, а ты на осяльник!

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

Как использовать

Классы ByteLevelBPETokenizer и TransformerForCausalLM определены в ноутбуке с заданием, поэтому их код нужно выполнить до загрузки.

import torch

repo = "Timur-th/llm-course-hw1"
tokenizer = ByteLevelBPETokenizer.from_pretrained(repo)
model = TransformerForCausalLM.from_pretrained(repo).eval()

prompt = "Заходит в бар"
ids = [tokenizer.eos_token_id] + tokenizer.encode(prompt, add_eos_token=False)
output = model.generate(
    torch.tensor(ids)[None, :],
    max_new_tokens=128,
    eos_token_id=tokenizer.eos_token_id,
    do_sample=True,
    top_k=10,
)
print(tokenizer.decode(output[0, 1:].tolist()))