Modele językowe (LLM Basics) – Lab 4

---

Czego nauczysz się w tej części?

• wymienić elementy bloku dekodera
• wyjaśnić sens wielu głów uwagi
• opisać rolę połączenia rezydualnego, normalizacji i MLP
• prześledzić kształty tensora przez mini-transformer
• zbudować przesunięte etykiety do przewidywania kolejnego tokenu
• obliczyć entropię krzyżową dla logitów
• odróżnić uczenie równoległe od generacji autoregresyjnej
• wyjaśnić, co przechowuje KV cache i czego nie usuwa

Nie musisz jeszcze: trenować użytecznego LLM ani implementować rozproszonego uczenia. Mini-model ma cel edukacyjny: wszystkie tensory mieszczą się w pamięci, a wynik nie będzie dobrym chatbotem.

Jak będziemy pracować?

• Krok 1powtórzenie mechanizmu uwagi i mapa bloku dekodera
• Krok 2implementacja mini-bloku w PyTorch
• Krok 3logity, przesunięte etykiety i entropia krzyżowa
• Krok 4krótkie przeuczenie jednej sekwencji
• Krok 5generacja token po tokenie i KV cache
• Krok 6rezultaty pracy

Zanim zaczniesz

Rozrysuj przepływ: embedding → [normalizacja → uwaga → dodanie rezydualne → normalizacja → MLP → dodanie rezydualne] × N → normalizacja → warstwa liniowa do słownika. Oznacz, które operacje mieszają informacje między pozycjami, a które działają niezależnie na każdym tokenie.

Mini-model w tym wpisie celowo przeuczy jedną sekwencję. Spadek funkcji straty pokaże mechanikę celu uczącego, ale nie będzie dowodem generalizacji na nowe dane.

---

Od jednej głowy uwagi do stosu bloków

W poprzedniej części obliczaliśmy wynik jednej głowy uwagi. Dla sekwencji o długości n otrzymywaliśmy macierz [n, d_v]. Każdy jej wiersz był nowym wektorem dla jednej pozycji: ważoną sumą wektorów wartości V z dozwolonych pozycji.

W bloku modelu językowego używamy kilku głów. Jeśli każda z G głów zwraca macierz [n, d_v], łączymy ich wyniki obok siebie, wzdłuż kolumn. Powstaje macierz [n, G · d_v]: nadal mamy n wierszy, ale każdy wiersz zawiera teraz wyniki wszystkich głów. Następnie mnożymy ją przez uczoną macierz projekcji W_O o kształcie [G · d_v, d_model]. Wynik ma kształt [n, d_model].

W naszym przykładzie d_v = d_model / G, więc G · d_v = d_model. Po połączeniu głów liczba kolumn jest już właściwa. Projekcja W_O nadal jest potrzebna: pozwala modelowi tworzyć nowe kombinacje składowych pochodzących z różnych głów. Dokładniejsze obliczenie uwagi wielogłowicowej zobaczymy w następnej sekcji.

Wynik uwagi dodajemy do wejścia tej podwarstwy. Potem macierz trafia do MLP, czyli perceptronu wielowarstwowego (ang. multilayer perceptron). MLP przekształca cały wektor danej pozycji, a nie każdą liczbę w macierzy osobno. Najpierw zwiększa długość wektora z d_model do d_ff, a następnie przywraca d_model. Wynik MLP również dodajemy do wejścia tej podwarstwy. Przed uwagą i przed MLP stosujemy normalizację.

Etap	Kształt dla jednej sekwencji	Co otrzymujemy?
Wejście bloku	[n, d_model]	Jeden wektor dla każdej pozycji.
Wynik każdej głowy uwagi	[n, d_v]	Wektory powstałe przez ważenie informacji z dozwolonych pozycji.
Połączenie wyników G głów	[n, G · d_v]	Wyniki głów ustawione obok siebie w każdym wierszu.
Projekcja przez W_O	[n, d_model]	Nowe kombinacje składowych z różnych głów.
Pierwsze dodanie rezydualne	[n, d_model]	Wynik uwagi dodany do wejścia jej podwarstwy.
Wewnątrz MLP	[n, d_ff]	Każdy wiersz jest przekształcany osobno; jego długość chwilowo rośnie.
Wyjście MLP i drugie dodanie rezydualne	[n, d_model]	MLP wraca do d_model, a jego wynik jest dodawany do wejścia tej podwarstwy.

Cały blok przyjmuje i zwraca macierz [n, d_model]. Dzięki temu możemy ułożyć wiele takich bloków jeden po drugim: wyjście jednego staje się wejściem następnego. W naszym ćwiczeniu użyjemy jednego bloku, aby łatwiej śledzić obliczenia. Pomijamy tu wymiar batcha; w kodzie dojdzie on na początku kształtu każdego tensora.

Co daje połączenie rezydualne?

Połączenie rezydualne oznacza dodanie wejścia podwarstwy do wyniku jej obliczeń. Jeśli na danej pozycji wektor wejściowy wynosi [1.0, 2.0, -1.0], a podwarstwa zwraca [0.3, -0.5, 0.2], po dodaniu dostajemy [1.3, 1.5, -0.8]. Dodajemy wektory składowa po składowej, dlatego muszą mieć tę samą długość.

W bloku robimy to dwa razy: po podwarstwie uwagi i po MLP. Wynik każdej podwarstwy może więc modyfikować reprezentację, a równocześnie istnieje bezpośrednia ścieżka od jej wejścia do wyjścia. Taka ścieżka ułatwia również przepływ gradientu podczas uczenia głębokiego stosu bloków.

W użytym dalej układzie pre-norm pierwszy krok można zapisać jako R = Z + Uwaga(Norm1(Z)), a drugi jako Y = R + MLP(Norm2(R)). Do wyniku MLP dodajemy zatem R – reprezentację sprzed MLP – a nie pierwotne Z.

