<mark id="oztn4"></mark>
0
作者丨鄭佳美
編輯丨岑 峰
這幾天，Kimi 研究員、RoPE 提出者蘇劍林發布了一篇名為《簡單談談 K3 的 MoE 和 Attention》的文章，集中討論了 Kimi K3 在專家架構、訓練穩定性、負載均衡和注意力設計上的幾項關鍵取舍。
蘇神在文章中提到 K3 擁有 2.8 萬億總參數，并配置了 896 個路由專家。但模型不會讓每個 Token 調用全部專家，而是只從中激活 16 個。在注意力結構上，K3 也沒有沿用單一方案，而是將 KDA 與 Gated MLA 交替排列。
這些配置共同指向了 K3 的基本設計思路：模型可以繼續擴大容量，但不能讓計算成本隨著總參數量和上下文長度同步增長。更多專家意味著模型擁有更細致的能力分工，混合注意力則讓模型能夠以更低成本處理長上下文。
然而，規模擴張也會帶來新的問題。 當專家數量增加時，Router 不僅要準確選擇專家，還要避免少數專家持續過載；當 Token 被發送到不同 GPU 上的專家時，跨設備通信也會迅速增加。
與此同時，更長的上下文會推高 KV Cache 和注意力計算成本，而 BF16、FP8 等低精度訓練雖然能夠節省資源，卻更容易受到異常激活的影響。
因此，K3 的關鍵并不是簡單加入更多專家或更換注意力模塊，而是為規模擴張建立一套配套的控制機制。
LatentMoE 負責降低專家計算和通信成本，RMSNorm 與 SiTU-GLU 用來穩定專家分支的數值，Quantile Balancing 負責協調近千個專家之間的負載，KDA 與 NoPE MLA 則重新分配長上下文中的記憶和檢索任務。
換句話說，K3 真正要解決的問題是：模型可以擁有更大的參數和信息空間，但每一步計算都必須控制實際調用的部分。
01
傳統 MoE 會在模型內部設置大量專家，再由 Router 根據 Token 內容選擇其中少數幾個參與計算。這樣，模型可以擁有很大的總參數量，同時將單次計算控制在較小范圍內。
但在真實訓練系統中，專家通常分布在不同 GPU 上。Router 完成選擇以后，系統還需要把 Token 的隱藏狀態發送到相應設備。隱藏維度越寬、每個 Token 激活的專家越多，需要傳輸的數據就越多。到了大規模 MoE 階段，通信往往比矩陣計算更早成為瓶頸。
LatentMoE 改變了 Token 進入專家的方式。K3 不會直接把完整隱藏狀態發送給路由專家，而是先將其壓縮到更窄的潛在空間。K3 的主隱藏維度為 7168，路由專家則在 3584 維空間中計算，完成后再將結果恢復到完整隱藏維度。
降維帶來的收益，不只是減少單個專家的計算量。專家需要讀取的權重、處理的激活以及跨設備傳輸的數據都會同步下降。K3 將這部分預算用于擴大專家池：原本可以采用從 448 個專家中選擇 8 個的方案，引入 LatentMoE 后，最終擴展為從 896 個路由專家中選擇 16 個。
兩種配置的激活比例相同，但后一種方案擁有更多可組合的專家。模型可以讓多個較窄的專家共同處理一個 Token，而不是依賴少數寬專家完成所有變換。專家分工由此變得更加細致。
潛在空間同時也是一道信息瓶頸。路由專家接收到的是壓縮表示，如果維度過窄，部分信息可能在進入專家前已經損失。
K3 因此還保留了 2 個始終參與計算的共享專家，負責處理語言結構、基礎語義和常見推理模式。路由專家則集中學習更加細分的能力，避免數百個專家重復承擔相同的通用計算。雷峰網
LatentMoE 的意義，因而不只是把專家做小，而是重新組織通用能力、專業能力與通信成本之間的關系。
02
LatentMoE 增加了降維、專家計算、結果聚合和重新升維等步驟，計算路徑比普通 MoE 更長。
路徑中的數值波動也更容易被后續矩陣放大。K3 因此加入 RMSNorm、SiTU-GLU 和 Quantile Balancing，將其改造成 Stable LatentMoE。
RMSNorm 被放在專家結果聚合之后、升維之前。由于不同 Token 會進入不同的專家組合，Router 分配的權重也不相同，聚合結果的尺度可能存在較大差異。如果模型直接將結果升維并送回主干，這些波動就可能繼續擴散。
RMSNorm 會先校準專家分支的整體尺度，再由升維矩陣恢復完整表示。它相當于在路由分支與主干網絡之間設置了一套統一的數值接口。
不過，整體尺度穩定，并不意味著局部沒有異常。SwiGLU 會生成兩個分支，再將它們逐位置相乘。如果兩個分支在同一位置同時產生較大數值，輸出就會被乘法迅速放大。
這類離群值對 BF16、FP8 等低精度訓練尤其危險。少量極端數值會占用較大的動態范圍，使大量正常數值只能被壓縮到更窄的精度區間。雷峰網(公眾號：雷峰網)
SiTU-GLU 通過 Soft Cap 限制這種增長。激活較小時，它盡量保留原有函數的計算行為；數值進入危險區域后，增長速度逐漸降低并趨于飽和。與直接 Clamp 相比，這種處理更加平滑，也不會把所有超過閾值的數值簡單壓成同一個結果。
RMSNorm 和 SiTU-GLU 解決的是數值問題，Quantile Balancing 處理的則是專家負載。
MoE Router 容易形成正反饋。某個專家早期獲得更多 Token，就會得到更多梯度；能力提升后，它又更容易被繼續選擇。長此以往，少數專家可能持續過載，其他專家則缺少訓練機會。
K2 已經采用無輔助損失的負載均衡方法。系統通過調整專家的選擇偏置控制流量，而不是額外加入均衡損失干擾 Router 的語義學習。不過，K2 的偏置更新類似 SignSGD，只會按照固定步長提高或降低專家被選中的概率。
當專家數量增加到 896 個后，固定步長很難兼顧調整速度和穩定性。Quantile Balancing 不再逐步試探，而是直接觀察 Router 分數與 Top-K 門檻之間的分布，估計每個專家的選擇門檻應該移動到什么位置，才能接收到目標數量的 Token。
這種變化讓負載均衡從經驗性的反饋調節，轉向更直接的分布求解。RMSNorm、SiTU-GLU 與 QB 分別控制整體尺度、局部極值和專家流量，共同構成了 Stable LatentMoE 的穩定機制。
03
K3 的注意力結構由 KDA 和 Gated MLA 共同組成，大致每經過 3 層 KDA，就插入 1 層 MLA。兩者的區別不僅在于計算成本，也在于它們保存歷史信息的方式。
MLA 保留了對完整歷史的全局訪問能力。雖然它會將 KV Cache 壓縮到潛在空間，但當前 Token 仍然可以根據 Query 回查歷史中的具體內容。它承擔的是全局檢索任務。
KDA 則不會保存所有歷史 Token 的完整表示，而是把過去的信息持續寫入固定大小的狀態，并通過更新、遺忘和覆蓋維持這份狀態。它能夠以較低成本處理長序列，但固定容量意味著部分歷史細節會被壓縮。
K3 沒有要求其中一種機制獨立承擔所有任務。KDA 負責高頻地維持連續狀態，MLA 則周期性訪問完整歷史，補充固定狀態可能遺漏的重要信息。
這種分工也解釋了 K3 為什么能夠在 MLA 中移除 RoPE。
純 MLA 模型需要顯式建模 Token 之間的位置關系，因此 K2 仍然依賴 RoPE。K3 中的 KDA 會按照 Token 順序遞歸更新狀態。較早的信息需要經歷更多次傳播、衰減和覆蓋，較近的信息則通過更短的路徑影響當前位置。順序與距離已經部分包含在狀態演化過程中。
因此，MLA 不再需要獨立承擔全部位置建模任務，可以將更多能力用于全局內容匹配。所謂“廣義 RoPE”，并不是說 KDA 與 RoPE 完全等價，而是強調位置信息也可以由具有順序性的狀態更新機制表達。
K3 還在 MLA 輸出后增加了 Gate。MLA 負責從歷史中找到相關內容，Gate 再根據當前輸入判斷哪些通道值得寫回主干。這樣，模型不僅能夠控制檢索對象，也能夠控制檢索結果的使用強度。
KDA、MLA 與 Gate 因此形成了明確分工：KDA 維持連續狀態，MLA 執行全局回查，Gate 篩選回查結果。
04
K3 的 NoPE MLA 仍然保留了原本用于 RoPE 的額外 64 維分支。既然模型已經不再對這部分施加 RoPE，從理論形式看，這段結構似乎可以刪除。
但刪除分支會改變 Query 和 Key 的張量形狀，并可能影響 KV Cache 布局、Attention Kernel、通信邏輯和推理框架。對于已經建立完整訓練與部署鏈路的 2.8 萬億參數模型，底層形狀變化意味著大量組件需要重新開發和驗證。
K3 保留原有結構，反映的是一種最小改動原則。它未必是數學上最簡潔的形式，卻可以繼續復用成熟的 MLA 基礎設施，降低訓練和部署風險。
Per-Head Muon 處理的則是優化器與注意力結構之間的關系。
K3 繼續使用 Muon 優化器，但會按照不同 Attention Head 分開處理相關參數。多個 Head 往往學習不同的信息模式，如果優化器將它們作為一個整體統一歸一化，梯度較大的 Head 可能影響其他 Head 的更新尺度。
Per-Head Muon 不一定直接帶來顯著的指標提升，但它讓優化方式與模型內部的結構邊界保持一致，減少了不同 Head 之間不必要的耦合。
從這些細節可以看到，K3 的最終架構并不是只由理論效果決定。通信成本、低精度數值、Kernel 復用、框架兼容和工程風險都會參與設計取舍。
K3 仍然存在尚未完全回答的問題。低維潛在空間可能損失信息，KDA 的固定狀態仍會遺忘細節，更多專家也不必然形成同等數量的清晰能力。模型的效率還依賴專家通信、低精度訓練和專用 Kernel，單獨復制架構未必能夠獲得相同收益。
但 K3 展示了一條較為明確的路線：當模型進入萬億參數和百萬上下文階段，Scaling 的重點已經不只是擴大容量，而是建立更有效的參數、記憶與計算調度機制。
參考鏈接：https://kexue.fm/archives/11848
上車，帶你看遍全球 AI 頂會精華
可獨家暢覽：
專家演講PPT
大會報告全文
熱門論文解讀
學術新星訪談
掃描上方二維碼
或點擊「閱讀原文」關注專區。
雷峰網原創文章，未經授權禁止轉載。詳情見轉載須知。
