Trzy Macierze Na Głowicę, Lub Jedna Duża Macierz

Widok Per-Head

Każda głowica potrzebuje własnej projekcji zapytania, projekcji klucza i projekcji wartości. Dla głowicy h:

Q_h = X · W_Q^h    gdzie W_Q^h ma kształt [d_model, head_dim]
K_h = X · W_K^h    gdzie W_K^h ma kształt [d_model, head_dim]
V_h = X · W_V^h    gdzie W_V^h ma kształt [d_model, head_dim]

X ma kształt wejściowy [batch, seq_len, d_model]. Po projekcji Q_h, K_h, V_h mają każdy kształt [batch, seq_len, head_dim].

Zintegrowany Widok

Macierze dla każdej głowy znajdują się obok siebie w pamięci. Pojedyncza zintegrowana macierz W_Q o kształcie [d_model, d_model] produkuje wszystkie głowy naraz:

Q_fused = X · W_Q              # [batch, seq_len, d_model]
Q_per_head = reshape(Q_fused)  # [batch, n_head, seq_len, head_dim]

Złączone mnożenie macierzowe wymaga jednego wywołania BLAS zamiast 12. Rdzenie tensorowe CUDA osiągają szczytową przepustowość na mnożeniach macierzowych tej wielkości; mnożenia macierzowe na głowę niedostatecznie wykorzystywałyby sprzęt.

Liczba parametrów

Trzy scalone macierze W_Q, W_K, W_V, każda d_model × d_model. Plus macierz projekcji wyjściowej W_O, również d_model × d_model. Dla ANDREA-120M:

parametrów na uwagę w warstwie = 4 × 768² = 2,359,296 ≈ 2.36M
parametrów w 12 warstwach      = 12 × 2.36M ≈ 28.3M

Około jedna czwarta wszystkich parametrów ANDREA-120M znajduje się w projekcjach uwagi. Pozostałe trzy czwarte to podwarstwa MLP i osadzenia.

Nazewnictwo Projekcji

Dwanaście Wektorów Staje się Jednym

Po obliczeniach każdej głowy

Każda głowa produkuje tensor wyjściowy o kształcie [batch, seq_len, head_dim]. Dwanaście głów produkuje dwanaście takich tensorów. Konkatenacja wzdłuż wymiaru cechowego składa je z powrotem razem:

concat_output = concat(head_1, head_2, ..., head_12)
shape         = [batch, seq_len, n_head × head_dim]
= [batch, seq_len, 768]    # dla ANDREA-120M

Concat odwraca split. Całkowity wymiar cech wraca do d_model. Brak utraty informacji w wymiarach; różnica tkwi w tym, co zawiera każdy chunk: chunk głowy 1 odzwierciedla wyuczony wzorzec uwagi głowy 1.

Projekcja wyjściowa W_O

Samodzielna konkatenacja pozostawia głowy izolowane: wyjście głowy 4 sąsiaduje z wyjściem głowy 7, żadna nie jest świadoma drugiej. Projekcja wyjściowa W_O o kształcie [d_model, d_model] je miesza:

attention_output = concat_output · W_O
shape            = [batch, seq_len, d_model]

Po W_O każdy wymiar wyjściowy zawiera wyuczoną kombinację liniową wszystkich dwunastu głów. Informacje płyną swobodnie między głowami poprzez tę pojedynczą mnożenie macierzowe.

Dlaczego głowy się specjalizują

W architekturze nic nie zmusza głowy 4 do uczenia się czasu czasowników ani głowy 9 do uczenia się parzystej interpunkcji. Specjalizacja wyłania się z ciśnienia gradientu: podczas treningu głowy, które przyczyniają się redundantnie, otrzymują mniejsze gradienty niż głowy, które przyczyniają się unikalnie. Przez tysiące kroków każda głowa osadza się w niszy, która najbardziej efektywnie redukuje całkowitą stratę.

Empirycznie, wytrenowane transformery pokazują głowy, które zajmują się: wzorcami pozycyjnymi (głowa patrzy na poprzedni token), wzorcami syntaktycznymi (głowa patrzy na pasujący nawias zamykający), wzorcami semantycznymi (głowa patrzy na najbardziej最近 nazwany podmiot). Żadne etykiety nie trenują tej specjalizacji. Sam sygnał treningowy, propagowany przez W_O, sortuje głowy.

Dlaczego Dwanaście Głów, a Nie Jedna Szersza Głowa

Jak CUDA przechowuje głowy

Pojedynczy tensor, zmieniony kształtem

Silnik treningowy ANDREA microgpt_cuda.cu nie alokuje dwunastu oddzielnych buforów dla dwunastu głów. Alokuje jeden scalony tensor i traktuje wymiar głowy jako wzorzec stride:

// po Q = X · W_Q (jedno mnożenie macierzowe, scalone przez głowy)
// Q ma kształt [batch, seq_len, d_model]
// przekształć w kształt    [batch, seq_len, n_head, head_dim]
// transponuj do  [batch, n_head, seq_len, head_dim]
// każda głowa jest teraz ciągła w dwóch wewnętrznych wymiarach

Transpozycja przesuwa n_head przed seq_len. Dlaczego? Ponieważ następna operacja (Q_h · K_h^T) wymaga, aby wycinek seq_len × head_dim każdej głowy był ciągły w pamięci. Mnożenia macierzowe CUDA działają szybciej na ciągłych tensorach.

Jeden kernel, wiele głów

Pojedyncze jądro CUDA uwagi działa równolegle na wszystkich głowach. Każdy blok wątków obsługuje jedną parę (batch, head); wątki wewnątrz bloku współpracują nad kafelkiem seq_len × head_dim. Jądro nigdy nie wie, że przetwarza wiele głów; siatka uruchomienia obsługuje równoległość.

Konfiguracja Odzwierciedla Sprzęt

Wybór ANDREA-120M n_head=12, head_dim=64 jest zgodny z tensor cores RTX 4090, które preferują kafelki matmul w wielokrotnościach 16. head_dim=64 = 4 × 16 idealnie pasuje do kształtu kafelka. head_dim=32 (ANDREA-12M) również pasuje, ale niedostatecznie wykorzystuje kafelek. head_dim=72 nie pasowałoby i wymusiłoby powrót do kernelów zapasowych.

Końcowy Obraz

Pięć kroków. Trzy operacje macierzowe dotykają wejścia (projekcje Q, K, V). Jedna operacja macierzowa dotyka scalonych głów (W_O). Jeden kernel attention obsługuje wszystkie głowy równolegle. ANDREA-120M wykonuje wszystkie pięć kroków raz na warstwę, dwanaście warstw w głąb, przy każdym przejściu w przód.

Czytanie potoku