Co robi normalizacja i jaki daje efekt?

Normalizacja działa na składowych jednego wektora na danej pozycji. Nie uśrednia tokenów w sekwencji ani przykładów w batchu. Pomaga utrzymać kontrolowaną skalę wartości przekazywanych do kolejnych obliczeń. W naszym bloku normalizujemy wektory przed podwarstwą uwagi i przed MLP; stąd nazwa pre-norm.

Prosty przykład LayerNorm: dla wektora [1.0, 3.0] średnia składowych wynosi 2.0, a odchylenie standardowe 1.0. Po odjęciu średniej i podzieleniu przez odchylenie otrzymalibyśmy [-1.0, 1.0]. Pomijamy tutaj niewielką stałą zabezpieczającą obliczenie oraz parametry uczone przez warstwę, które pozwalają później zmienić skalę i przesunięcie wyniku.

RMSNorm działa inaczej: skaluje wektor na podstawie średniej kwadratów jego składowych, ale nie odejmuje ich średniej. Dla tego samego wektora [1.0, 3.0] i przy pominięciu parametrów uczonych dzielilibyśmy obie składowe przez √5 ≈ 2.236, otrzymując około [0.447, 1.342]. Obie metody pomagają kontrolować skalę reprezentacji, choć wykonują różne obliczenia.

Dlaczego MLP ma dwie warstwy i po co wymiar d_ff?

W naszym mini-modelu MLP składa się z dwóch warstw liniowych oraz funkcji aktywacji GELU pomiędzy nimi. Pierwsza warstwa przyjmuje wektor o długości d_model i tworzy wektor o długości d_ff. Druga przekształca go z powrotem do d_model. Dla wartości użytych później w kodzie jest to przejście 16 → 64 → 16.

Każda składowa wektora po pierwszej warstwie może zależeć od wszystkich składowych wejściowego wektora tej pozycji. Większy wymiar pośredni daje miejsce na różne wyuczone kombinacje tych cech. Funkcja GELU wprowadza nieliniowość: bez niej dwie następujące po sobie warstwy liniowe można byłoby zastąpić jedną warstwą liniową.

Druga warstwa przywraca długość d_model, aby wynik MLP dało się dodać do jego wejścia i przekazać do następnego bloku. To samo MLP, z tymi samymi wagami, działa osobno na każdym wierszu macierzy. Nie pobiera bezpośrednio wektorów z innych pozycji, choć wektor wejściowy może już zawierać informacje z wcześniejszych tokenów dzięki podwarstwie uwagi.

Prześledź liczby dla jednej pozycji

Weźmy mały przykład: n = 3, d_model = 4, dwie głowy, d_v = 2 dla każdej głowy oraz d_ff = 8. Śledzimy pozycję 2, która może korzystać z pozycji 0, 1, 2. Liczby są dobrane tylko po to, aby pokazać rachunek.

1. Pierwsza głowa. Załóżmy, że jej wagi dla pozycji 2 wynoszą [0.2, 0.5, 0.3], a wektory wartości to [1.0, 0.0], [0.0, 1.0] i [0.0, 1.0]. Ważona suma wynosi [0.2, 0.8].
2. Druga głowa. Niech wagi wynoszą [0.5, 0.25, 0.25], a jej wektory wartości: [2.0, 0.0], [0.0, -1.0] i [0.0, -1.0]. Wynik drugiej głowy to [1.0, -0.5].
3. Połączenie głów. Ustawiamy ich wyniki obok siebie: [0.2, 0.8, 1.0, -0.5]. Wektor ma teraz 2 · 2 = 4 składowe. Uczona projekcja W_O może mieszać składowe obu głów. Przykładowo, jeśli w tym uproszczeniu tworzy wektor [c₀ + c₂, c₁ + c₃, c₂, c₃], otrzymujemy [1.2, 0.3, 1.0, -0.5].
4. Pierwsze dodanie rezydualne. Jeśli wejściowy wektor pozycji 2 wynosił [0.5, -0.5, 1.0, 0.0], dodajemy do niego wynik projekcji. Dostajemy [1.7, -0.2, 2.0, -0.5].
5. MLP. Po normalizacji MLP przekształca ten czteroskładowy wektor najpierw w wektor o długości 8, a potem znów o długości 4. Załóżmy, że zwraca [0.1, 0.2, -0.3, 0.4]. Te liczby zależałyby od uczonych wag obu warstw MLP.
6. Drugie dodanie rezydualne. Dodajemy wynik MLP do wektora sprzed MLP: [1.7, -0.2, 2.0, -0.5] + [0.1, 0.2, -0.3, 0.4] = [1.8, 0.0, 1.7, -0.1]. Jest to wyjściowy wektor bloku dla pozycji 2.

Takie obliczenie wykonujemy dla każdej pozycji. Wyjściem całego bloku w tym przykładzie jest więc macierz [3, 4]: trzy zaktualizowane wektory, każdy o czterech składowych.

Po ostatnim bloku osobna głowa słownikowa (LM head) zamienia końcowy wektor każdej pozycji na logity dla całego słownika. Jest częścią całego modelu językowego, ale nie każdego powtarzanego bloku.

Oryginalny transformer miał enkoder i dekoder oraz dodatkowy mechanizm uwagi łączący obie części. Tutaj zajmujemy się modelem typu decoder-only, który nie ma osobnego enkodera.

Wiele głów tworzy różne sposoby dopasowania

W poprzedniej sekcji zobaczyliśmy, jak połączyć wyniki kilku głów. Teraz przyjrzyjmy się temu, dlaczego ich wyniki mogą być różne. Wszystkie głowy otrzymują tę samą macierz wejściową Z, ale każda ma własne uczone macierze projekcji W_Q, W_K i W_V. Może więc tworzyć inne zapytania, klucze i wartości, a w konsekwencji inaczej rozdzielać wagi między dozwolone pozycje.

