Postav si vlastní model

Čtení o transformeru je jedna věc. Vidět, jak ti na vlastním počítači klesá loss a jak se z náhodných znaků postupně stává skoro-čeština, je věc druhá. Tahle kapitola je jeden soubor v NumPy, který si zkopíruješ a za dvacet sekund ti běží. Žádný PyTorch, žádný autograd, forward i backward jsou napsané ručně, takže je v tom vidět úplně všechno.

Co budeš potřebovat

pip install numpy          # to je celé

A nějaký text. Klidně cokoli, kniha, tvoje poznámky, dokumentace. Čím víc, tím líp; pár set kilobajtů stačí. Model níž byl natrénovaný na 200 000 znacích textu téhle učebnice.

Celý model

▶ Spustitelné. Ulož jako model.py a pusť python3 model.py. Potřebuješ jen NumPy a soubor text.txt.

"""Znakový transformer v NumPy. Ruční forward i backward, žádné frameworky.
Spusť:  python3 model.py text.txt"""
import numpy as np, sys

rng = np.random.default_rng(1337)

# ---------- data ----------
text = open(sys.argv[1] if len(sys.argv) > 1 else 'text.txt', encoding='utf8').read()
chars = sorted(set(text))
V = len(chars)                                   # velikost slovníku
stoi = {c: i for i, c in enumerate(chars)}
data = np.array([stoi[c] for c in text])
print(f'{len(data)} znaků, slovník {V}')

T, B, C, H = 16, 32, 48, 48                      # kontext, dávka, d_model, head_dim
MASK = np.triu(np.full((T, T), -1e9), 1)         # kauzální maska

def batch():
    i = rng.integers(0, len(data) - T - 1, B)
    return (np.stack([data[j:j+T] for j in i]),
            np.stack([data[j+1:j+T+1] for j in i]))

def softmax(z):
    e = np.exp(z - z.max(-1, keepdims=True))
    return e / e.sum(-1, keepdims=True)

# ---------- parametry ----------
P = {'E':  rng.normal(0, .02, (V, C)),           # embedding znaků
     'Pos': rng.normal(0, .02, (T, C)),          # poziční embedding
     'Wq': rng.normal(0, .02, (C, H)),
     'Wk': rng.normal(0, .02, (C, H)),
     'Wv': rng.normal(0, .02, (C, H)),
     'Wo': rng.normal(0, .02, (H, C)),           # výstupní projekce hlavy
     'Wu': rng.normal(0, .02, (C, V))}           # unembedding na logits
M = {k: np.zeros_like(v) for k, v in P.items()}  # Adam: první moment
S = {k: np.zeros_like(v) for k, v in P.items()}  # Adam: druhý moment

def forward(x):
    e = P['E'][x] + P['Pos']                                  # (B,T,C)
    q, k, v = e @ P['Wq'], e @ P['Wk'], e @ P['Wv']           # (B,T,H)
    s = q @ k.transpose(0, 2, 1) / np.sqrt(H) + MASK          # (B,T,T)
    a = softmax(s)                                            # váhy pozornosti
    h = e + (a @ v) @ P['Wo']                                 # attention + residual
    return h @ P['Wu'], (x, e, q, k, v, a, h)                 # logits (B,T,V)

def loss_and_grad(logits, y):
    p = softmax(logits)
    n = y.size
    l = -np.log(p.reshape(-1, V)[np.arange(n), y.ravel()] + 1e-9).mean()
    d = p.reshape(-1, V).copy()
    d[np.arange(n), y.ravel()] -= 1
    return l, (d / n).reshape(logits.shape)

