破解 Transformer：深入探討現代大型語言模型的數學原理

大型語言模型（LLM）的內部運作常常給人「黑盒」的感覺，因為它被高層抽象和複雜的矩陣運算所遮蔽。對於開發者與研究者而言，了解資料的精確流向——從 token ID 到機率分布——對於優化效能以及掌握像多頭潛在注意力（MLA）或專家混合（MoE）等架構創新實際運作方式至關重要。

Transformer Math Explorer 提供了對這些過程的罕見、細緻觀察。它刻意拋棄高層的矩陣乘法，改以明確的加總與索引呈現，揭示了驅動現代 AI 產生每個 token 的基礎數學。

高層流程：從 Token 到機率

從本質上看，transformer 模型（例如 GPT-2）可視為一個函式，輸入為 token ID，輸出為下一個 token 的機率分布。此過程遵循特定的結構管線：

  1. Embeddings：Token ID 會被轉換成連續向量。
  2. Transformer Blocks：這些嵌入向量會流經 $L$ 個相同的區塊，每個區塊都會在上下文中細化 token 的表示。
  3. Output Head：最終處理過的向量會映射回詞彙表大小，產生下一個 token 的機率分布。

Transformer 區塊的結構

每個 transformer 區塊的設計旨在維持訊號穩定性，同時提升模型學習複雜模式的容量。標準區塊由兩個主要子層組成：

  • Causal Self-Attention：讓模型衡量序列中不同 token 的重要性。
  • Multi-Layer Perceptron (MLP)：處理注意力機制提取的資訊。

兩個子層皆被 residual connections 包裹，以防止訓練時的梯度消失問題，並搭配 normalization 層，使激活值維持在穩定範圍內。

Causal Self-Attention 的機制

Self-attention 是 transformer 的核心引擎。對於單一頭而言，計算遵循嚴謹的流程：將輸入 $x$ 投影為 Query（$Q$）、Key（$K$）與 Value（$V$）向量，計算它們之間的相容性分數，遮蔽未來位置以確保因果性，套用 softmax 函式，最後計算值的加權總和。

Scaled Dot Product

位置 $t$ 的 query 與位置 $s$ 的 key 之間的「相容性分數」使用縮放點積計算。為防止隨著頭寬 ($d_h$) 增大而導致分數方差爆炸，結果會除以 $\sqrt{d_{h}}$：

$$z_{t , s} = \frac{1}{\sqrt{d_{h}}} \sum_{0 \leq a < d_{h}} q_{t , a} k_{s , a}$$

KV 快取

在自回歸生成過程中，模型一次產生一個 token。由於注意力機制具因果性，位置 $t$ 的 token 的 Key（$k_t$）與 Value（$v_t$）向量僅依賴於位置 $\leq t$ 的輸入。

這表示隨著後續 token 被加入序列時，這些向量不會改變。為避免重複計算，模型會使用 KV cache，將這些向量儲存起來並在每個解碼步驟中重複使用。此舉大幅降低產生長序列的計算負擔。

超越基礎：現代架構變體

雖然基本數學保持一致，現代模型已引入多項最佳化以提升效率與推理能力。Transformer Math Explorer 允許使用者在這些進階配置之間切換：

  • MLA (Multi-Head Latent Attention)：一種減少 KV 快取記憶體占用的最佳化。
  • MoE (Mixture of Experts)：將密集的 MLP 層替換為稀疏、專門的「專家」網路，以在不成比例提升計算成本的情況下增加模型容量。
  • RoPE (Rotary Positional Embeddings)：透過在複數空間中旋轉向量來編碼位置資訊的高階方法，提升模型處理更長上下文的能力。
  • MTP (Multi-Token Prediction)：一種架構轉變，使模型同時預測多個未來 token，而非僅預測單一 token。

透過將這些概念拆解為基礎的加總與索引，我們得以超越「線性代數」的抽象，直接觀察使機器能預測句子中下一個詞的實際算術運算。

Sources