Macierz Z ma kształt [n, d_model]. W pierwszym bloku jest to reprezentacja tokenów z informacją o pozycji – macierz, którą w części 02A oznaczaliśmy przez H. W kolejnych blokach Z oznacza wynik poprzedniego bloku. Literą G oznaczamy liczbę głów, aby nie używać ponownie symbolu H.

Dla głowy o numerze g mnożymy Z przez trzy osobne macierze parametrów. Otrzymujemy Q, K i V tej konkretnej głowy:

Zapytania i klucze głowy g wyznaczają jej własną macierz wyników dopasowania [n, n]. Po nałożeniu maski przyczynowej i zastosowaniu softmax głowa otrzymuje wagi, którymi łączy swoje wektory wartości. To ten sam mechanizm, który obliczaliśmy w części 02A – teraz wykonujemy go osobno dla każdej głowy.

Co mogą zmienić osobne projekcje? Dla zapytania z pozycji 2 jedna głowa może na przykład przydzielić pozycjom 0, 1, 2 wagi [0.2, 0.5, 0.3], a druga [0.5, 0.25, 0.25]. Pierwsza w większym stopniu wykorzysta wtedy pozycję 1, a druga pozycję 0. Ponieważ głowy mają też osobne projekcje V, różnić się mogą nie tylko wagi, lecz również wektory, które są za ich pomocą łączone.

W naszym uproszczeniu przyjmujemy d_head = d_model / G oraz d_v = d_head. Każda głowa zwraca wtedy macierz [n, d_head]. Łączymy wyniki głów wzdłuż kolumn i mnożymy przez uczoną macierz W_O:

Połączenie G macierzy [n, d_head] daje [n, G · d_head] = [n, d_model]. Macierz W_O ma w tym przykładzie kształt [d_model, d_model]: miesza składowe wyników różnych głów, zachowując wymiar potrzebny do dodania rezydualnego i przekazania wyniku dalej.

Obiekt	Kształt przy n = 5, d_model = 16 i G = 4	Znaczenie
Z	[5, 16]	Pięć wektorów wejściowych, po jednym na pozycję.
W_Q^(g), W_K^(g), W_V^(g)	Każda [16, 4]	Parametry projekcji jednej głowy; każda projekcja korzysta z całego wektora o długości 16.
Q^(g), K^(g), V^(g)	Każda [5, 4]	Wektory zapytań, kluczy i wartości wyznaczone przez tę głowę.
Wagi jednej głowy	[5, 5]	Dla każdej pozycji zapytania: wagi przypisane pozycjom źródłowym.
Wynik jednej głowy	[5, 4]	Pięć wektorów, każdy o czterech składowych.
Połączenie wyników czterech głów	[5, 16]	Cztery czteroskładowe wyniki ustawione obok siebie w każdym wierszu.
W_O i wynik całego mechanizmu	[16, 16] oraz [5, 16]	Projekcja miesza wyniki głów i zwraca wektory o długości d_model.

Czy każda głowa widzi tylko fragment wektora wejściowego?

Nie. W powyższym przykładzie każda głowa oblicza swoje czteroskładowe Q, K i V z całego wejściowego wektora o długości 16. Mówią o tym kształty macierzy projekcji [16, 4]. Liczba 4 jest długością wektora otrzymanego po projekcji, a nie liczbą składowych wejścia widocznych dla głowy.

Biblioteka może dla wydajności obliczać projekcje wszystkich głów za pomocą większych, wspólnych operacji macierzowych, a potem dzielić wynik na głowy. Nie zmienia to przedstawionej tu interpretacji: każda głowa ma własny zestaw uczonych parametrów projekcji.

W użytej niżej klasie nn.MultiheadAttention wartość d_model musi być podzielna przez liczbę głów G. Dla d_model = 16 i G = 4 projekcje każdej głowy zwracają wektory o czterech składowych. Różne wagi głów można analizować, ale nie należy z góry przypisywać każdej głowie jednej stałej, łatwej do nazwania funkcji językowej.

Złóż uwagę wielogłowicową z mnożeń macierzy

Zanim użyjemy nn.MultiheadAttention, wykonajmy jeden krok uwagi wielogłowicowej za pomocą zwykłych operacji na tensorach. Nie będziemy tutaj uczyć parametrów. Macierze projekcji wylosujemy raz, aby zobaczyć, co i przez co mnożymy oraz jaki kształt ma wynik każdego działania.

Dla czytelności pomijamy wymiar batcha. Mamy n = 3 pozycje, wektory wejściowe długości d_model = 4 oraz G = 2 głowy. Każda głowa tworzy zapytania, klucze i wartości długości d_head = 2. Każda jej macierz projekcji ma więc kształt [4, 2]: bierze cały czteroskładowy wektor pozycji i zwraca wektor dwuskładowy.

Obliczenie	Kształt	Co oznacza wynik?
Wejście Z	[3, 4]	Trzy pozycje, każda reprezentowana przez wektor długości cztery.
Projekcje Q, K, V jednej głowy	[3, 4] @ [4, 2] → [3, 2]	Trzy nowe wektory dwuskładowe dla każdej z trzech ról.
Wyniki dopasowania jednej głowy	[3, 2] @ [2, 3] → [3, 3]	Porównanie zapytania z każdej pozycji z kluczem z każdej pozycji.
Wynik jednej głowy	[3, 3] @ [3, 2] → [3, 2]	Ważone sumy wektorów wartości.
Połączenie dwóch głów	[3, 2] + [3, 2] → [3, 4]	Wyniki ustawione obok siebie w kolumnach – znak + oznacza tu złączenie, nie dodawanie liczb.
Projekcja wyjściowa	[3, 4] @ [4, 4] → [3, 4]	Mieszanie składowych z obu głów; wynik gotowy do dodania rezydualnego.

