Parallax：參數化局部線性注意力機制，突破大型語言模型的運算瓶頸與效能極限

本研究推出 Parallax，透過參數化局部線性注意力解決傳統 LLA 的擴充難題，並藉由硬體感知演算法提升運算效率，在維持模型規模下顯著優於 FlashAttention。

Problem

大型語言模型（LLMs）的核心注意力機制長期缺乏結構性創新。雖然局部線性注意力（LLA）在理論上能提供更佳的偏差-變異數權衡，但受限於數值穩定性與運算瓶頸，過去難以應用於大規模預訓練任務中。

研發團隊開發了 Parallax 機制，捨棄 LLA 原有的複雜數值求解器，改用參數化的查詢投影器來探測鍵值（KV）共變異數。此外，提出一種硬體感知的演算法以提升算術強度，並透過 Muon 優化器實現架構與優化器的協同設計。

在 0.6B 與 1.7B 規模的預訓練實驗中，Parallax 在困惑度（Perplexity）與下游評測任務均展現優勢。其實作的解碼內核在多種批次大小與上下文長度下，效能均追平或超越 FlashAttention 2/3，達成運算與準度的 Pareto 改進。

這項研究是學術界首次展示注意力機制架構與優化器協同設計（Codesign）的顯著成效。Parallax 不僅解決了長序列處理的效率難題，更為未來 LLM 架構結合非參數統計理論開闢了實務可行的新路徑。