maf1oz1/llm-course-hw1-mla

Russian Jokes LM — mini с Multi-Head Latent Attention (бонус)

Вариант модели из домашнего задания №1 курса по LLM (HSE), в котором Grouped-Query Attention заменён на Multi-Head Latent Attention (MLA) — механизм из DeepSeek-V2. Остальное — токенизатор, ALiBi, SwiGLU, данные и гиперпараметры — совпадает с основной моделью: maf1oz1/llm-course-hw1.

Зачем MLA

При генерации модель хранит ключи и значения всех предыдущих токенов (KV-кэш), и на длинных текстах он занимает много памяти. MLA уменьшает его: вектор токена сначала сжимается в маленький латентный вектор c_kv, а ключи и значения для всех голов восстанавливаются из него. Кэшировать нужно только латент.

Как реализовано

x ──kv_down──► c_kv (128) ──RMSNorm──┬──k_up──► K для всех 6 голов
                                     └──v_up──► V для всех 6 голов

• kv_down: 384 → 128, k_up и v_up: 128 → 384; GQA-проекция kv_proj удалена.
• Каждая голова получает свои K и V из общего латента — размножать KV-головы, как в GQA, не нужно.
• Параметров в attention столько же, сколько у GQA-версии (384·128 + 2·128·384 = 384·384), так что сравнение честное.
• Упрощённый вариант: без decoupled RoPE и без сжатия запросов — позиции кодируются ALiBi, который с латентом не конфликтует.
• Реализация — классы-наследники MLACausalSelfAttention, MLABlock, MLATransformerForCausalLM; основной код модели не меняется.

KV-кэш

Схема	Чисел на токен в каждом слое	Во сколько раз меньше MHA
MHA (6 KV-голов)	768	1×
GQA (3 KV-головы, основная модель)	384	2×
MLA (латент 128) — эта модель	128	6×

Архитектура и обучение

Параметр	Значение
Параметров	10.54M
Слоёв / скрытая размерность	6 / 384
Голов внимания	6, K/V из латента размера 128
Позиционное кодирование	ALiBi
Feed-forward	SwiGLU, 1024
Нормализация	RMSNorm, pre-norm
Контекст / словарь	128 токенов / 1024 (byte-level BPE)
Обучение	3 эпохи (25 404 шага), batch 16, AdamW lr 3e-4, warmup 10% + linear decay

Результаты

Модель	KV-кэш на токен в слое	Val loss	Perplexity
ALiBi + GQA (основная)	384	2.651	14.2
RoPE + GQA	384	2.612	13.6
ALiBi + MLA — эта модель	128	2.682	14.6

MLA уменьшает KV-кэш в 3 раза по сравнению с GQA ценой небольшого ухудшения лосса (+0.03). На коротких анекдотах выигрыш по памяти незаметен — в нашем generate нет KV-кэша, — но на длинном контексте и больших батчах он измерялся бы сотнями мегабайт даже для модели такого размера.

Как запустить

Классы ByteLevelBPETokenizer, TransformerForCausalLM и MLATransformerForCausalLM определены в тетрадке домашнего задания — их нужно выполнить перед загрузкой.

import torch

tokenizer = ByteLevelBPETokenizer.from_pretrained("maf1oz1/llm-course-hw1-mla")
model = MLATransformerForCausalLM.from_pretrained("maf1oz1/llm-course-hw1-mla").eval()

ids = torch.tensor([tokenizer.encode("Штирлиц пришел домой", add_eos_token=False)])
out = model.generate(ids, max_new_tokens=150, eos_token_id=tokenizer.eos_token_id,
                     do_sample=True, top_k=20, temperature=0.8)
print(tokenizer.decode(out[0].tolist()))

Ограничения

Те же, что у основной модели: около 10M параметров и ~10M токенов данных — форма анекдота воспроизводится, смысл быстро теряется; маленький словарь иногда порождает несуществующие слова; модель может повторять грубые шутки из датасета.