自注意力
3 秒看懂
自注意力是 Transformer 架構的核心,讓模型在處理一個序列(比如一句話)時,能夠同時關注序列中所有位置的資訊,並動態計算每個位置對其他位置的重要程度。它不依賴順序一步步處理,而是通過全域性“配對打分”的方式,直接建模任意兩個元素之間的關係,從而高效捕獲長距離依賴。
關鍵類比:傳統 RNN 像序列讀報紙,只能記住附近幾行;自注意力像瞬間掃視整頁,並可任意關聯相隔很遠的詞彙。
3 分鐘產業解釋
自注意力機制徹底改變了自然語言處理、計算機視覺和多模態模型。它的核心思想是:對每個輸入 token,生成 Query、Key、Value 三個向量,通過 Query 與所有 Key 的點積計算注意力權重,再用這些權重對 Value 加權求和,得到融合了全域性上下文的輸出表徵。多頭注意力進一步將多個不同的“子空間”平行計算,增強表達能力。
在產業中,自注意力是 ChatGPT、BERT、Vision Transformer (ViT)、擴散模型中的 U‑Net 注意力塊等一切明星模型的基石。它使模型引數量可達千億甚至萬億級別,支撐了當前大語言模型的湧現能力。其計算複雜度為序列長度的平方 O(L²),因此長序列場景下的稀疏注意力、FlashAttention、KV Cache等最佳化是關鍵工程方向。
15 分鐘專家深入
自注意力的形式化定義(Scaled Dot‑Product Attention):
給定輸入矩陣 X \in \mathbb{R}^{L \times d}(L 為序列長度,d 為隱藏維度),通過可學習矩陣 W^Q, W^K, W^V 線性投影得到:
Q = X W^Q, \quad K = X W^K, \quad V = X W^V
其中 Q, K, V \in \mathbb{R}^{L \times d_k} (通常 d_k = d / h,h 為頭數)。注意力輸出為:
\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V
縮放因子 \sqrt{d_k} 防止點積隨維度增大進入 softmax 的飽和區,穩定梯度。多頭注意力將 d 切成 h 個頭分別計算,再拼接後線性投影:
\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h) W^O
自注意力即 Q, K, V 來自同一輸入源(通常加上位置編碼),交叉注意力則在解碼器中以編碼器輸出為 K、V,解碼器輸入為 Q。
複雜度與最佳化:
- 原始自注意力的記憶體與計算均為
O(L^2 d),當L > 2048時成為瓶頸。 - FlashAttention(Dao et al. 2022)利用分塊(tiling)和運算元融合,在 GPU SRAM 中完成 softmax 重計算,大幅降低 HBM 訪問,實現 2~4 倍加速且記憶體線性相關於
L,已廣泛用於 GPT‑4、LLaMA 等模型的訓練與推論。 - KV Cache:自迴歸生成時,將歷史 Key、Value 快取在視訊記憶體中,避免每步重新計算,將推論複雜度從
O(L^3)降至O(L^2)。 - 稀疏注意力:如 Longformer、BigBird,通過區域性視窗+全域性 token 或隨機模式,將複雜度壓至
O(L),支援萬級 token 輸入。
幾何與訊號處理視角: 自注意力可視為一種非區域性均值濾波(non‑local means)或圖神經網路中的親和傳播。Transformer 中的注意力矩陣視覺化常呈現某些頭負責句法關係、某些頭捕獲語義全域性模式等。位置編碼(絕對/相對)彌補序列順序的缺失,其中旋轉位置編碼(RoPE)成為主流,能為注意力注入相對位置資訊且在長距離外推中表現更優。
技術原理
核心機制:基於內容的軟定址 自注意力其實是一個可微分的鍵值儲存查詢過程。Query 相當於搜尋詞,Key 相當於資料庫索引,Value 為儲存內容。輸出是各個 Value 的加權和,權重由 Query 與相應 Key 的匹配程度決定。無需固定視窗大小,模型自動決定“關注哪裡”。
注意力權重計算流(以單頭為例):
輸入 tokens: [T1, T2, ..., Tn]
每個 token 投影: q_i, k_i, v_i
對於每個 token i:
分數 s_ij = q_i · k_j^T / sqrt(d_k) (j = 1..n)
權重 a_ij = softmax(s_i1, ..., s_in)
輸出 o_i = Σ_j a_ij v_j
矩陣形式:
Q [L x d_k]
K^T [d_k x L]
QK^T = S [L x L] -> softmax -> A [L x L]
V [L x d_k]
AV = O [L x d_k]
多頭並行: 將 Q、K、V 各切分 h 份,獨立計算,最後拼接:
Q → Q1, Q2, ..., Qh
K → K1, K2, ..., Kh
V → V1, V2, ..., Vh
head_i = Attention(Q_i, K_i, V_i)
MultiHead = Concat(head_1, ..., head_h) W^O
關鍵引數:
- 隱藏維度
d(如 768, 1024, 4096) - 頭數
h(如 12, 16, 32;通常使得d_k = 64 \sim 128) - 序列長度
L - 引數量:Q、K、V 投影各為
d \times d_k \times h(或直接d \times d),輸出投影d \times d,總計約4d^2(忽略 bias)。例如 GPT‑3 中的多頭自注意力塊引數約4 \times (12288)^2 \approx 603 M,僅一層。
梯度流與初始化:
殘差連線與 LayerNorm 包圍子層,避免梯度消失。Q、K、V 矩陣常用 Xavier 初始化,輸出矩陣初始化通常乘以因子 1/\sqrt{2N} (N 為層數)以保證訓練穩定。對於大規模模型,混合精度(FP16/BF16)與注意力層的重計算(gradient checkpointing)必不可少。
技術演進史
| 時間 | 里程碑 | 要點 |
|---|---|---|
| 2014 | Bahdanau 注意力 | 在 RNN seq2seq 中引入加性注意力,動態對齊源句與目標詞 |
| 2015 | Luong 注意力 | 提出乘性注意力(點積),效率更高 |
| 2017 | Transformer 提出 (Vaswani 等) | 完全基於自注意力,拋棄迴圈結構,實現高度並行訓練;定義縮放點積注意力和多頭機制 |
| 2018 | BERT/GPT | BERT 使用 Transformer 編碼器的雙向自注意力;GPT 使用自迴歸解碼器的因果掩碼自注意力 |
| 2019 | Sparse Transformers (Child 等) | 引入稀疏注意力模式,降低複雜度 |
| 2020 | 視覺 Transformer (Dosovitskiy 等) | 將影像切割為 patch,應用純自注意力,撼動 CNN 霸權 |
| 2021 | Performer, Linformer | 通過核近似、低秩分解將注意力複雜度降至線性 |
| 2022 | FlashAttention (Dao 等) | IO‑aware 精準實現,解決記憶體瓶頸,使長序列訓練成為可能 |
| 2023 | FlashAttention-2, RingAttention | 進一步最佳化並行效率,支援百萬 token 上下文 |
| 2024 | FlashAttention-3, MoE + 長上下文模型 | Gemini 1.5 Pro (1M tokens)、GPT‑4 Turbo (128K) 等,依賴高效注意力實現 |
技術路線對比
| 機制 | 典型網路 | 複雜度(單層) | 長距離依賴 | 並行性 | 侷限性 |
|---|---|---|---|---|---|
| 自注意力 | Transformer | O(L^2 d) | 強(全域性) | 高(逐層並行) | 序列長度平方開銷 |
| 線性注意力 | Linformer, Performer | O(L d) | 近似全域性 | 高 | 表達能力弱於標準注意力 |
| 稀疏注意力 | Longformer, BigBird | O(L \log L) \sim O(L) | 較強(區域性+稀疏全域性) | 高 | 模式設計需先驗知識 |
| 迴圈網路 (RNN/LSTM) | Seq2Seq | O(L d^2) | 弱(梯度消失) | 低(時間步序列) | 難以建模超長距離關係 |
| 卷積網路 (CNN) | WaveNet, TCN | O(k L d^2) | 感受野受層數限制 | 較高 | 需堆疊多層,引數效率較低 |
| 狀態空間模型 (SSM) | Mamba (S4, S6) | O(L d) | 強(選擇性機制) | 訓練高效,推論可比肩 | 上下文特化能力尚在實證比較 |
量化比較(以典型中等設定 L=2048, d=1024, h=16 為例,均為 [估算]):
- 標準自注意力:約
2 \times L^2 \times d計算量 ≈8.6\times10^9FLOPS,記憶體佔用接近L^2矩陣。 - 稀疏注意力(視窗大小 w=512):約
L \times w \times d級別,降低 4~8 倍。 - FlashAttention:實際視訊記憶體訪問降低 20~30 倍,在長序列上取得與理論 FLOPs 縮減近似的 wall‑clock 加速。
上下游
上游:資料與表徵
- 輸入 token 的嵌入(Embedding)+ 位置編碼(正弦、可學習、RoPE)
- 對於非文本模態,如 ViT 的 patch embedding,或語音的梅爾頻譜特徵再投影至序列
- 訓練架構:PyTorch, JAX, TensorFlow;分散式庫:Megatron‑LM, DeepSpeed
核心:自注意力計算模組
- GPU/NPU 上的矩陣乘法 (GEMM) 和 softmax 運算元;FlashAttention 等定製核心
- 記憶體與頻寬:HBM 容量與頻寬決定最大序列長度和批次大小([未充分揭露])
- 關鍵最佳化:張量並行(列切 QKV, 行切輸出)、序列並行(RingAttention)
下游:模型架構與應用
- 僅編碼器 (BERT/RoBERTa):文本理解,分類,抽取式 QA
- 僅解碼器 (GPT, LLaMA, Gemini):自迴歸語言模型,創意寫作,程式碼生成
- 編碼器‑解碼器 (T5, BART):翻譯,摘要
- 視覺:ViT, Swin Transformer, 擴散模型中的交叉注意力(文本控制影像生成)
- 多模態:CLIP 式雙塔、BLIP‑2 的 Q‑Former(使用交叉注意力彌合模態 gap)
端側部署:
- 量化 (INT8/INT4)、蒸餾、剪枝減少注意力矩陣計算
- KV Cache 壓縮技術:Grouped Query Attention (GQA)、Multi‑Query Attention (MQA) 降低記憶體讀取
- NPU/邊緣晶片適配自定義注意力運算元
關鍵指標
- 序列長度 (Context Length):從 512(BERT)到 128K(GPT‑4 Turbo)、1M(Gemini 1.5 Pro),當前前沿百萬 token 級別。
- 頭數 (h) 和每頭維度 (d_head):常見 d_head 為 64, 128;頭數從基礎模型的 12 到大型模型的 64~128,GQA 中海量 KV 頭數縮減。
- 注意力複雜度/吞吐量:用 FLOPs/token 表徵;FlashAttention 實現下,單 A100 GPU 可處理約 64K token 訓練([估算])。
- 推論時延:首個 token 延遲(prefill 階段)與每 token 延遲(decode 階段),後者主要受 KV Cache 存取限制。GQA 相比 MHA 在 batch size 較大時可將 decode 延遲降低 30%~50% [估算]。
- 困惑度 (PPL):評估語言模型質量,自注意力設計(ALiBi vs. RoPE)直接影響長序列外推時的 PPL。
- 視訊記憶體佔用:總引數量決定模型權重;KV Cache 大小 ≈
2 \times \text{層數} \times \text{頭數} \times d_{\text{head}} \times L(per token),是大批次推論的主要瓶頸。
供需與市場資料
- 需求端:全球大語言模型訓練和推論需求呈指數增長,推動對能效比更高的注意力加速硬體需求。據 [行業報告估算],2024 年大型 AI 叢集中 60%~80% 的計算時間消耗在 Transformer 的注意力和前饋網路層。
- 供給端:NVIDIA H100/B200 的 Transformer Engine 提供對 FP8/BF16 混合精度的原生支援;AMD MI300X 等也引入類似最佳化。雲端廠商(AWS, GCP, Azure)提供大規模 GPU 租賃。專用 ASIC/FPGA 如 SambaNova, Cerebras, Groq 針對自注意力設計了資料流架構以破除記憶體牆。
- 價格:雲端端 GPU 算力價格約 $2~$4/GPU·時(H100 市場均價 [估算]),訓練千億引數模型需數千 GPU·月。推論 API 價格(如 GPT‑4 Turbo)已降至 $0.01/1K tokens 級別。
- 趨勢:長上下文需求持續推動 FlashAttention 類演算法落地,混合頭(GQA)成為標配。線性注意力與狀態空間模型(Mamba)在特定場景分流。端側裝置對輕量自注意力的需求催生 MobileViT、EdgeNeXt 等。
代表公司與資本對映
| 層級 | 公司/組織 | 與自注意力的關聯 | 資本屬性 |
|---|---|---|---|
| 核心模型 | OpenAI (GPT 系列), Google (Gemini, T5), Meta (LLaMA), Anthropic (Claude) | 架構中深度使用自注意力及其變體,引領技術趨勢 | 非上市(OpenAI, Anthropic 估值得天價),上市公司 (Google, Meta) |
| 架構與硬體 | NVIDIA (CUDA, TensorRT, H100/B200), AMD, Intel | 提供訓練/推論加速,CUDA 生態壟斷自注意力計算效率最佳化 | 上市公司,市值萬億美元級 |
| 演算法最佳化 | Together AI, Mistral, 清華/北大等學術機構 | 提出 FlashAttention‑3, RingAttention, 稀疏注意力等關鍵最佳化 | 初創公司融資數億美元 |
| 應用落地 | Microsoft (Copilot), Salesforce (Einstein GPT), Duolingo | 將自注意力驅動的生成式 AI 嵌入產品 | 上市公司 |
| 端側推論 | Apple, Qualcomm, MediaTek | 在晶片中整合 Neural Engine/DSP 最佳化 Transformer 運算元 | 上市公司 |
投資邏輯
- 算力底層受益:自注意力仍是未來 2~3 年大型模型主導範式,訓練和推論的 FLOPs 需求隨模型規模和上下文長度雙增,直接利好 GPU/TPU 晶片、HBM 儲存、光模組互聯。
- 高效注意力創新企業的機遇:能大幅降低自注意力計算成本的技術公司(如稀疏注意力 IP、定製 AI 晶片)可能顛覆現有生態,但風險在於需與 CUDA 生態相容。
- 長上下文催生新應用:自注意力實現百萬 token 上下文,開啟程式碼庫理解、全基因組分析、長影片理解等新市場,軟體層價值大。
- 風險:若狀態空間模型等新生架構在多項基準上全面超越而不再依賴自注意力,當前龐大的 Transformer 訓練投入可能折價。此外,量化、蒸餾等技術降低了推論門檻,可能使硬體專用的注意力加速優勢被稀釋。
- 配置策略:短期(1‑2 年)重標配 GPU/光模組;中長期關注能效比更高的替代算力架構與上下文技術突破,可在初創二級市場或同業對比中分散版面配置。
常見誤讀糾偏
-
誤讀1:“自注意力就是全域性注意力,所以一定能解決所有長距離依賴問題。” 糾偏:自注意力從機制上提供了全域性視野,但在訓練樣本有限或最佳化不充分時,注意力分佈可能崩塌到區域性或無關 token,尤其長序列的中間資訊可能被稀釋。還需要位置編碼、初始化策略、正則化等配合,才能使模型真正學會利用長距離資訊。
-
誤讀2:“多頭注意力中每個頭一定會學到語法、語義等可解釋的角色。” 糾偏:雖然某些頭確實可能專注特定模式(句法依存、共指等),但這種解耦並非必然。大量頭呈現冗餘或混合模式。頭數過多反而可能導致過擬合或訓練變慢,GQA 便是通過壓縮 KV 頭數實現效率與效果平衡。
學習路徑
- 基礎理論:閱讀 “Attention Is All You Need” 論文原文,重點掌握 Scaled Dot‑Product Attention 公式、多頭機制、自注意力和交叉注意力的區別。
- 動手實現:用 PyTorch/NumPy 從零實現單頭、多頭自注意力,在小資料集上除錯梯度流。對比不使用注意力、使用加性注意力等實驗。
- 進階最佳化:深入理解 FlashAttention 論文和實現原理(分塊,online softmax)。學習 KV Cache 的儲存管理與 GQA/MQA 的原理。
- 變體探索:瞭解 Sparse Attention (Longformer, BigBird)、線性注意力 (Performer)、核注意力。
- 跨模態遷移:閱讀 ViT 論文,理解自注意力從 NLP 到 CV 的適配(patch + 位置編碼)。探索擴散模型中的交叉注意力。
- 系統與落地:學習 Megatron‑LM 的張量並行、序列並行切分策略;嘗試用 vLLM/TGI 部署模型,分析 attention 的視訊記憶體佔用與排程。
一句話總結
自注意力是賦予序列模型全域性互動能力的基礎算符,其優雅的“查詢‑鍵‑值”體系與注意力重加權的設計,構成了當代大型模型理解上下文、跨越長距關聯的核心引擎,並通過軟硬體協同最佳化持續拓寬智慧邊界。
延伸閱讀與來源
- Vaswani et al. (2017) “Attention Is All You Need”, NeurIPS. https://arxiv.org/abs/1706.03762
- Dao et al. (2022) “FlashAttention: Fast and Memory‑Efficient Exact Attention with IO‑Awareness”, NeurIPS. https://arxiv.org/abs/2205.14135
- Dao (2023) “FlashAttention‑2: Faster Attention with Better Parallelism and Work Partitioning”. https://arxiv.org/abs/2307.08691
- Brown et al. (2020) “Language Models are Few‑Shot Learners” (GPT‑3), NeurIPS. https://arxiv.org/abs/2005.14165
- Dosovitskiy et al. (2020) “An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale”, ICLR. https://arxiv.org/abs/2010.11929
- Chowdhery et al. (2022) “PaLM: Scaling Language Modeling with Pathways”. https://arxiv.org/abs/2204.02311 (涉及 GQA/多查詢注意力)
- 各大科技公司技術部落格:Google AI Blog, Meta AI Blog, OpenAI Research, 輝達開發者部落格 (FlashAttention 與 Transformer Engine)
- 行業估算資料來源:NVIDIA 財務簡報、IDC AI 基礎設施預測、Semianalysis 等第三方分析平台(已標註 [估算])