def backward(dlogits, cache):
    x, e, q, k, v, a, h = cache
    g = {}
    g['Wu'] = h.reshape(-1, C).T @ dlogits.reshape(-1, V)
    dh = dlogits @ P['Wu'].T
    de = dh.copy()                                            # residual větev
    g['Wo'] = (a @ v).reshape(-1, H).T @ dh.reshape(-1, C)
    do = dh @ P['Wo'].T
    da = do @ v.transpose(0, 2, 1)
    dv = a.transpose(0, 2, 1) @ do
    ds = a * (da - (da * a).sum(-1, keepdims=True)) / np.sqrt(H)   # softmax zpět
    dq, dk = ds @ k, ds.transpose(0, 2, 1) @ q
    for n_, d_ in (('Wq', dq), ('Wk', dk), ('Wv', dv)):
        g[n_] = e.reshape(-1, C).T @ d_.reshape(-1, H)
    de += dq @ P['Wq'].T + dk @ P['Wk'].T + dv @ P['Wv'].T
    g['Pos'] = de.sum(0)
    g['E'] = np.zeros_like(P['E']); np.add.at(g['E'], x.ravel(), de.reshape(-1, C))
    return g

def step(g, t, lr=3e-3):                                       # AdamW light
    for n_ in P:
        M[n_] = .9 * M[n_] + .1 * g[n_]
        S[n_] = .999 * S[n_] + .001 * g[n_] ** 2
        P[n_] -= lr * (M[n_] / (1 - .9**t)) / (np.sqrt(S[n_] / (1 - .999**t)) + 1e-8)

def generate(n=200, temp=.8, seed='Model '):
    ids = [stoi.get(c, 0) for c in seed]
    for _ in range(n):
        ctx = [0] * (T - len(ids[-T:])) + ids[-T:]
        logits, _ = forward(np.array([ctx]))
        ids.append(int(rng.choice(V, p=softmax(logits[0, -1] / temp))))
    return ''.join(chars[i] for i in ids)

# ---------- trénink ----------
for t in range(1, 6001):
    x, y = batch()
    logits, cache = forward(x)
    l, dl = loss_and_grad(logits, y)
    step(backward(dl, cache), t)
    if t % 1000 == 0 or t == 1:
        print(f'krok {t:5d}   loss {l:.3f}')
        if t in (1, 2000, 6000):
            print('   ' + generate(90).replace('\n', ' '))

Skutečný běh

Takhle to dopadlo u mě (procesor, žádná grafika, 17 sekund):

200000 znaků, slovník 120

krok     1   loss 4.787
   Model _;o₃ďYTjR⏩wkŘŽpC;mMai„G️Š.7 w?rA×c2ýČ+wV⚠8;²C️M…3Č→ W⅔jřď+Na→XěčK—VKž0éo..č4E0H⅔5@Řaig

krok  1000   loss 2.703
krok  2000   loss 2.503
   Model tře, kovaběl zelace nachá tohodlené ler, výsto hrosto číže zudijá nélo darý nodní povy

krok  3000   loss 2.490
krok  4000   loss 2.387
krok  5000   loss 2.370
krok  6000   loss 2.355
   Model tor baledůž dostoky jizku 1. Kdybrebez a. U : Detak postí pál model sto. - Mododdě ta

Stojí za to se u toho zastavit:

• Krok 1 je čirá náhoda. Loss 4,79 je přesně ln(120), model hádá rovnoměrně mezi všemi znaky.
• Krok 2000 už drží slabiky a délku slov. „kovaběl“, „tohodlené“ nejsou česká slova, ale mají českou strukturu, model se naučil, jak čeština vypadá.
• Krok 6000 přidává interpunkci na rozumných místech, čísla, dokonce fragmenty skutečných slov („model“, „Kdyby“).

Model má zhruba dvacet tisíc parametrů. Pro srovnání: nejmenší modely, které se dnes dají používat, jich mají miliardu, padesát tisíckrát víc.

Podívej se, kam se dívá