import math
import torch

torch.manual_seed(7)

# Trzy pozycje, po cztery składowe na pozycję.
z = torch.tensor([
    [1.0, 0.0, 1.0, 0.0],
    [0.0, 1.0, 0.0, 1.0],
    [1.0, 1.0, 0.0, 0.0],
])  # [n=3, d_model=4]

n, d_model = z.shape
n_heads = 2
assert d_model % n_heads == 0
d_head = d_model // n_heads

# Pierwszy wymiar wybiera głowę.
# Każda głowa ma własne projekcje całego wektora wejściowego.
w_q = torch.randn(n_heads, d_model, d_head) * 0.5
w_k = torch.randn(n_heads, d_model, d_head) * 0.5
w_v = torch.randn(n_heads, d_model, d_head) * 0.5

# Projekcja wyjściowa miesza składowe po połączeniu głów.
w_o = torch.randn(n_heads * d_head, d_model) * 0.5

# Dla każdej głowy: [n, d_model] @ [d_model, d_head].
q = torch.stack([z @ w_q[g] for g in range(n_heads)])
k = torch.stack([z @ w_k[g] for g in range(n_heads)])
v = torch.stack([z @ w_v[g] for g in range(n_heads)])
# q, k, v: [n_heads, n, d_head]

# Osobna macierz dopasowań dla każdej głowy.
scores = q @ k.transpose(-2, -1) / math.sqrt(d_head)
# scores: [n_heads, n, n]

future = torch.triu(
    torch.ones(n, n, dtype=torch.bool),
    diagonal=1,
)  # [n, n]

# Ta sama maska pozycji przyszłych obowiązuje w obu głowach.
masked_scores = scores.masked_fill(future.unsqueeze(0), float("-inf"))
weights = torch.softmax(masked_scores, dim=-1)
# weights: [n_heads, n, n]

heads = weights @ v
# heads: [n_heads, n, d_head]

# Ustaw wyniki głów obok siebie dla każdej pozycji.
combined = heads.transpose(0, 1).reshape(n, n_heads * d_head)
# combined: [n, n_heads * d_head] = [3, 4]

output = combined @ w_o
# output: [n, d_model] = [3, 4]

torch.set_printoptions(precision=3, sci_mode=False)
print("Q:", q.shape, "K:", k.shape, "V:", v.shape)
print("wyniki dopasowania:", scores.shape)
print("wagi:", weights.shape)
print("wagi pierwszej głowy:\n", weights[0])
print("wagi drugiej głowy:\n", weights[1])
print("wyniki głów:", heads.shape)
print("po połączeniu:", combined.shape)
print("po projekcji wyjściowej:", output.shape)

assert torch.allclose(
    weights.sum(dim=-1),
    torch.ones_like(weights.sum(dim=-1)),
)
assert torch.all(weights[:, future] == 0)

Jak czytać kod? Wymiary q, k i v to [głowy, pozycje, składowe]. Operator @ oblicza osobne mnożenie macierzy dla każdej głowy. Dlatego q @ k.transpose(-2, -1) zwraca dwie macierze dopasowania [3, 3], a nie jedną wspólną macierz dla wszystkich głów.

Funkcja softmax działa po ostatnim wymiarze, czyli osobno na każdym wierszu każdej głowy. Następnie każda głowa mnoży swoją macierz wag przez swoje V. Dopiero wyniki głów są ustawiane obok siebie i przekształcane przez w_o. Otrzymany tensor ma taki sam kształt jak z, co pozwoli nam w następnym kroku wykonać dodanie rezydualne.

1. Odczytaj q[0, 2] i k[0, 1]. Oblicz ręcznie ich iloczyn skalarny, podziel przez √2 i porównaj z scores[0, 2, 1].
2. Porównaj wagi w wierszu 2 obu głów. Czy mają takie same wartości? Z czego może wynikać różnica?
3. Sprawdź scores[0, 0, 2] oraz weights[0, 0, 2]. Dlaczego pierwsza wartość jest liczbą, a druga wynosi 0?
4. Zastąp w_o przez torch.eye(d_model) i ponownie oblicz output. Porównaj go z combined. Co robiła projekcja wyjściowa?

Co powinno wyjść?

q, k i v mają kształt [2, 3, 2]. Wyniki dopasowania i wagi mają kształt [2, 3, 3], a wyniki głów [2, 3, 2]. Po ustawieniu wyników głów obok siebie otrzymujemy [3, 4]; projekcja wyjściowa zachowuje ten kształt.

Przed nałożeniem maski scores[0, 0, 2] jest obliczonym wynikiem dopasowania. Pozycja 0 nie może jednak korzystać z późniejszej pozycji 2, więc przed softmax zastępujemy ten wynik przez −∞. Odpowiadająca mu waga wynosi 0. Po ustawieniu w_o na macierz jednostkową otrzymujemy output = combined: projekcja przestaje mieszać składowe złączonych głów.

To dydaktyczna implementacja samego przejścia w przód dla jednej sekwencji. Macierze projekcji są tu wylosowane i nie są uczone. W kolejnej sekcji nn.MultiheadAttention wykona odpowiadające operacje wewnątrz gotowej warstwy, z uwzględnieniem batcha.

Mini-blok pre-norm w PyTorch

Zapiszemy teraz blok dekodera jako klasę PyTorch. W układzie pre-norm najpierw normalizujemy wejście podwarstwy uwagi, obliczamy jej wynik i dodajemy go do wejścia. Następnie normalizujemy otrzymaną reprezentację, przepuszczamy ją przez MLP i wykonujemy drugie dodanie rezydualne. Dropout pomijamy dla czytelności.

Wejście x ma kształt [B, n, d_model]. Dzięki batch_first=True pierwszy wymiar oznacza batch, drugi pozycje w sekwencji, a trzeci składowe wektorów. Maska boolowska ma wartość True dla par pozycji, których nie wolno połączyć.

