Parallax 保留 softmax,加入學得的協方差校正,1.7B 上 +1pt

Northwestern、Tilde Research 和 University of Washington 的研究者發布了 Parallax(arxiv 2605.29157,程式碼在 github.com/Yifei-Zuo/Parallax),一種注意力機制,它在架構上做的是大多數線性注意力工作的反面。Parallax 不是為了獲得亞二次複雜度而替換 softmax,而是精確保留 softmax 並在其上加入一個學得的協方差校正分支。動機是品質而非速度:底層理論(Local Linear Attention,從 test-time regression 框架中的非參數統計衍生)將 softmax 的「局部常數估計」升級為「局部線性估計」,對聯想記憶有可證明更好的偏差-方差權衡。在 1.7B 規模配合 Muon 優化器,Parallax 在下游精度上擊敗 Transformer 基線 1.02 分(62.45 vs 61.43),在 MAD-Benchmark 上以 0.716 平均精度超越 Mamba、Gated DeltaNet、MesaNet 和 Kimi DeltaAttention。

機制是加性的。輸出等於 softmax 注意力減去一個投影的協方差項:KV 協方差乘以一個學得的探針 ρᵢ = WR·xᵢ,其中 WR 是一個可學習的投影矩陣,直接從層輸入中探測 KV 協方差。關鍵的設計屬性:設置 WR = 0 精確恢復標準 softmax。這意味著你可以通過添加 WR 並微調來轉換預訓練的 Transformer 檢查點,而無需從頭開始。計算故事:Parallax 通過為兩個分支重用相同的 KV 流來加倍算術強度,在 NVIDIA H200 解碼內核上匹配或超越 FlashAttention 2 和 3(在 compute-matched 設置下 1.54x 加速,I/O-matched 下 1.14x,跨越 batch 1 到 2048 和上下文 128 到 32768)。關鍵警告:增益取決於優化器。在 Muon 下,該架構具有顯著優勢。在 AdamW 下,優勢顯著縮小或消失,這是一種不尋常的架構-優化器協同設計依賴,限制了你可以應用此 drop-in 的地方。重要:這不是線性複雜度注意力。解碼仍需要 KV 快取;結構是 softmax-attention-shaped,而不是 Mamba-shaped。

兩個值得停下來思考的線索。首先,敘事反轉本身就很有趣。多年來線性注意力的主導故事一直是「替換 softmax 以逃避二次複雜度」,Performer、Linformer、Hyena、Mamba 加上 MiniMax M3 的 MSA 都生活在這個 frame 中。Parallax 說:softmax 是更一般事物的局部常數估計,而更一般的事物是局部線性估計,理論上更好。修復不是替換 softmax,而是擴展它。這種重新框定對思考注意力設計的 builders 來說很重要:在經驗景觀之下有一層理論,表明 softmax 不是天花板,而是一個起點。其次,優化器-架構協同設計的發現是誠實且罕見的。大多數論文不發布「我們的增益在 AdamW 下蒸發了」。Tilde Research 和合作者標記這一點意味著未來 LLA 家族的工作需要追蹤它居住在哪個優化器制度中,而 Muon 剛剛獲得了一個具體的理由,與其使能的架構創新一起被研究。

週一早上,如果你從頭訓練 1 到 2B 規模的語言模型並且有 Muon 可用:Parallax 值得作為 vanilla softmax 的 drop-in baseline 來處理,鑑於 github.com/Yifei-Zuo/Parallax 上的開源程式碼和 WR=0 與預訓練檢查點轉換的向後相容性。H200 解碼內核匹配 FlashAttention 2 和 3,所以推理端成本不是障礙。如果你由於相容性原因卡在 AdamW 上,增益可能不會實現;在投入之前在你自己的訓練執行上驗證。如果你從原語開始建構並關心注意力設計作為第一原則,LLA 理論(test-time regression,局部線性估計)是要讀的論文;Parallax 實現是一種參數化,不是唯一可能的。如果你做前沿規模工作,1.7B 的數字令人鼓舞,但架構-優化器協同設計在你的規模上需要驗證,然後才能押注你的基礎設施。