4 位元 AdamW 的捨入空間之爭：為什麼要在預條件子空間裡做隨機捨入

AdamW 是當今大模型訓練的預設優化器，但它的代價常被低估。除了權重和梯度，每個參數還要保存一階矩與二階矩兩份狀態，按 32 位元儲存時，這部分常駐記憶體與權重本身相當甚至更多。把優化器狀態壓到 4 位元，是削減訓練顯存最直接的辦法之一。TorchAO 已經提供了 4 位元 AdamW 的實作，但它與 32 位元 AdamW 之間仍有可測的驗證損失差距。Hanyang Li 等五位作者在 2026 年 10 月 8 日發布於 arXiv cs.LG 的這篇論文，沒有去調碼本的形狀或區塊的大小，而是追問一個更基礎的問題：量化器在相鄰兩個重建電平之間做選擇時，究竟應該在哪個座標裡捨入。 作者稱這個座標為「捨入空間」。問題的根源在於誤差不是一次性的。量化誤差會沿著矩的遞推傳播，擾動之後每一步的自適應更新。對一階矩而言，更新量對狀態近似線性，在狀態空間裡捨入問題不大。二階矩則不同：Adam 的更新與二階矩平方根的倒數成正比，這個倒數就是預條件子。論文對緊鄰零的那個量化單元做了局部分析，結論是：狀態的平均誤差很小，並不意味著下一步預條件子的平均誤差也很小。狀態稍微偏離零，倒數函數就被急劇放大，誤差在預條件子一側被成倍拉高。作者還構造了一個一維二次問題，顯示在狀態空間捨入與在預條件子空間捨入，會產生性質上不同的優化動力學，而不只是數值上的差異。 這些觀察催生了第一個方案 ZIP-SR，全稱「保留零的預條件子空間隨機捨入」。它在二階矩碼本裡保留零這個電平，並且在預條件子空間而不是狀態空間裡計算隨機捨入的機率。隨機捨入的好處是期望無偏，把確定性的系統誤差變成零均值雜訊；而選對空間，則保證無偏性落在真正決定更新幅度的那個量上。第二個方案 ZE-EDEN 走互補路線：它使用不含零的二階矩碼本，因此量化值存在一個正的下限，避免了接近零時倒數發散。代價是這個正下限本身會扭曲預條件子，於是論文用 EDEN 式校準，對量化後的二階矩區塊重新縮放，來抵消這種畸變。兩種配置的一階矩都使用 4 位元 NormalFloat，即 NF4，並且在訓練的最後 10% 階段，對語言模型輸出頭的一階矩施加有針對性的隨機捨入。

實驗涵蓋了 GPT 與 Llama 風格的預訓練，模型規模從 1.3 億到 27 億參數。結果是，兩種方法在每一個被評估的規模上，都縮小了 TorchAO 4 位元 AdamW 相對 32 位元 AdamW 的平均驗證損失差距，最大的一次縮小達到 70%。在全參數監督微調中，兩種配方的驗證損失都低於 TorchAO，同時在下游任務上與 32 位元 AdamW 保持接近。值得注意的是，「每個規模都改善」比單點的大幅領先更有說服力，因為它說明收益來自機制，而不是某次執行的運氣。 把這些機制放回工程語境，可以看到幾層含義。第一，4 位元狀態量化的難點從來不只是位元數少，而是誤差如何被後續計算放大。以往的做法預設把狀態當作要逼近的對象，用均方誤差衡量好壞；這篇論文提醒我們，真正被消費的是由狀態派生出來的預條件子，評價與捨入都應當圍繞它展開。第二，隨機捨入與確定性捨入的取捨在這裡有了新的依據：在狀態空間裡無偏，並不等於在預條件子空間裡無偏，所以同樣是隨機捨入，換一個座標就是另一種估計器。第三，兩條路線給出了一個有趣的對照：一條保留零並改變機率的計算方式，另一條排除零並靠重新縮放修正下限。它們殊途同歸，都在處理同一個病灶，即零附近那個量化單元裡的預條件子失真。對工程團隊而言，這意味著可以按自己的碼本限制和核心實作選擇更容易落地的一條。此外，對輸出頭一階矩在訓練末段做針對性隨機捨入，也提示了一個經驗：並非所有張量對量化同樣敏感，把有限的精度預算花在最敏感的位置，往往比整體均勻降低誤差更划算。這也解釋了為何論文不滿足於單一配方，而是同時給出兩條可互相印證的路線，讓讀者能夠判斷收益究竟來自哪一環。

我們的判斷是，這篇論文最大的價值在於把優化器狀態量化的設計空間重新劃分了一次。過去的討論多集中在碼本怎麼設計、區塊怎麼切；這裡提出的是第三根軸，即捨入發生在哪個座標。這個視角很可能遷移到其他對狀態做非線性變換的優化器。同時也要保持克制：摘要沒有給出顯存節省的具體數字和訓練吞吐的變化，評估規模止於 27 億參數，更大模型與更長訓練是否維持同樣的收益，還需要獨立重現來回答。對正在為訓練顯存發愁的團隊，合理的做法是把 ZIP-SR 與 ZE-EDEN 當作值得在自己的規模上對照 TorchAO 做小範圍驗證的候選，而不是直接替換生產配置。