import torch
from torch import nn


class TinyDecoderBlock(nn.Module):
    def __init__(self, d_model: int, n_heads: int, d_ff: int) -> None:
        super().__init__()

        self.norm1 = nn.LayerNorm(d_model)
        self.attn = nn.MultiheadAttention(
            embed_dim=d_model,
            num_heads=n_heads,
            batch_first=True,
        )

        self.norm2 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Linear(d_ff, d_model),
        )

    def forward(
        self, x: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        seq_len = x.shape[1]

        # True nad przekątną: pozycja i nie może użyć pozycji j > i.
        future = torch.triu(
            torch.ones(
                seq_len,
                seq_len,
                device=x.device,
                dtype=torch.bool,
            ),
            diagonal=1,
        )

        # Pierwsza podwarstwa: normalizacja, uwaga, dodanie rezydualne.
        attention_input = self.norm1(x)
        attended, weights = self.attn(
            attention_input,
            attention_input,
            attention_input,
            attn_mask=future,
            need_weights=True,
            average_attn_weights=False,
        )
        after_attention = x + attended

        # Druga podwarstwa: normalizacja, MLP, dodanie rezydualne.
        mlp_input = self.norm2(after_attention)
        mlp_output = self.mlp(mlp_input)
        output = after_attention + mlp_output

        return output, weights


block = TinyDecoderBlock(d_model=16, n_heads=4, d_ff=64)
x = torch.randn(2, 5, 16)  # B=2, n=5, d_model=16

y, weights = block(x)

print("wejście:", x.shape)
print("wyjście:", y.shape)
print("wagi uwagi:", weights.shape)

assert y.shape == x.shape

Dlaczego trzy razy przekazujemy attention_input? Jest to uwaga własna (ang. self attention): zapytania, klucze i wartości powstają z tej samej sekwencji. Wewnątrz nn.MultiheadAttention działają jednak osobne, uczone projekcje, więc powstałe Q, K i V nie są automatycznie identyczne. Biblioteka wykonuje też połączenie wyników głów i projekcję wyjściową omawianą w poprzedniej sekcji.

Śledź tensor przez blok

Dodaj tymczasowe instrukcje print albo hooki i zapisz kształty: attention_input, attended, weights, after_attention, mlp_input, wyniku pierwszej i drugiej warstwy liniowej MLP oraz output. Warstwy liniowe MLP znajdziesz pod indeksami block.mlp[0] i block.mlp[2].

1. W której warstwie MLP długość wektora rośnie z 16 do 64, a w której wraca do 16?
2. Dlaczego attended musi mieć taki sam kształt jak x?
3. Dlaczego tensor weights ma dodatkowy wymiar, którego nie ma output?

Oczekiwane kształty i odpowiedzi

Tensor	Kształt
x	[2, 5, 16]
future	[5, 5]
attention_input	[2, 5, 16]
attended	[2, 5, 16]
weights	[2, 4, 5, 5]
after_attention	[2, 5, 16]
mlp_input	[2, 5, 16]
Po block.mlp[0]	[2, 5, 64]
Po block.mlp[2]	[2, 5, 16]
output	[2, 5, 16]

Pierwsza warstwa liniowa MLP zmienia ostatni wymiar z 16 na 64; druga przywraca 16. Kształty attended i x muszą być zgodne, ponieważ dodajemy te tensory składowa po składowej. W weights wymiar 4 rozróżnia głowy uwagi: dla każdej sekwencji i głowy dostajemy macierz wag [5, 5]. Ich wyniki są następnie łączone, dlatego output znów ma po jednym wektorze długości 16 na pozycję.

Jak model uczy się przewidywać następny token?

Załóżmy, że model otrzymuje początek zdania „Ala ma”. Jego zadaniem może być przewidzenie kolejnego tokenu, na przykład „kota”. Podczas uczenia sprawdzamy jednak nie tylko koniec zdania: na każdej pozycji sekwencji model przewiduje token, który występuje zaraz po niej.

Poprawne odpowiedzi otrzymujemy przez przesunięcie sekwencji tokenów o jedną pozycję. Dla czytelności przypiszemy identyfikatorom konkretne symbole: 0 = <BOS> oznacza początek sekwencji, 1 = Ala, 3 = ma, 4 = kota, a 2 = <EOS> oznacza koniec sekwencji. Pełne mapowanie na słownik 20 tokenów pojawi się w kodzie:

Element	Identyfikatory tokenów
Pełna sekwencja	[0, 1, 3, 4, 8, 7, 5, 6, 9, 10, 11, 12, 5, 2]
Wejście modelu	[0, 1, 3, 4, 8, 7, 5, 6, 9, 10, 11, 12, 5]
Poprawne następne tokeny	[1, 3, 4, 8, 7, 5, 6, 9, 10, 11, 12, 5, 2]

W tym przykładzie pełny tekst brzmi: „Ala ma kota i psa. Ola czyta książkę o zwierzętach.”. Po znaczniku <BOS> poprawną odpowiedzią jest Ala. Po <BOS> Ala poprawną odpowiedzią jest ma. Po pierwszym zdaniu kolejnym tokenem jest Ola, a po ostatniej kropce – <EOS>. Maska przyczynowa zapewnia, że przy przewidywaniu na pozycji i model korzysta tylko z tokenów na pozycjach j ≤ i. Choć do modelu podajemy jednocześnie trzynaście tokenów, pierwszy wynik nie może wykorzystać dwunastu późniejszych.

Od wektorów w bloku do logitów słownikowych

W poprzedniej sekcji sam blok dekodera zwracał y, weights = block(x). Wektorów y nie należy mylić z identyfikatorami przewidywanych tokenów: są to nowe reprezentacje pozycji, nadal o długości d_model. Z kolei weights zawiera wagi uwagi, które mówią, jak pozycje w sekwencji łączą informacje. Żaden z tych dwóch wyników nie jest jeszcze rozkładem po tokenach słownika. Dla czytelności dalej nazwiemy je block_output i attention_weights.

Etap w naszym przykładzie	Kształt	Co zawiera?
input_ids	[1, 13]	Trzynaście ID tokenów w jednej sekwencji.
Embeddingi tokenów + wektory pozycji	[1, 13, 16]	Po jednym wektorze długości d_model = 16 na pozycję.
attention_weights	[1, 4, 13, 13]	W każdej z czterech głów: wagi łączące trzynaście pozycji zapytań z trzynastoma pozycjami źródłowymi.
block_output	[1, 13, 16]	Reprezentacje po przejściu przez uwagę, połączenia resztkowe i MLP.
final_norm(block_output)	[1, 13, 16]	Wektory po końcowej normalizacji.
lm_head(...) → logits	[1, 13, 20]	Na każdej z trzynastu pozycji: po jednym wyniku dla każdego z 20 tokenów słownika.
targets	[1, 13]	Trzynaście ID poprawnych następnych tokenów.

Na przykład attention_weights[0, 0, 2] opisuje, jak pierwsza głowa uwagi rozdziela wagi między pozycje sekwencji przy obliczaniu wyniku dla pozycji 2. Późniejsze pozycje mają wtedy wagę zero. Natomiast logits[0, 2] zawiera 20 surowych wyników dla możliwych następnych tokenów. Te dwie tablice opisują różne rzeczy. Wagi uwagi powstają wewnątrz bloku; logity słownikowe powstają dopiero po zastosowaniu lm_head do końcowych reprezentacji. Większy model może przekazywać wynik jednego bloku do następnego, zanim wyznaczy logity.

Każdy wektor logitów można przeliczyć funkcją softmax na prawdopodobieństwa tokenów słownika. Podczas uczenia interesuje nas przede wszystkim to, jakie prawdopodobieństwo otrzymał poprawny następny token. Mierzymy je za pomocą entropii krzyżowej (ang. cross-entropy).

Tutaj tᵢ₊₁ to poprawny następny token po pozycji i, a pᵢ(tᵢ₊₁) to prawdopodobieństwo przypisane mu przez model na tej pozycji. n oznacza liczbę ocenianych pozycji; w pierwszym przykładzie n = 13. Liczymy stratę dla każdej z nich i uśredniamy. Symbol pᵢ oznacza rozkład po tokenach słownika, a nie wagi uwagi.

Trzy przewidywania w jednym przebiegu

Aby dało się zobaczyć cały rozkład prawdopodobieństwa, na chwilę użyjmy osobnego, mniejszego przykładu ze słownikiem złożonym z czterech symboli: 0, 1, 2, 3. Pełna sekwencja to [1, 3, 0, 2]. Podajemy wejście [1, 3, 0] i oczekujemy kolejno tokenów [3, 0, 2]. Numery w tej tabeli należą do tego małego, niezależnego przykładu. Poniższe prawdopodobieństwa są ilustracyjne: pokazują, jak odczytać trzy wiersze wyników po softmax, a nie wyjście nieuczonego modelu z kodu poniżej.

Pozycja	Dostępne tokeny wejścia	Prawdopodobieństwa tokenów [0, 1, 2, 3]	Poprawny token	Strata pozycji
0	[1]	[0,2; 0,1; 0,1; 0,6]	3	−ln(0,6)
1	[1, 3]	[0,8; 0,1; 0,05; 0,05]	0	−ln(0,8)
2	[1, 3, 0]	[0,1; 0,1; 0,7; 0,1]	2	−ln(0,7)

Średnia strata z tych trzech pozycji wynosi (−ln(0,6) − ln(0,8) − ln(0,7)) / 3 ≈ 0,364. Nie uruchamiamy dodatkowo modelu osobno dla wejść [1] i [1, 3]: dzięki masce przyczynowej pierwszy i drugi wiersz jednego przebiegu odpowiadają właśnie przewidywaniu z tych prefiksów. Osobne wywołania powieliłyby te same oceny.

Czy potrzebujemy kodowania one-hot i osobnego softmax?

Nie w tej implementacji. W czteroelementowym przykładzie token o ID 3 można obrazowo zapisać jako wektor one-hot [0, 0, 0, 1], ale nn.Embedding przyjmuje samo ID i wybiera odpowiedni wektor z tabeli embeddingów. Również poprawną odpowiedź przekazujemy jako ID tokenu, np. 3, a nie jako wektor one-hot.

F.cross_entropy przyjmuje surowe logity i ID poprawnych tokenów. Wewnątrz oblicza odpowiednik log_softmax oraz stratę dla wskazanych ID, więc wcześniej nie stosujemy softmax. W tabeli powyżej pokazaliśmy prawdopodobieństwa jedynie po to, aby rachunek −ln(p) był czytelny. Logity nie są wagami uwagi ani parametrami modelu.

Teraz złożymy prosty model z embeddingów, jednego bloku dekodera i głowy słownikowej. Kod jest kontynuacją poprzedniej sekcji: torch, nn i TinyDecoderBlock są już dostępne.

Poniższe mapowanie służy wyłącznie ćwiczeniu: przypisujemy ręcznie całe wyrazy i znaki interpunkcyjne do ID. W prawdziwym modelu takie przypisanie przygotowuje tokenizer, a token może być także fragmentem wyrazu.

import torch.nn.functional as F

class TinyCausalLM(nn.Module):
    def __init__(
        self,
        vocab_size: int = 20,
        max_len: int = 32,
        d_model: int = 16,
    ) -> None:
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, d_model)
        self.position_embedding = nn.Embedding(max_len, d_model)
        self.block = TinyDecoderBlock(d_model, n_heads=4, d_ff=64)
        self.final_norm = nn.LayerNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        positions = torch.arange(
            input_ids.shape[1],
            device=input_ids.device,
        )
        token_vectors = self.token_embedding(input_ids)
        position_vectors = self.position_embedding(positions).unsqueeze(0)
        block_input = token_vectors + position_vectors  # [B, n, d_model]

        block_output, _ = self.block(block_input)  # [B, n, d_model]
        normalized = self.final_norm(block_output)
        return self.lm_head(normalized)  # [B, n, vocab_size]