Nejlepší na vlastním modelu je, že do něj vidíš. Matice a ve funkci forward jsou váhy pozornosti, stačí je vytisknout. Tohle je skutečná mapa z natrénovaného modelu pro vstup "pozornost je ":

      p  o  z  o  r  n  o  s  t  _  j  e  _
  p   @  ·  ·  ·  ·  ·  ·  ·  ·  ·  ·  ·  ·
  o   @  .  ·  ·  ·  ·  ·  ·  ·  ·  ·  ·  ·
  z   @     .  ·  ·  ·  ·  ·  ·  ·  ·  ·  ·
  o         @     ·  ·  ·  ·  ·  ·  ·  ·  ·
  r         @  @  -  ·  ·  ·  ·  ·  ·  ·  ·
  n            .     @  ·  ·  ·  ·  ·  ·  ·
  o               =  @  @  ·  ·  ·  ·  ·  ·
  s                     -  @  ·  ·  ·  ·  ·
  t                     .  -  @  ·  ·  ·  ·
  _                              @  ·  ·  ·
  j                              @     ·  ·
  e                                 @     ·
  _                                       @

  řádek = token, který se dívá      sloupec = na koho se dívá
  @ = velká váha, tečka = malá, · = zakázáno maskou

Dvě věci jsou tu vidět na první pohled:

1. Kauzální trojúhelník. Pravá horní půlka je prázdná, maska funguje, nikdo nevidí dopředu.
2. Diagonála je nejsilnější. Znakový model potřebuje hlavně sám sebe a nejbližší okolí.

A když se podíváš na poslední řádek v procentech:

    n    0.0 %
    o    0.0 %
    s    0.1 %
    t    0.2 %
    _    3.3 %  █
    j    0.3 %
    e    7.1 %  ███
    _   89.0 %  ████████████████████████████████████████████

Poslední mezera si 89 % pozornosti nechává pro sebe, 7 % dává předchozímu e a zbytek rozdělí mezi starší kontext. Model se sám rozhodl, že pro předpověď po mezeře je nejdůležitější, že tam mezera je.

Co s tím dělat dál

Nejvíc se naučíš tím, že do toho začneš šťourat. Vyzkoušej:

1. Vyhoď residual, z řádku h = e + (a @ v) @ P['Wo'] udělej h = (a @ v) @ P['Wo']. Co se stane s loss?
2. Vyhoď masku, nastav MASK = np.zeros((T, T)). Loss spadne dramaticky dolů. Proč to ale není zlepšení? (Nápověda: model vidí odpověď.)
3. Vyhoď dělení √H a sleduj, jak se attention mapa změní na téměř binární.
4. Zvětši kontext T z 16 na 64. O kolik se zlepší loss a o kolik se prodlouží krok?
5. Přidej druhou hlavu, rozděl H na dvě půlky, spočítej pozornost zvlášť a výsledky spoj.
6. Přidej druhou vrstvu, zabal attention do funkce a zavolej ji dvakrát za sebou.

Body 5 a 6 z tohohle kódu udělají skutečný (i když malinký) transformer, pak už chybí jen FFN a máš architekturu z anatomie modelu.

Proč zrovna NumPy

Protože v PyTorchi bys napsal loss.backward() a nic bys neviděl. Tady je funkce backward napsaná ručně, řádek po řádku, a je v ní vidět přesně to, co popisuje loss a gradient: gradient teče od loss zpátky, každá operace ho jen „přeloží“ pro tu předchozí. Až tohle jednou napíšeš sám, autograd v jakémkoli frameworku přestane být kouzlo.

Kam dál

Tenhle model nemá RMSNorm, RoPE, SwiGLU ani víc hlav, tedy nic z toho, co mají dnešní modely. V kapitole vylepši si model je všechny postupně přidáme a u každého změříme, co to udělalo s loss. Jeden z nich model překvapivě zhorší.

Shrnutí

• Funkční znakový transformer se vejde do ~90 řádků NumPy a natrénuje se za 17 sekund.
• Loss klesne z 4,79 (náhoda) na 2,36; výstup přejde od náhodných znaků k české struktuře slov.
• Váhy pozornosti se dají vytisknout, vidíš kauzální trojúhelník i to, kam model kouká.
• Nejvíc se naučíš odebíráním součástí: bez residualu, bez masky, bez škálování.