生成式AI時代下的機器學習 — Transformer 的競爭者們
生成式AI時代下的機器學習 — Transformer 的競爭者們¶
本講次探討了 Self-Attention (SA) 的局限性及其潛在的替代架構,特別是 Mamba 及其相關的 RNN 變形。
- Self-Attention (SA) 的優勢與挑戰
• SA 的架構基礎:Self-Attention 透過將輸入向量轉換為 Query (Q)、Key (K) 和 Value (V),並計算 Q 與 K 的內積來決定權重,最終對 V 進行加權求和得到輸出。
• SA 的核心優勢 — 訓練加速:SA/Transformer 的主要優勢在於訓練時的高度平行化 (Parallelization)。
◦ 當給定一個完整的輸入序列時,Transformer 可以一次性平行計算出序列中每個時間點應該輸出的下一個 token。
◦ 所有運算主要基於矩陣乘法,對 GPU 資源的利用極為高效。
• SA 的核心挑戰 — 長序列處理:
◦ 推理時的計算量與記憶體需求:SA 的計算複雜度隨著序列長度 T 呈平方級增長 (O(T2)),推理時需要記住所有先前的 token,記憶體需求極高。
◦ 記憶體有限性:雖然 SA 看起來可以存取無限長的序列,但由於 Query 和 Key 的維度 D 是固定的,當序列長度 T 遠大於 D 時,記憶體會開始衝突或混亂,儲存資訊的能力實則有限。
- Recurrent Neural Network (RNN) 的重現
• 傳統 RNN 流:使用 Hidden State (H) 來存儲當前所有輸入的混合資訊。新 Hidden State (Ht) 由前一狀態 (Ht−1) 和當前輸入 (Xt) 共同決定。
◦ 推理優勢:RNN 在推理時,每一步的運算量和記憶體需求是固定的,非常適合生成長序列。
◦ 訓練瓶頸:傳統 RNN 由於 Ht 依賴 Ht−1,難以實現平行運算,導致訓練效率低。
• Linear Attention (LA):
◦ 透過移除廣義 RNN 架構中的「Reflection」部分(即 Fa 函式),RNN 的遞歸公式可以被完全展開。
◦ 數學證明顯示,Linear Attention 實質上就是沒有 Softmax 函數的 Self-Attention。
◦ LA 的優勢:LA 在推理時具備 RNN 的效率,但在訓練時可以像 SA 一樣展開為矩陣運算,從而實現平行化和加速訓練。
• LA 的限制與 Softmax 的作用:
◦ 記憶體靜態:原始 LA 架構中,一旦資訊被寫入 Hidden State (Memory),它將永遠不會被改變或遺忘。
◦ Softmax 的功能:在 SA 中,Softmax 使注意力權重具有相對性。當序列中出現更重要的資訊時,原本重要的資訊的權重會被相對降低,從而實現隱式的「遺忘」或「記憶調整」。
- Linear Attention 的進化:Mamba 與記憶體控制
為了克服 LA 記憶體靜態的限制,新的架構著重於將記憶體控制機制重新引入 RNN:
• Retention Network (RetNet):引入一個衰減常數 (γ<1),乘在 Ht−1 上,使舊記憶體隨著時間推移而自然淡忘。
• Gated Retention (Gated RetNet):將衰減常數 γ 參數化為 Gt,使其成為一個學習到的、與上下文相關的門控機制,讓模型可以選擇性地遺忘或保留資訊。
• Mamba:目前最知名的 Linear Attention 變體。
◦ Mamba 的設計相對複雜,但其核心是通過對 Ht−1 進行元素級(element-wise)的複雜乘法,實現對記憶體的精確控制。
◦ Mamba 的成就:Mamba 是第一個能夠在性能上超越 Transformer 架構(在多數基準測試中微幅勝出)的 Linear Attention 模型。同時,Mamba 在推理時的加速效率遠高於 Transformer。
• DeltaNet:這類架構(如 Titan)將 Hidden State H 視為訓練參數,將記憶體更新 (Reflection) 視為一次 Gradient Descent 的步驟。更新的目標是確保 Memory 在被 Key 查詢時,能夠精確地提取出其應有的 Value,實現「在測試時學習記憶」。
- 總結:隨著 Mamba 等架構的出現,Linear Attention 已能用於訓練極大規模的語言模型(如 52B 或 400B 級別),且在性能上與 Transformer 持平或略優。這使得關於「Attention 是否是唯一所需」的爭論持續延燒。
參考資訊:
第四講:Transformer 的時代要結束了嗎?介紹 Transformer 的競爭者們https://www.youtube.com/watch?v=gjsdVi90yQo&list=PLJV_el3uVTsNZEFAdQsDeOdzAaHTca2Gi&index=4
Comments
Loading comments…
Leave a Comment