id_to_token = {
    0: "<BOS>", 1: "Ala", 2: "<EOS>", 3: "ma", 4: "kota",
    5: ".", 6: "Ola", 7: "psa", 8: "i", 9: "czyta",
    10: "książkę", 11: "o", 12: "zwierzętach", 13: "lubi",
    14: "student", 15: "planuje", 16: "naukę", 17: "jutro",
    18: "?", 19: "<PAD>",
}

def decode(token_ids: torch.Tensor) -> str:
    words = []
    for token_id in token_ids.tolist():
        token = id_to_token[token_id]
        if token == "<EOS>":
            break
        if token in {"<BOS>", "<PAD>"}:
            continue
        if token in {".", ",", "?", "!"}:
            if words:
                words[-1] += token
        else:
            words.append(token)
    return " ".join(words)

torch.manual_seed(3)
model = TinyCausalLM(vocab_size=len(id_to_token))

sequence = torch.tensor([[0, 1, 3, 4, 8, 7, 5, 6, 9, 10, 11, 12, 5, 2]])
inputs = sequence[:, :-1]   # [1, 13]
targets = sequence[:, 1:]   # [1, 13]

logits = model(inputs)  # [1, 13, 20]

losses_by_position = F.cross_entropy(
    logits.reshape(-1, logits.shape[-1]),  # [13, 20]
    targets.reshape(-1),                   # [13]
    reduction="none",
)
loss = losses_by_position.mean()

