Современные архитектуры LLM
GQA/MLA, SwiGLU, gated attention, embedding sharing, RMSNorm — как устроены frontier модели.
Modern LLM Architectures — от MHA до MLA и далее
Архитектура frontier LLM — это не "берём трансформер и добавляем слоёв". Каждый компонент — от механизма внимания до функции активации — прошёл годы итераций и боли. Тут разбираем конкретные решения из Kimi K2 (1.06T), gpt-oss-120b (116.83B), SmolLM3 (3B) и Arcee Trinity (400B) — с числами, формулами и результатами ablation-ов. Никакой воды, только мясо.
Attention Variants: MHA → MQA → GQA → MLA
Классический Multi-Head Attention (MHA) — у каждой головы свои Q, K, V проекции. Звучит логично, но KV-cache при инференсе разрастается до неприличия — это главный боттлнек для длинных контекстов. MQA (Multi-Query) качнул маятник в другую сторону: один K,V на всех. Дёшево, но головы теряют специализацию. GQA — золотая середина: K,V общие внутри небольших групп (2/4/8 голов).
Базовый scaled dot-product attention — основа всех вариантов
HuggingFace прогнали ablation на 1B модели: GQA с 2/4/8 группами стабильно бьёт MHA по HellaSwag, MMLU и ARC. MHA лучше MQA и GQA с 16 группами. А дальше появился MLA (Multi-Latent Attention) от DeepSeek — это вообще другой уровень. Вместо полного KV-cache хранится сжатая латентная переменная. При инференсе она декомпрессируется в K,V — 4-8× сжатие KV-cache при качестве не хуже GQA. Магия линейной алгебры.
📊 Frontier Model Comparison
import torch
import torch.nn as nn
import torch.nn.functional as F
class GQA(nn.Module):
"""Grouped Query Attention — KV heads shared across groups."""
def __init__(self, d_model=2048, n_heads=32, n_kv_heads=8):
super().__init__()
self.n_heads = n_heads
self.n_kv_heads = n_kv_heads
self.head_dim = d_model // n_heads
self.n_rep = n_heads // n_kv_heads # сколько Q голов делят один KV
self.q_proj = nn.Linear(d_model, n_heads * self.head_dim, bias=False)
self.k_proj = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
self.v_proj = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
self.o_proj = nn.Linear(d_model, d_model, bias=False)
def forward(self, x):
B, S, _ = x.shape
q = self.q_proj(x).view(B, S, self.n_heads, self.head_dim)
k = self.k_proj(x).view(B, S, self.n_kv_heads, self.head_dim)
v = self.v_proj(x).view(B, S, self.n_kv_heads, self.head_dim)
# Repeat KV heads to match Q heads (GQA → MHA broadcast)
k = k.repeat_interleave(self.n_rep, dim=2) # (B, S, n_heads, d)
v = v.repeat_interleave(self.n_rep, dim=2)
# Standard attention (+ apply RoPE here in practice)
attn = F.scaled_dot_product_attention(
q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2),
is_causal=True
)
return self.o_proj(attn.transpose(1, 2).reshape(B, S, -1))
# KV cache: GQA(8) = 8 KV heads vs MHA(32) = 32 KV heads → 4x меньше!Gated Attention
Gated attention — это фильтр на выходе attention. Сигмоида вычисляет гейт-вектор из входа и покомпонентно прижимает выход каждой головы. Зачем? Без него attention sinks (токены, которые воруют всё внимание) дестабилизируют тренировку. Гейт мягко давит эти аномалии. При масштабе 100B+ это разница между «тренировка сошлась» и «loss улетел в NaN».
Гейт-вектор: σ(W^G · x_t) — обучаемое гейтирование для каждой позиции
Выход attention покомпонентно умножается на гейт — подавляет нестабильные активации
SwiGLU и функции активации
SwiGLU убила ReLU/GeLU в FFN блоках и заняла их место. Идея: вход делится надвое, одна часть идёт через swish, другая — через linear gate, результаты перемножаются. Почему это работает лучше? Честный ответ — никто толком не знает, но ablation-ы убедительны. gpt-oss-120b использует gated SwiGLU. Исключения: Gemma 2 (GeGLU), NVIDIA (relu²) — бунтари.
Embedding Sharing (Tied vs Untied)
Input и output embeddings — это две матрицы размером vocab_size × d_hidden каждая. Для мелких моделей это до 20% всех параметров! В Llama 3.2 1B пятая часть модели — просто embedding. Tied embeddings (одна матрица на оба) экономит 18% параметров (с 1.46B до 1.2B) при сопоставимом качестве. Бесплатный обед? Почти — для маленьких моделей точно да.
При untied embeddings: в Llama 3.2 1B это ~20% всех параметров, в 70B — только 3%
RMSNorm и Sandwich Norm
RMSNorm — это LayerNorm на диете: убрали центрирование по среднему, оставили только нормализацию по RMS. Быстрее, проще, а работает не хуже. Arcee Trinity пошли дальше — depth-scaled sandwich norm: нормализация ДО и ПОСЛЕ attention/MLP, причём gain второго RMSNorm = 1/√L. Чем глубже слой — тем мягче нормализация. Учитывает, что активации на разных глубинах живут по-разному.
Sandwich norm: pre-norm (RMSNorm¹) и post-norm (RMSNorm²) вокруг sublayer M_ℓ
Width vs Depth и Embedding Scaling
Глубина vs ширина — вечный спор. При одинаковом бюджете параметров глубокие модели бьют широкие на language modeling. Особенно заметно на мелких моделях. Но широкие проще параллелить на кластере. Embedding scaling (× √d) — используется в Grok-1/2, Trinity Large, Gemma 1-2 для стабилизации residual stream. Без него активации с embedding слоя слишком мелкие по сравнению с остальными.
Embedding scaling: активации эмбеддингов масштабируются на √d для стабильного residual stream
Architecture Decision Heuristics
- Memory/infra-constrained → dense + GQA + RoPE/RNoPE (MoE требует загрузки всех экспертов)
- Inference efficiency at scale → MoE с load balancing, если инфра позволяет
- Long context — ключевое требование → document masking + RoPE scaling (ABF/YaRN) или RNoPE
- Нужны простые kernels и быстрая итерация → избегайте novel attention (MLA), если нет возможности ablate
- Новичок в LLM training → dense, базовый рецепт, focus on basics
💡 Takeaway