Hugging Face Infini-Attention 再現分析

Infini-Attention 機制

Infini-Attention 旨在透過將標準自注意力的二次記憶增長替換為固定大小的壓縮記憶緩衝區，以實現理論上無限的上下文長度。其流程遵循以下技術步驟：

Segmentation（分段）：將輸入序列切分為固定大小的段落。
Local Attention（局部注意力）：在當前段落內計算標準的因果點積注意力。
Memory Retrieval（記憶檢索）：模型使用當前段落的查詢向量 ($Q$) 從壓縮記憶矩陣 ($M_{s-1}$) 中檢索長期上下文，並通過非線性激活函數（ELU + 1）處理。
Integration（整合）：可學習的標量參數 ($\beta$) 作為門控機制，使用 sigmoid 函數在檢索到的長期記憶 ($A_{\text{mem}}$) 與局部點積注意力 ($A_{\text{dot}}$) 之間取得平衡。
Memory Update（記憶更新）：透過將當前段落的鍵值狀態加入現有緩衝區，更新壓縮記憶。
State Transfer（狀態轉移）：丟棄前一段的注意力狀態，僅將更新後的壓縮記憶傳遞至下一段。

在使用 200M Llama 模型與 Llama 3 8B 進行再現的過程中，Hugging Face 發現了多項關鍵的收斂障礙：

最初的實驗顯示約 95% 的平衡權重集中在 0.5 附近，表明門控機制未能收斂。分析發現，標準的 Llama 3 8B 超參數（學習率為 $3.0 \times 10^{-4}$）不足以讓平衡因子從初始化值顯著移動。

為了解決此問題，團隊為門控函數設定了較高的獨立學習率（0.01），同時保持全局學習率為 $3.0 \times 10^{-4}$。這使得平衡因子能達到更理想的範圍，儘管在 200M 模型中於 200 億個 token 後最初出現 NaN 損失。

進一步檢查發現，權重衰減會促使平衡因子的 L2 範數變小，導致 sigmoid 值集中在 0.5 附近。為了抵消這一影響，Hugging Face 採取了以下調整：

這些調整使得全局權重分佈於 0 到 1 的完整範圍，且有 10% 的 heads 其權重達到 0.9 至 1.0 之間。

評估使用「通行碼檢索任務」進行，其中特定的「針」(例如通行碼) 隱藏在上下文中不同位置的無關文字裡。

Small-scale（小規模，200M 模型）：早期訊號顯示模型能生成與較早段落相關的內容，但在檢索到精確的針時仍有困難。
Llama 3 8B：當針放置於較早段落時，模型最初未通過針的評估。
Optimized Setup（最佳化設定）：在調整學習率並移除權重衰減後，模型在延續較早段落的精確內容方面有所提升，並在特定提示下通過了一些針通行碼測試，但仍表現不穩定。

儘管門控收斂已有改善，Hugging Face 結論認為 Infini-Attention 尚未達到可在生產環境使用的可靠程度。團隊仍認為 Ring Attention、YaRN 與 RoPE scaling 是延伸預訓練模型至更長上下文的最佳選擇。

Gating Sensitivity（門控敏感度）：門控函數的正確收斂需要特定的學習率調整，並移除權重衰減以避免權重集中於 0.5 附近。
Loss vs. Utility（損失與效用）：即使模型條件不佳或注意力輸出存在維度錯誤，梯度下降仍可能降低訓練損失，這表示僅靠損失無法充分衡量一致性；必須持續進行評估。
Compression Trade-off（壓縮取捨）：隨著記憶壓縮次數的增加，性能必然下降。