print("wejście:", inputs)
print("cele:   ", targets)
print("tekst pełnej sekwencji:", decode(sequence[0]))
print("logity:", logits.shape)
print("straty na pozycjach:", losses_by_position)
print("średnia strata:", loss.item())

Model zwraca tensor [1, 13, 20]: jedna sekwencja, trzynaście pozycji i po 20 logitów na każdej pozycji. Po zmianie kształtów przekazujemy do F.cross_entropy trzynaście wierszy logitów [13, 20] i trzynaście ID poprawnych tokenów [13]. Parametr reduction="none" pozwala obejrzeć osobne straty. Bierzemy ich średnią; bez tego parametru funkcja domyślnie zwróciłaby średnią od razu.

Na razie tylko mierzymy stratę: samo wywołanie F.cross_entropy nie zmienia parametrów. Model zacznie się uczyć dopiero wtedy, gdy wywołamy loss.backward(), aby obliczyć gradienty, a następnie optimizer.step(), aby zaktualizować parametry.

Przesunięcie sekwencji mówi modelowi, co ma przewidzieć. Maska przyczynowa określa, które tokeny wolno mu przy tym wykorzystać. Strata łączy oceny ze wszystkich przewidywanych pozycji.

Uczenie jest równoległe, generacja jest autoregresyjna

Podczas uczenia znamy całą sekwencję i obliczamy przewidywania dla wielu pozycji w jednym przebiegu. Po uśrednieniu strat z bieżącej sekwencji lub batcha wywołujemy propagację wsteczną i wykonujemy jedną aktualizację parametrów. Powtarzamy ten proces w kolejnych krokach treningowych. Nie aktualizujemy parametrów osobno po każdej pozycji tej samej sekwencji.

Podczas generacji kolejny token nie jest jeszcze znany. Bierzemy logity z ostatniej pozycji, wybieramy token, dopisujemy go i dopiero wtedy wykonujemy kolejny krok. Przy losowaniu zamieniamy logity na rozkład prawdopodobieństwa; przy wyborze tokenu o największym logicie softmax nie jest potrzebny.

Celowo przeucz mini-model

Wykonaj około dwustu kroków optymalizacji na jednej sekwencji. Strata powinna wyraźnie spaść. Obejrzyj przewidywania dla znanych prefiksów, a potem podaj modelowi pierwsze zdanie jako prompt i wygeneruj drugie zdanie token po tokenie. Następnie zmień token w pierwszym zdaniu i porównaj otrzymaną kontynuację. Jedna zapamiętana sekwencja nie wystarczy, by oczekiwać poprawnego uogólnienia.

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3)
model.train()

for step in range(201):
    optimizer.zero_grad()
    logits = model(inputs)
    loss = F.cross_entropy(
        logits.reshape(-1, logits.shape[-1]),
        targets.reshape(-1),
    )  # średnia strata z trzynastu pozycji
    loss.backward()   # gradient średniej straty
    optimizer.step()  # jedna aktualizacja w tym kroku

    if step % 40 == 0:
        print(step, round(loss.item(), 4))

model.eval()
@torch.no_grad()
def greedy_continue(prompt_ids: torch.Tensor, max_new_tokens: int = 10) -> torch.Tensor:
    generated = prompt_ids.clone()
    available_steps = model.position_embedding.num_embeddings - generated.shape[1]
    for _ in range(min(max_new_tokens, available_steps)):
        next_logits = model(generated)[:, -1, :]
        next_id = next_logits.argmax(dim=-1, keepdim=True)
        generated = torch.cat((generated, next_id), dim=1)
        if next_id.item() == 2:  # <EOS>
            break
    return generated

with torch.no_grad():
    predictions = model(inputs).argmax(dim=-1)

prompt_ids = sequence[:, :7].clone()  # <BOS> Ala ma kota i psa .
generated = greedy_continue(prompt_ids)

print("przewidywania dla znanych prefiksów:",
      [id_to_token[int(i)] for i in predictions[0]])
print("poprawne następne tokeny:",
      [id_to_token[int(i)] for i in targets[0]])
print("prompt:", decode(prompt_ids[0]))
print("oczekiwana kontynuacja:", decode(sequence[0, prompt_ids.shape[1]:]))
print("wygenerowana kontynuacja:", decode(generated[0, prompt_ids.shape[1]:]))
print("cały wygenerowany tekst:", decode(generated[0]))

Lista predictions dotyczy pozycji wejścia, które model już znał podczas obliczenia; pokazuje przewidywania dla kolejnych poprawnych prefiksów. Funkcja greedy_continue wykonuje właściwą generację: zaczyna od całego pierwszego zdania, za każdym razem pobiera logity z ostatniej pozycji, dopisuje wybrany token i kończy po <EOS> albo po ustalonej liczbie kroków. Funkcja decode zamienia identyfikatory na czytelny tekst i dołącza kropkę do poprzedniego wyrazu.

Wpisz własny prompt z naszego słownika

Teraz sprawdź model na tekście, którego nie widział podczas treningu. Możesz wpisać na przykład Ola lubi kota. albo student planuje naukę jutro.. Używaj dokładnie słów widocznych na liście, z zachowaniem wielkich liter w imionach. To prosta zamiana całych wyrazów na ID, więc nie rozpozna innych form gramatycznych ani słów spoza słownika.

token_to_id = {
    token: token_id
    for token_id, token in id_to_token.items()
    if token not in {"<BOS>", "<EOS>", "<PAD>"}
}

def encode(text: str) -> torch.Tensor:
    pieces = text.replace(".", " . ").replace("?", " ? ").split()
    if not pieces:
        raise ValueError("Wpisz przynajmniej jeden token.")
    unknown = sorted({piece for piece in pieces if piece not in token_to_id})
    if unknown:
        raise ValueError(f"Tokeny spoza słownika: {unknown}")
    ids = [0] + [token_to_id[piece] for piece in pieces]  # dodaj <BOS>
    if len(ids) >= model.position_embedding.num_embeddings:
        raise ValueError("Prompt jest za długi dla tego mini-modelu.")
    return torch.tensor([ids], dtype=torch.long)

print("Dostępne tokeny:", ", ".join(sorted(token_to_id)))
user_text = input("Wpisz prompt: ")
user_prompt_ids = encode(user_text)
user_generated = greedy_continue(user_prompt_ids, max_new_tokens=10)

print("Twój prompt:", decode(user_prompt_ids[0]))
print("Dopisany tekst:", decode(user_generated[0, user_prompt_ids.shape[1]:]))
print("Cały tekst:", decode(user_generated[0]))

Porównaj wyniki dla promptu treningowego Ala ma kota i psa. oraz dla kilku nowych połączeń znanych słów. Model uczony tylko na jednej sekwencji może powtarzać fragmenty, dopisywać coś niepasującego do początku lub zakończyć wypowiedź w dziwnym miejscu. Sprawdzenie innego promptu nie zmienia jego parametrów: wywołujemy jedynie generację.

Jeżeli model zapamiętał jedną sekwencję, ale zawodzi po jej zmianie, nie jest to paradoks. Zminimalizował stratę na danych treningowych bez nauczenia się reguły, która działałaby również dla nowych przykładów.

KV cache nie jest pamięcią rozmowy

Przy generowaniu tokenu numer t klucze i wartości wcześniejszych pozycji nie zmieniają się. Bez cache model ponownie obliczałby ich projekcje w każdym kroku. KV cache przechowuje wcześniejsze K i V dla każdej warstwy, dzięki czemu dla nowego tokenu trzeba dopisać tylko nową parę.

Bez cache	Z KV cache
ponowne projekcje wcześniejszych pozycji	wcześniejsze K i V są przechowane
więcej powtarzanych obliczeń	więcej zajętej pamięci urządzenia
łatwiejszy kod demonstracyjny	szybsza typowa generacja autoregresyjna

Cache nie usuwa faktu, że nowe zapytanie musi porównać się z wcześniejszymi kluczami. Nie zmienia też okna kontekstu w trwałą wiedzę modelu. To optymalizacja wykonania, a nie mechanizm uczenia.

def projections_without_cache(generated_tokens: int) -> int:
    return sum(range(1, generated_tokens + 1))

def new_token_projections_with_cache(generated_tokens: int) -> int:
    return generated_tokens

for n in [8, 32, 128]:
    print(n, projections_without_cache(n), new_token_projections_with_cache(n))

Ten rachunek dotyczy wyłącznie liczby ponownych projekcji kluczy i wartości w uproszczonym przykładzie. Nie jest pomiarem czasu ani pełnego kosztu obliczeń. Nasza funkcja greedy_continue nie implementuje KV cache: przy każdym nowym tokenie ponownie przekazuje do modelu cały dotychczasowy tekst. Nawet po dodaniu cache nowe zapytanie nadal trzeba porównać z kluczami dostępnych wcześniejszych pozycji. Przy zwykłej pełnej uwadze i stale wydłużanym kontekście liczba takich porównań dla kolejnych nowych tokenów rośnie jak 1 + 2 + … + n, a więc KV cache nie czyni kosztu całej generacji liniowym względem n.

Praca z AI: możesz poprosić o diagram przepływu danych przez blok lub kod wypisujący kształty tensorów. Sprawdź otrzymany wynik, odpowiadając na trzy pytania: gdzie nakładamy maskę przyczynową, które tensory dodawane w połączeniach resztkowych muszą mieć taki sam kształt i co dokładnie oszczędza KV cache? Zwróć uwagę, czy opis cache nie myli ograniczenia powtarzanych obliczeń z liniowym kosztem całej generacji.

---