制全解析：從SwiGLU到Gated MLA與MoE路由)
最近在梳理 Kimi K3 的架構(gòu)資料時(shí)有一個(gè)詞被反復(fù)提及門控Gating。Gated MLA、KDA 的門控分支、MoE 中的 SiTU-GLU、SwiGLU……看起來每個(gè)模塊都在做“門控”但每個(gè)地方的門控含義又不太一樣。如果只是零散地看很容易把一個(gè)概念套到另一個(gè)模塊上最后越看越亂。這篇文章就把 LLM 里的門控機(jī)制集中梳理一遍結(jié)合 Kimi K3 的公開架構(gòu)討論拆開 Gated MLA、KDA 門控分支以及 MoE 里的 SiTU-GLU 與 SwiGLU。我會(huì)盡量不堆術(shù)語先講清楚門控本身是什么再逐層進(jìn)入現(xiàn)代大模型的門控設(shè)計(jì)并給出可運(yùn)行的 PyTorch 簡(jiǎn)化實(shí)現(xiàn)。需要提前說明的是Kimi K3 的官方技術(shù)報(bào)告尚未完整公開本文的架構(gòu)解析基于公開模型信息、社區(qū)復(fù)現(xiàn)討論和技術(shù)博客整理代碼示例是教學(xué)用簡(jiǎn)化版本不完全等價(jià)于官方實(shí)現(xiàn)。1. 門控機(jī)制LLM 架構(gòu)里的“信息閥門”1.1 從 LSTM 說起門控是什么門控不是新概念。早在 LSTM長(zhǎng)短期記憶網(wǎng)絡(luò)時(shí)代門控就是核心機(jī)制。LSTM 里有三個(gè)門遺忘門決定上一時(shí)刻的記憶要保留多少。輸入門決定當(dāng)前輸入有多少寫入記憶。輸出門決定當(dāng)前記憶有多少輸出到隱藏狀態(tài)。每個(gè)門都輸出一個(gè) 0 到 1 之間的數(shù)值用 sigmoid 激活函數(shù)實(shí)現(xiàn)。0 表示“完全關(guān)閉”1 表示“完全打開”中間值則表示“部分通過”。用一句話概括門控的本質(zhì)學(xué)習(xí)一個(gè)軟開關(guān)決定信息以什么比例通過。為什么要用“軟”開關(guān)因?yàn)橛查_關(guān)if-else不可導(dǎo)梯度無法回傳軟開關(guān)用 sigmoid 這類連續(xù)函數(shù)梯度能順暢流動(dòng)模型就能通過反向傳播自動(dòng)學(xué)習(xí)開關(guān)的開合程度。1.2 Transformer 時(shí)代的三類門控Transformer 出現(xiàn)后門控并沒有消失而是演化出了更豐富的形態(tài)。在當(dāng)今的大模型架構(gòu)里門控至少出現(xiàn)在三個(gè)層面第一層是注意力內(nèi)部的門控。標(biāo)準(zhǔn)的縮放點(diǎn)積注意力用 softmax 對(duì) attention score 做歸一化這本身可以理解為一種“競(jìng)爭(zhēng)式門控”所有位置競(jìng)爭(zhēng)權(quán)重權(quán)重和為 1。但這是隱式門控模型不能靈活地表達(dá)“某個(gè) token 的 KV 信息整體要減弱”這類需求。第二層是前饋網(wǎng)絡(luò)中的門控。從 GLU 到 SwiGLU、SiTU-GLU本質(zhì)都是用一個(gè)非線性變換生成門控信號(hào)去調(diào)控另一路線性變換的輸出。這是目前大模型最常用的門控形式。第三層是MoE 中的路由門控。MoEMixture of Experts混合專家模型里有一個(gè) router它決定每個(gè) token 被送到哪幾個(gè)專家網(wǎng)絡(luò)去計(jì)算。這層門控做的是“稀疏調(diào)度”和注意力門控關(guān)注“通道強(qiáng)度”不同路由門控關(guān)注“選擇性激活”。理解了這三層門控再看 Kimi K3 的各個(gè)模塊就容易對(duì)號(hào)入座了。1.3 Kimi K3 中的門控體系根據(jù)公開資料和社區(qū)討論Kimi K3 的架構(gòu)中門控機(jī)制分布得很廣Gated MLA在多頭潛在注意力MLA中引入門控分支對(duì) KV 信息做軟過濾。KDA 的門控分支在鍵解耦注意力Key-Decoupled Attention路徑中用門控融合不同來源的 key 信息。MoE 中的 SiTU-GLU 與 SwiGLU專家網(wǎng)絡(luò)內(nèi)部使用門控激活函數(shù)路由門控負(fù)責(zé)專家選擇。這三個(gè)方向剛好對(duì)應(yīng)上面說的三層門控。下面逐個(gè)拆開講。2. 環(huán)境準(zhǔn)備與符號(hào)約定2.1 Python 環(huán)境與依賴本文的代碼示例全部使用 PyTorch 編寫。如果你的電腦上還沒有環(huán)境可以這樣準(zhǔn)備conda create -n gating python3.10 conda activate gating pip install torch版本方面以下代碼基于 Python 3.8 與 PyTorch 2.x 編寫不同小版本之間的 API 基本一致。如果你使用 CPU 版本的 PyTorch也能運(yùn)行本文所有示例因?yàn)槭纠簧婕靶∫?guī)模張量計(jì)算。2.2 本文使用的數(shù)學(xué)符號(hào)為了不讓公式讀起來太吃力先約定幾個(gè)符號(hào)(x)輸入向量或輸入序列shape 為[batch, seq_len, d_model](W)線性投影權(quán)重矩陣(\sigma)sigmoid 激活函數(shù)(\otimes)逐元素相乘Hadamard product(d_model)隱藏層維度(n_heads)注意力頭數(shù)(head_dim)每個(gè)注意力頭的維度(latent_dim)MLA 中低秩壓縮后的潛在維度3. 從激活函數(shù)到門控前饋SwiGLU 與 SiTU-GLU3.1 為什么 LLM 不再只用 ReLU早期 Transformer 的前饋網(wǎng)絡(luò)FFN非常簡(jiǎn)單[ FFN(x) \text{ReLU}(xW_1 b_1)W_2 b_2 ]ReLU 把負(fù)數(shù)直接置 0正數(shù)原樣通過。這個(gè)“硬門控”簡(jiǎn)單高效但有兩個(gè)問題負(fù)數(shù)區(qū)間梯度為 0容易造成神經(jīng)元死亡。輸出均值不為 0不利于深層網(wǎng)絡(luò)的訓(xùn)練穩(wěn)定性。后來研究者發(fā)現(xiàn)用平滑激活函數(shù)替代 ReLU 能提升訓(xùn)練穩(wěn)定性和模型質(zhì)量。于是有了 GELU、Swish / SiLU 等替代方案。但真正帶來質(zhì)變的是把激活函數(shù)放進(jìn)一個(gè)“門控結(jié)構(gòu)”里也就是 GLU 系列。3.2 GLU 與 SwiGLU 的原理GLUGated Linear Unit門控線性單元最早由 Dauphin 等人提出它的形式是[ GLU(x) (xW_1 b_1) \otimes \sigma(xW_2 b_2) ]這里有兩路線性變換一路直接輸出叫做“值分支”。另一路經(jīng)過 sigmoid輸出 0 到 1 的門控信號(hào)。兩個(gè)結(jié)果逐元素相乘。模型通過訓(xùn)練可以決定每個(gè)維度上信息通過的比例。SwiGLU 是 LLaMA 等大模型使用的變體核心改動(dòng)是把 sigmoid 換成 SwishSiLU門控。[ SwiGLU(x) (xW_1) \otimes \text{SiLU}(xW_2) ]其中[ \text{SiLU}(x) x \cdot \sigma(x) ]SiLU 的形狀比 sigmoid 更豐富當(dāng)輸入為負(fù)且絕對(duì)值較大時(shí)SiLU 輸出會(huì)先輕微下降再趨近 0這種非單調(diào)性給門控帶來了更強(qiáng)的表達(dá)能力。3.3 SiTU-GLU門控激活的另一種形態(tài)SiTU-GLU 是 Kimi K3 相關(guān)討論中頻繁出現(xiàn)的一個(gè)詞。SiTU 通常被理解為 Sigmoid-Tanh Unit 的縮寫。雖然官方準(zhǔn)確公式還沒有完整公開但社區(qū)討論中較多提到的一種形式是[ SiTU(x) x \cdot \left( \frac{\sigma(x) \tanh(x)}{2} \right) ]也就是說把 sigmoid 和 tanh 兩種門控曲線做平均再與輸入相乘。sigmoid 的取值范圍是 0 到 1tanh 的取值范圍是 -1 到 1兩者結(jié)合后門控信號(hào)不再局限于 0~1而是可以輸出負(fù)值相當(dāng)于“允許信息反向調(diào)制”。把 SiTU 放進(jìn) GLU 結(jié)構(gòu)就得到 SiTU-GLU[ SiTU_GLU(x) (xW_1) \otimes SiTU(xW_2) ]注意這個(gè)推導(dǎo)是社區(qū)討論形式不等同于官方實(shí)現(xiàn)。對(duì)于學(xué)習(xí)來說重要的不是記住某個(gè)公式而是理解 SiTU-GLU 屬于“門控激活 線性變換”的家族。3.4 代碼實(shí)現(xiàn)PyTorch 版 SwiGLU / SiTU-GLU下面用 PyTorch 實(shí)現(xiàn) SwiGLU 和 SiTU-GLU。import torch import torch.nn as nn import torch.nn.functional as F class SwiGLU(nn.Module): SwiGLU 門控前饋層 輸入 x 會(huì)經(jīng)過兩條分支 - 分支1線性變換 w - 值分支 - 分支2線性變換 v SiLU - 門控分支 最后逐元素相乘。 def __init__(self, dim_in, dim_out, biasFalse): super().__init__() self.w nn.Linear(dim_in, dim_out, biasbias) self.v nn.Linear(dim_in, dim_out, biasbias) def forward(self, x): return self.w(x) * F.silu(self.v(x)) class SiTUGLU(nn.Module): SiTU-GLU 教學(xué)簡(jiǎn)化版 門控分支使用 (sigmoid tanh) / 2 作為門控信號(hào)。 這里只是社區(qū)討論形式的示例實(shí)際參數(shù)以官方為準(zhǔn)。 def __init__(self, dim_in, dim_out, biasFalse): super().__init__() self.w nn.Linear(dim_in, dim_out, biasbias) self.v nn.Linear(dim_in, dim_out, biasbias) def forward(self, x): gate_input self.v(x) gate (torch.sigmoid(gate_input) torch.tanh(gate_input)) / 2.0 return self.w(x) * gate # 測(cè)試示例 if __name__ __main__: torch.manual_seed(42) x torch.randn(4, 16, 128) # [batch, seq_len, d_model] swiglu SwiGLU(128, 256) situ_glu SiTUGLU(128, 256) out1 swiglu(x) out2 situ_glu(x) print(SwiGLU output:, out1.shape) # torch.Size([4, 16, 256]) print(SiTU-GLU output:, out2.shape) # torch.Size([4, 16, 256])這段代碼里最關(guān)鍵的一行是return self.w(x) * F.silu(self.v(x))它體現(xiàn)了 SwiGLU 的靈魂一個(gè)分支負(fù)責(zé)內(nèi)容變換另一個(gè)分支負(fù)責(zé)生成門控權(quán)重。4. Gated MLA門控多頭潛在注意力拆解4.1 先回顧 MLA 要解決的問題要理解 Gated MLA得先理解 MLAMulti-head Latent Attention多頭潛在注意力。MLA 最初由 DeepSeek-V2 引入核心動(dòng)機(jī)是降低 KV cache 的顯存占用。在標(biāo)準(zhǔn) MHAMulti-Head Attention中每個(gè)注意力頭都要緩存一份完整的 K 和 V推理時(shí)長(zhǎng)序列時(shí)顯存開銷非常大。MLA 的思路是先用一個(gè)低秩壓縮矩陣 (W_{DKV}) 把隱藏狀態(tài)壓縮成維度很小的 latent 向量。在注意力計(jì)算時(shí)再把這個(gè) latent 向量解壓回完整的 K、V 矩陣。這樣一來推理時(shí)只需要緩存低維的 latent 向量KV cache 的大小大幅降低。但這里有一個(gè)問題RoPE 位置編碼需要在 K 上疊加位置信息而低秩壓縮后的 latent 空間不方便直接做這個(gè)操作。DeepSeek-V2 的解法是“解耦 RoPE”把一部分維度拿出來單獨(dú)做位置編碼另一部分維度保持原本的低秩壓縮邏輯。4.2 Gate 加在哪里Gated MLA 可以理解成“在 MLA 的 latent 空間里插入一個(gè)門控分支”。Kimi K3 的公開討論中Gated MLA 的門控分支通常被認(rèn)為放在低秩壓縮之后、解壓之前或解壓過程中它的作用是對(duì) latent 向量的每個(gè)維度生成一個(gè) 0 到 1 之間的權(quán)重。用這個(gè)權(quán)重去調(diào)制 K、V 的信息強(qiáng)弱。為什么要這么做因?yàn)樵跇?biāo)準(zhǔn)注意力中所有 token 的 KV 信息都會(huì)被同等對(duì)待但真實(shí)場(chǎng)景并非如此。有些 token 的信息對(duì)當(dāng)前 query 是冗余的甚至是有干擾的有些 token 的 KV 信息則需要被放大。門控分支讓模型學(xué)會(huì)“按需放行”。這比直接用 attention score 調(diào)節(jié)更精細(xì)。attention score 調(diào)節(jié)的是“query 對(duì) key 的匹配程度”而 Gated MLA 調(diào)節(jié)的是“key/value 本身要保留多少信息”。4.3 可運(yùn)行的 SimplifiedGatedMLA 示例下面實(shí)現(xiàn)一個(gè)教學(xué)用的 Gated MLA 簡(jiǎn)化版本不包含完整的 RoPE 處理但保留低秩壓縮和門控分支兩個(gè)核心思想。import math import torch import torch.nn as nn import torch.nn.functional as F class SimplifiedGatedMLA(nn.Module): 簡(jiǎn)化版 Gated MLA 流程 1. 將輸入 x 壓縮為低維 latent 2. 從 latent 解壓出 K、V 3. 根據(jù) latent 生成門控權(quán)重調(diào)制 K、V 4. 與 Q 做標(biāo)準(zhǔn)縮放點(diǎn)積注意力。 def __init__(self, d_model, n_heads, latent_dim, head_dim): super().__init__() assert d_model n_heads * head_dim, d_model 需等于 n_heads * head_dim self.n_heads n_heads self.head_dim head_dim self.latent_dim latent_dim # 壓縮隱藏狀態(tài) - 低維 latent self.w_dkv nn.Linear(d_model, latent_dim, biasFalse) # 解壓latent - K、V self.w_uk nn.Linear(latent_dim, n_heads * head_dim, biasFalse) self.w_uv nn.Linear(latent_dim, n_heads * head_dim, biasFalse) # Q 直接投影 self.w_q nn.Linear(d_model, n_heads * head_dim, biasFalse) # 門控分支從 latent 生成門控信號(hào) self.gate_proj nn.Linear(latent_dim, n_heads * head_dim, biasFalse) # 輸出投影 self.w_o nn.Linear(n_heads * head_dim, d_model, biasFalse) def forward(self, x): B, T, D x.shape # 1. 低秩壓縮 latent self.w_dkv(x) # [B, T, latent_dim] # 2. 解壓得到 K、V k self.w_uk(latent) # [B, T, n_heads * head_dim] v self.w_uv(latent) # [B, T, n_heads * head_dim] # 3. 門控調(diào)制 gate torch.sigmoid(self.gate_proj(latent)) # [B, T, n_heads * head_dim] k k * gate v v * gate # 4. Query 投影 q self.w_q(x) # [B, T, n_heads * head_dim] # 5. 多頭拆分 def reshape_to_heads(t): return t.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) q reshape_to_heads(q) # [B, n_heads, T, head_dim] k reshape_to_heads(k) v reshape_to_heads(v) # 6. 縮放點(diǎn)積注意力 attn_scores (q k.transpose(-2, -1)) / math.sqrt(self.head_dim) attn_weights F.softmax(attn_scores, dim-1) out attn_weights v # [B, n_heads, T, head_dim] # 7. 合并多頭并輸出 out out.transpose(1, 2).reshape(B, T, -1) return self.w_o(out) # 運(yùn)行測(cè)試 if __name__ __main__: torch.manual_seed(0) model SimplifiedGatedMLA( d_model128, n_heads4, latent_dim64, head_dim32 ) x torch.randn(2, 10, 128) # [batch2, seq_len10, d_model128] y model(x) print(Gated MLA output:, y.shape) # torch.Size([2, 10, 128])這個(gè)簡(jiǎn)化實(shí)現(xiàn)有幾個(gè)可以繼續(xù)深挖的點(diǎn)實(shí)際 MLA 中解壓后的 K、V 會(huì)先拆分出“內(nèi)容部分”和“RoPE 部分”分別處理再拼接。門控分支可以加在 latent 上也可以加在解壓后的 K、V 上。不同位置效果不同。真實(shí)實(shí)現(xiàn)里 (Q) 也會(huì)做低秩投影這里為了可讀性直接使用完整 Q 投影。4.4 門控分支對(duì)推理成本的影響Gated MLA 的工程價(jià)值在于門控分支在 latent 空間計(jì)算維度通常遠(yuǎn)小于完整的 K/V 維度。以latent_dim64、n_heads * head_dim128為例門控分支需要的計(jì)算量大約是 K/V 解壓后做門控的 1/2。在大模型推理場(chǎng)景下這部分額外計(jì)算成本是可以接受的因?yàn)楣?jié)省的 KV cache 顯存遠(yuǎn)大于新增門控的算力開銷。5. KDA 的門控分支注意力路徑的信息分流5.1 KDA 的基本思路KDA 通常被理解為 Key-Decoupled Attention鍵解耦注意力。社區(qū)中關(guān)于 Kimi K3 KDA 的討論重點(diǎn)在于“把 key 路徑拆成多個(gè)分支”。標(biāo)準(zhǔn)注意力中每個(gè) token 只有一個(gè) key 向量它同時(shí)承擔(dān)內(nèi)容語義匹配和位置關(guān)系匹配兩個(gè)職責(zé)。這實(shí)際上是一個(gè)隱式的多任務(wù)耦合content matching 和 position matching 共享同一個(gè)向量模型很難獨(dú)立調(diào)節(jié)兩者的權(quán)重。KDA 的思路是把 key 拆開內(nèi)容分支負(fù)責(zé) token 本身的語義相關(guān)性。位置分支負(fù)責(zé) token 之間的相對(duì)位置關(guān)系。門控分支學(xué)習(xí)一個(gè)權(quán)重決定最終 attention score 在多大程度上依賴內(nèi)容分支、多大程度上依賴位置分支。5.2 門控分支如何參與融合假設(shè)我們有(q)query 向量(k_{content})內(nèi)容分支的 key(k_{position})位置分支的 key傳統(tǒng)做法可能是直接拼接或者相加[ score q \cdot (k_{content} k_{position}) ]KDA 門控分支的做法是學(xué)習(xí)一個(gè)門控參數(shù) (g)讓模型自動(dòng)決定兩種信息的占比[ score g \cdot (q \cdot k_{content}) (1 - g) \cdot (q \cdot k_{position}) ]當(dāng) (g) 接近 1 時(shí)注意力主要由內(nèi)容語義驅(qū)動(dòng)當(dāng) (g) 接近 0 時(shí)注意力主要由位置關(guān)系驅(qū)動(dòng)。這種門控的優(yōu)勢(shì)是讓模型在不同層、不同注意力頭之間形成分工。有的頭可能更偏向內(nèi)容匹配有的頭更偏向位置匹配門控機(jī)制讓這種分工變得顯式可學(xué)。5.3 偽代碼與實(shí)現(xiàn)思路下面給出一個(gè)教學(xué)級(jí)別的偽代碼展示 KDA 門控分支的數(shù)據(jù)流def kda_attention(q, k_content, k_position, v, gate_logits): 教學(xué)簡(jiǎn)化版 KDA 門控分支 gate_logits: [B, n_heads, T, 1] 或 [B, n_heads, T, T] 這里簡(jiǎn)化為每個(gè) token 一個(gè)門控權(quán)重。 score_content q k_content.transpose(-2, -1) score_position q k_position.transpose(-2, -1) # 門控權(quán)重在 0~1 之間 gate torch.sigmoid(gate_logits) # 信息融合 total_score gate * score_content (1 - gate) * score_position attn_weight F.softmax(total_score, dim-1) out attn_weight v return out, gate實(shí)際實(shí)現(xiàn)中g(shù)ate_logits可以由 query 生成也可以由 query 和 key 的交互生成。前者計(jì)算量更小后者表達(dá)能力更強(qiáng)具體怎么選是一個(gè)效率和效果的權(quán)衡。6. MoE 里的門控路由與專家內(nèi)部門控6.1 MoE 路由門控的本質(zhì)MoEMixture of Experts混合專家是當(dāng)前超大模型的主流架構(gòu)。一個(gè) MoE 層通常包含一個(gè)路由網(wǎng)絡(luò)router。若干專家網(wǎng)絡(luò)experts。路由網(wǎng)絡(luò)本質(zhì)上就是一個(gè)門控分類器輸入 token 的表示輸出每個(gè)專家的選擇概率。通常的做法是用線性層把 token 映射到專家數(shù)量維度的 logits。softmax 得到概率分布。top-k 選出得分最高的 k 個(gè)專家。對(duì)選中的概率做歸一化作為加權(quán)系數(shù)。路由門控的數(shù)學(xué)形式如下[ p_i \frac{\exp((x \cdot W_r)i)}{\sum{j1}^{N}\exp((x \cdot W_r)_j)} ]然后取 top-k得到最終組合權(quán)重。6.2 路由門控的負(fù)載均衡挑戰(zhàn)路由門控最經(jīng)典的問題是負(fù)載不均衡如果某個(gè)專家總被選中其他專家?guī)缀醪槐皇褂谜麄€(gè) MoE 層就退化成單專家模型稀疏性的收益全部消失。解決辦法是加輔助均衡損失。常見的做法是計(jì)算每個(gè)專家的平均路由概率。計(jì)算每個(gè)專家的平均被選中次數(shù)。讓兩個(gè)分布盡量接近。公式可以表達(dá)為[ L_{balance} \alpha \cdot N \cdot \sum_{i1}^{N} f_i \cdot p_i ]其中 (f_i) 是第 (i) 個(gè)專家被選中的頻率(p_i) 是路由平均概率(N) 是專家數(shù)量(\alpha) 是平衡系數(shù)。6.3 專家內(nèi)部門控SiTU-GLU 與 SwiGLU 的定位MoE 里的門控其實(shí)存在兩層第一層是路由門控決定“選擇誰”。第二層是專家內(nèi)部的前饋門控決定“信息怎么變換”。SwiGLU 和 SiTU-GLU 屬于第二層。專家網(wǎng)絡(luò)本質(zhì)上是一個(gè) FFN而 FFN 的內(nèi)部結(jié)構(gòu)正好可以用 GLU 系列激活函數(shù)來強(qiáng)化。所以一個(gè)典型的 MoE 專家可以寫成Expert(x) OutputProj( GateActivation(InputProj(x)) * ValueProj(x) )其中GateActivation可以是 SiLU對(duì)應(yīng) SwiGLU或 SiTU對(duì)應(yīng) SiTU-GLU。6.4 完整示例Top-K Router SwiGLU 專家下面把路由門控和 SwiGLU 專家組合成一個(gè)可運(yùn)行的簡(jiǎn)化 MoE 層。import torch import torch.nn as nn import torch.nn.functional as F class TopKRouter(nn.Module): 簡(jiǎn)化版 Top-K 路由器 輸入: [B, T, d_model] 輸出: top_idx: [B, T, top_k] 每個(gè) token 選中的專家編號(hào) top_probs: [B, T, top_k] 歸一化后的路由權(quán)重 def __init__(self, d_model, n_experts, top_k2): super().__init__() self.top_k top_k self.router nn.Linear(d_model, n_experts, biasFalse) def forward(self, x): logits self.router(x) # [B, T, n_experts] probs F.softmax(logits, dim-1) top_probs, top_idx torch.topk(probs, self.top_k, dim-1) # 對(duì)選中的概率重新歸一化 top_probs top_probs / top_probs.sum(dim-1, keepdimTrue) return top_idx, top_probs class SimpleMoE(nn.Module): 簡(jiǎn)化版 MoE Layer 每個(gè)專家內(nèi)部使用 SwiGLU 風(fēng)格的前饋網(wǎng)絡(luò)。 def __init__(self, d_model, n_experts, top_k, hidden_dim): super().__init__() self.top_k top_k self.router nn.Linear(d_model, n_experts, biasFalse) # 每個(gè)專家: 兩個(gè)線性層 SiLU 門控 self.experts nn.ModuleList([ nn.Sequential( SwiGLU(d_model, hidden_dim), nn.Linear(hidden_dim, d_model) ) for _ in range(n_experts) ]) def forward(self, x): B, T, D x.shape x_flat x.reshape(-1, D) # [B*T, D] # 路由 logits self.router(x_flat) # [B*T, n_experts] probs F.softmax(logits, dim-1) top_probs, top_idx torch.topk(probs, self.top_k, dim-1) top_probs top_probs / top_probs.sum(dim-1, keepdimTrue) out torch.zeros_like(x_flat) # 逐專家計(jì)算貢獻(xiàn) for e, expert in enumerate(self.experts): # 哪些 token 選中了專家 e mask (top_idx e).any(dim-1) if not mask.any(): continue expert_out expert(x_flat[mask]) # 提取每個(gè) token 對(duì)專家 e 的權(quán)重 e_probs torch.where( top_idx[mask] e, top_probs[mask], torch.zeros_like(top_probs[mask]) ) weight e_probs.sum(dim-1, keepdimTrue) # [n_selected, 1] out[mask] expert_out * weight return out.view(B, T, D) # 測(cè)試 if __name__ __main__: torch.manual_seed(42) moe SimpleMoE( d_model128, n_experts8, top_k2, hidden_dim256 ) x torch.randn(2, 10, 128) y moe(x) print(MoE output:, y.shape) # torch.Size([2, 10, 128])如果你希望訓(xùn)練更穩(wěn)定可以把top_k換成可學(xué)習(xí)的 soft 門控權(quán)重但在稠密場(chǎng)景下 Top-K 的效果已經(jīng)足夠好。7. 常見誤區(qū)與排查清單門控相關(guān)概念多且相似下面整理幾個(gè)高頻誤區(qū)和排查思路。問題現(xiàn)象常見原因解決思路把 Gated MLA 理解成 MHA 后面加個(gè) sigmoidMLA 的核心是低秩 KV 壓縮門控是在 latent 空間調(diào)制 K/V先理解 MLA 的壓縮-解壓過程再看門控插入位置認(rèn)為 KDA 就是 GQAGQA 減少 KV 頭數(shù)量KDA 是把 key 路徑拆成多分支并用門控融合畫出 key 分支結(jié)構(gòu)對(duì)比把 SwiGLU 和 SiLU 當(dāng)成同一個(gè)東西SwiGLU 是門控線性單元SiLU 是激活函數(shù)二者不是同一層面SwiGLU 值分支 x SiLU(門控分支)訓(xùn)練 MoE 時(shí)某個(gè)專家始終不被選中路由初始化不當(dāng)或缺乏負(fù)載均衡損失調(diào)整 router 初始化加入 balance loss門控輸出經(jīng)常飽和梯度消失sigmoid 輸入絕對(duì)值過大檢查 gate_proj 的初始化或使用帶下界的門控變體推理時(shí)門控分支拖慢速度門控在完整 K/V 維度計(jì)算把門控放在 latent 維度或與 KV 解壓算子融合8. 工程實(shí)踐建議與學(xué)習(xí)路線8.1 架構(gòu)設(shè)計(jì)建議在實(shí)際設(shè)計(jì)或復(fù)現(xiàn)帶門控的 LLM 模塊時(shí)有幾點(diǎn)值得留意。首先門控分支的位置直接影響信息選擇效果。Gated MLA 中門控放在 latent 空間更高效但表達(dá)力可能不如放在解壓后的 K/V 空間如果顯存允許可以選擇在兩組位置同時(shí)加門控再用殘差連接做兜底。其次初始化很重要。sigmoid 在輸入為 0 時(shí)輸出 0.5如果門控權(quán)重初始化過大信號(hào)會(huì)接近飽和梯度難以流動(dòng)。實(shí)踐中可以讓gate_proj的輸出初始偏向 0 附近甚至給 bias 設(shè)置一個(gè)負(fù)初始值讓門控初始偏向“關(guān)閉”以穩(wěn)定早期訓(xùn)練。8.2 訓(xùn)練與推理注意事項(xiàng)訓(xùn)練階段要關(guān)注門控是否存在“退化成常數(shù)”的問題。如果訓(xùn)練后期門控輸出始終固定在某個(gè)值附近說明門控沒有學(xué)到有效的信息選擇邏輯可以檢查一下梯度或嘗試給門控分支添加正則。推理階段Gated MLA 和 KDA 門控分支的額外算子會(huì)帶來 kernel launch 開銷。真實(shí)部署時(shí)建議把gate_proj、w_uk、w_uv等線性層合并成一個(gè)大矩陣乘法減少訪存和 kernel 啟停次數(shù)。對(duì) MoE 來說路由門控的數(shù)值精度要留心。在低精度推理時(shí)router logits 的微小波動(dòng)可能導(dǎo)致 top-k 選擇結(jié)果變化進(jìn)而影響生成穩(wěn)定性。必要時(shí)用更高精度計(jì)算 router logits或者對(duì) router logits 做縮放防止 softmax 后概率過于集中。8.3 學(xué)習(xí)路線如果你想系統(tǒng)掌握門控機(jī)制可以按這個(gè)順序走從激活函數(shù)入手實(shí)現(xiàn)并對(duì)比 ReLU、GELU、SiLU、Mish。實(shí)現(xiàn) GLU 和 SwiGLU觀察門控分支對(duì)梯度流動(dòng)的影響。實(shí)現(xiàn)標(biāo)準(zhǔn) MHA再實(shí)現(xiàn) DeepSeek-V2 的 MLA最后加上門控分支。閱讀 LLaMA 系列和 Mistral 的代碼看 SwiGLU 在實(shí)際模型中如何落地。嘗試實(shí)現(xiàn)一個(gè)包含 router 和專家網(wǎng)絡(luò)的 MoE 層加入負(fù)載均衡 loss?；氐?Kimi K3 的公開架構(gòu)資料對(duì)照本文的模塊圖逐層印證。如果條件允許可以下載一個(gè)中等規(guī)模的開源 MoE 模型用推理框架觀察不同專家被激活的頻率配合日志分析路由門控的行為。這會(huì)比單純看論文更有體感。門控機(jī)制是理解現(xiàn)代大模型的一條重要線索。從 LSTM 的遺忘門到 SwiGLU 的門控前饋再到 Gated MLA 和 MoE 路由本質(zhì)上都在做同一件事讓模型自己學(xué)會(huì)信息該如何選擇、何時(shí)放行、以什么比例放行。建議你打開編輯器把文中的 SimplifiedGatedMLA 和 SimpleMoE 各跑一遍然后試著調(diào)整門控位置、初始化和 top-k 大小直觀感受這些改動(dòng)對(duì)訓(xùn)練和輸出的影響。動(dòng)手跑通之后再回去看 Kimi K3 的架構(gòu)分析會(huì)順手很多。 大模型門控機(jī)制全解析：從SwiGLU到Gated MLA與MoE路由
最近在梳理 Kimi K3 的架構(gòu)資料時(shí)有一個(gè)詞被反復(fù)提及門控Gating。Gated MLA、KDA 的門控分支、MoE 中的 SiTU-GLU、SwiGLU……看起來每個(gè)模塊都在做“門控”但每個(gè)地方的門控含義又不太一樣。如果只是零散地看很容易把一個(gè)概念套到另一個(gè)模塊上最后越看越亂。這篇文章就把 LLM 里的門控機(jī)制集中梳理一遍結(jié)合 Kimi K3 的公開架構(gòu)討論拆開 Gated MLA、KDA 門控分支以及 MoE 里的 SiTU-GLU 與 SwiGLU。我會(huì)盡量不堆術(shù)語先講清楚門控本身是什么再逐層進(jìn)入現(xiàn)代大模型的門控設(shè)計(jì)并給出可運(yùn)行的 PyTorch 簡(jiǎn)化實(shí)現(xiàn)。需要提前說明的是Kimi K3 的官方技術(shù)報(bào)告尚未完整公開本文的架構(gòu)解析基于公開模型信息、社區(qū)復(fù)現(xiàn)討論和技術(shù)博客整理代碼示例是教學(xué)用簡(jiǎn)化版本不完全等價(jià)于官方實(shí)現(xiàn)。1. 門控機(jī)制LLM 架構(gòu)里的“信息閥門”1.1 從 LSTM 說起門控是什么門控不是新概念。早在 LSTM長(zhǎng)短期記憶網(wǎng)絡(luò)時(shí)代門控就是核心機(jī)制。LSTM 里有三個(gè)門遺忘門決定上一時(shí)刻的記憶要保留多少。輸入門決定當(dāng)前輸入有多少寫入記憶。輸出門決定當(dāng)前記憶有多少輸出到隱藏狀態(tài)。每個(gè)門都輸出一個(gè) 0 到 1 之間的數(shù)值用 sigmoid 激活函數(shù)實(shí)現(xiàn)。0 表示“完全關(guān)閉”1 表示“完全打開”中間值則表示“部分通過”。用一句話概括門控的本質(zhì)學(xué)習(xí)一個(gè)軟開關(guān)決定信息以什么比例通過。為什么要用“軟”開關(guān)因?yàn)橛查_關(guān)if-else不可導(dǎo)梯度無法回傳軟開關(guān)用 sigmoid 這類連續(xù)函數(shù)梯度能順暢流動(dòng)模型就能通過反向傳播自動(dòng)學(xué)習(xí)開關(guān)的開合程度。1.2 Transformer 時(shí)代的三類門控Transformer 出現(xiàn)后門控并沒有消失而是演化出了更豐富的形態(tài)。在當(dāng)今的大模型架構(gòu)里門控至少出現(xiàn)在三個(gè)層面第一層是注意力內(nèi)部的門控。標(biāo)準(zhǔn)的縮放點(diǎn)積注意力用 softmax 對(duì) attention score 做歸一化這本身可以理解為一種“競(jìng)爭(zhēng)式門控”所有位置競(jìng)爭(zhēng)權(quán)重權(quán)重和為 1。但這是隱式門控模型不能靈活地表達(dá)“某個(gè) token 的 KV 信息整體要減弱”這類需求。第二層是前饋網(wǎng)絡(luò)中的門控。從 GLU 到 SwiGLU、SiTU-GLU本質(zhì)都是用一個(gè)非線性變換生成門控信號(hào)去調(diào)控另一路線性變換的輸出。這是目前大模型最常用的門控形式。第三層是MoE 中的路由門控。MoEMixture of Experts混合專家模型里有一個(gè) router它決定每個(gè) token 被送到哪幾個(gè)專家網(wǎng)絡(luò)去計(jì)算。這層門控做的是“稀疏調(diào)度”和注意力門控關(guān)注“通道強(qiáng)度”不同路由門控關(guān)注“選擇性激活”。理解了這三層門控再看 Kimi K3 的各個(gè)模塊就容易對(duì)號(hào)入座了。1.3 Kimi K3 中的門控體系根據(jù)公開資料和社區(qū)討論Kimi K3 的架構(gòu)中門控機(jī)制分布得很廣Gated MLA在多頭潛在注意力MLA中引入門控分支對(duì) KV 信息做軟過濾。KDA 的門控分支在鍵解耦注意力Key-Decoupled Attention路徑中用門控融合不同來源的 key 信息。MoE 中的 SiTU-GLU 與 SwiGLU專家網(wǎng)絡(luò)內(nèi)部使用門控激活函數(shù)路由門控負(fù)責(zé)專家選擇。這三個(gè)方向剛好對(duì)應(yīng)上面說的三層門控。下面逐個(gè)拆開講。2. 環(huán)境準(zhǔn)備與符號(hào)約定2.1 Python 環(huán)境與依賴本文的代碼示例全部使用 PyTorch 編寫。如果你的電腦上還沒有環(huán)境可以這樣準(zhǔn)備conda create -n gating python3.10 conda activate gating pip install torch版本方面以下代碼基于 Python 3.8 與 PyTorch 2.x 編寫不同小版本之間的 API 基本一致。如果你使用 CPU 版本的 PyTorch也能運(yùn)行本文所有示例因?yàn)槭纠簧婕靶∫?guī)模張量計(jì)算。2.2 本文使用的數(shù)學(xué)符號(hào)為了不讓公式讀起來太吃力先約定幾個(gè)符號(hào)(x)輸入向量或輸入序列shape 為[batch, seq_len, d_model](W)線性投影權(quán)重矩陣(\sigma)sigmoid 激活函數(shù)(\otimes)逐元素相乘Hadamard product(d_model)隱藏層維度(n_heads)注意力頭數(shù)(head_dim)每個(gè)注意力頭的維度(latent_dim)MLA 中低秩壓縮后的潛在維度3. 從激活函數(shù)到門控前饋SwiGLU 與 SiTU-GLU3.1 為什么 LLM 不再只用 ReLU早期 Transformer 的前饋網(wǎng)絡(luò)FFN非常簡(jiǎn)單[ FFN(x) \text{ReLU}(xW_1 b_1)W_2 b_2 ]ReLU 把負(fù)數(shù)直接置 0正數(shù)原樣通過。這個(gè)“硬門控”簡(jiǎn)單高效但有兩個(gè)問題負(fù)數(shù)區(qū)間梯度為 0容易造成神經(jīng)元死亡。輸出均值不為 0不利于深層網(wǎng)絡(luò)的訓(xùn)練穩(wěn)定性。后來研究者發(fā)現(xiàn)用平滑激活函數(shù)替代 ReLU 能提升訓(xùn)練穩(wěn)定性和模型質(zhì)量。于是有了 GELU、Swish / SiLU 等替代方案。但真正帶來質(zhì)變的是把激活函數(shù)放進(jìn)一個(gè)“門控結(jié)構(gòu)”里也就是 GLU 系列。3.2 GLU 與 SwiGLU 的原理GLUGated Linear Unit門控線性單元最早由 Dauphin 等人提出它的形式是[ GLU(x) (xW_1 b_1) \otimes \sigma(xW_2 b_2) ]這里有兩路線性變換一路直接輸出叫做“值分支”。另一路經(jīng)過 sigmoid輸出 0 到 1 的門控信號(hào)。兩個(gè)結(jié)果逐元素相乘。模型通過訓(xùn)練可以決定每個(gè)維度上信息通過的比例。SwiGLU 是 LLaMA 等大模型使用的變體核心改動(dòng)是把 sigmoid 換成 SwishSiLU門控。[ SwiGLU(x) (xW_1) \otimes \text{SiLU}(xW_2) ]其中[ \text{SiLU}(x) x \cdot \sigma(x) ]SiLU 的形狀比 sigmoid 更豐富當(dāng)輸入為負(fù)且絕對(duì)值較大時(shí)SiLU 輸出會(huì)先輕微下降再趨近 0這種非單調(diào)性給門控帶來了更強(qiáng)的表達(dá)能力。3.3 SiTU-GLU門控激活的另一種形態(tài)SiTU-GLU 是 Kimi K3 相關(guān)討論中頻繁出現(xiàn)的一個(gè)詞。SiTU 通常被理解為 Sigmoid-Tanh Unit 的縮寫。雖然官方準(zhǔn)確公式還沒有完整公開但社區(qū)討論中較多提到的一種形式是[ SiTU(x) x \cdot \left( \frac{\sigma(x) \tanh(x)}{2} \right) ]也就是說把 sigmoid 和 tanh 兩種門控曲線做平均再與輸入相乘。sigmoid 的取值范圍是 0 到 1tanh 的取值范圍是 -1 到 1兩者結(jié)合后門控信號(hào)不再局限于 0~1而是可以輸出負(fù)值相當(dāng)于“允許信息反向調(diào)制”。把 SiTU 放進(jìn) GLU 結(jié)構(gòu)就得到 SiTU-GLU[ SiTU_GLU(x) (xW_1) \otimes SiTU(xW_2) ]注意這個(gè)推導(dǎo)是社區(qū)討論形式不等同于官方實(shí)現(xiàn)。對(duì)于學(xué)習(xí)來說重要的不是記住某個(gè)公式而是理解 SiTU-GLU 屬于“門控激活 線性變換”的家族。3.4 代碼實(shí)現(xiàn)PyTorch 版 SwiGLU / SiTU-GLU下面用 PyTorch 實(shí)現(xiàn) SwiGLU 和 SiTU-GLU。import torch import torch.nn as nn import torch.nn.functional as F class SwiGLU(nn.Module): SwiGLU 門控前饋層 輸入 x 會(huì)經(jīng)過兩條分支 - 分支1線性變換 w - 值分支 - 分支2線性變換 v SiLU - 門控分支 最后逐元素相乘。 def __init__(self, dim_in, dim_out, biasFalse): super().__init__() self.w nn.Linear(dim_in, dim_out, biasbias) self.v nn.Linear(dim_in, dim_out, biasbias) def forward(self, x): return self.w(x) * F.silu(self.v(x)) class SiTUGLU(nn.Module): SiTU-GLU 教學(xué)簡(jiǎn)化版 門控分支使用 (sigmoid tanh) / 2 作為門控信號(hào)。 這里只是社區(qū)討論形式的示例實(shí)際參數(shù)以官方為準(zhǔn)。 def __init__(self, dim_in, dim_out, biasFalse): super().__init__() self.w nn.Linear(dim_in, dim_out, biasbias) self.v nn.Linear(dim_in, dim_out, biasbias) def forward(self, x): gate_input self.v(x) gate (torch.sigmoid(gate_input) torch.tanh(gate_input)) / 2.0 return self.w(x) * gate # 測(cè)試示例 if __name__ __main__: torch.manual_seed(42) x torch.randn(4, 16, 128) # [batch, seq_len, d_model] swiglu SwiGLU(128, 256) situ_glu SiTUGLU(128, 256) out1 swiglu(x) out2 situ_glu(x) print(SwiGLU output:, out1.shape) # torch.Size([4, 16, 256]) print(SiTU-GLU output:, out2.shape) # torch.Size([4, 16, 256])這段代碼里最關(guān)鍵的一行是return self.w(x) * F.silu(self.v(x))它體現(xiàn)了 SwiGLU 的靈魂一個(gè)分支負(fù)責(zé)內(nèi)容變換另一個(gè)分支負(fù)責(zé)生成門控權(quán)重。4. Gated MLA門控多頭潛在注意力拆解4.1 先回顧 MLA 要解決的問題要理解 Gated MLA得先理解 MLAMulti-head Latent Attention多頭潛在注意力。MLA 最初由 DeepSeek-V2 引入核心動(dòng)機(jī)是降低 KV cache 的顯存占用。在標(biāo)準(zhǔn) MHAMulti-Head Attention中每個(gè)注意力頭都要緩存一份完整的 K 和 V推理時(shí)長(zhǎng)序列時(shí)顯存開銷非常大。MLA 的思路是先用一個(gè)低秩壓縮矩陣 (W_{DKV}) 把隱藏狀態(tài)壓縮成維度很小的 latent 向量。在注意力計(jì)算時(shí)再把這個(gè) latent 向量解壓回完整的 K、V 矩陣。這樣一來推理時(shí)只需要緩存低維的 latent 向量KV cache 的大小大幅降低。但這里有一個(gè)問題RoPE 位置編碼需要在 K 上疊加位置信息而低秩壓縮后的 latent 空間不方便直接做這個(gè)操作。DeepSeek-V2 的解法是“解耦 RoPE”把一部分維度拿出來單獨(dú)做位置編碼另一部分維度保持原本的低秩壓縮邏輯。4.2 Gate 加在哪里Gated MLA 可以理解成“在 MLA 的 latent 空間里插入一個(gè)門控分支”。Kimi K3 的公開討論中Gated MLA 的門控分支通常被認(rèn)為放在低秩壓縮之后、解壓之前或解壓過程中它的作用是對(duì) latent 向量的每個(gè)維度生成一個(gè) 0 到 1 之間的權(quán)重。用這個(gè)權(quán)重去調(diào)制 K、V 的信息強(qiáng)弱。為什么要這么做因?yàn)樵跇?biāo)準(zhǔn)注意力中所有 token 的 KV 信息都會(huì)被同等對(duì)待但真實(shí)場(chǎng)景并非如此。有些 token 的信息對(duì)當(dāng)前 query 是冗余的甚至是有干擾的有些 token 的 KV 信息則需要被放大。門控分支讓模型學(xué)會(huì)“按需放行”。這比直接用 attention score 調(diào)節(jié)更精細(xì)。attention score 調(diào)節(jié)的是“query 對(duì) key 的匹配程度”而 Gated MLA 調(diào)節(jié)的是“key/value 本身要保留多少信息”。4.3 可運(yùn)行的 SimplifiedGatedMLA 示例下面實(shí)現(xiàn)一個(gè)教學(xué)用的 Gated MLA 簡(jiǎn)化版本不包含完整的 RoPE 處理但保留低秩壓縮和門控分支兩個(gè)核心思想。import math import torch import torch.nn as nn import torch.nn.functional as F class SimplifiedGatedMLA(nn.Module): 簡(jiǎn)化版 Gated MLA 流程 1. 將輸入 x 壓縮為低維 latent 2. 從 latent 解壓出 K、V 3. 根據(jù) latent 生成門控權(quán)重調(diào)制 K、V 4. 與 Q 做標(biāo)準(zhǔn)縮放點(diǎn)積注意力。 def __init__(self, d_model, n_heads, latent_dim, head_dim): super().__init__() assert d_model n_heads * head_dim, d_model 需等于 n_heads * head_dim self.n_heads n_heads self.head_dim head_dim self.latent_dim latent_dim # 壓縮隱藏狀態(tài) - 低維 latent self.w_dkv nn.Linear(d_model, latent_dim, biasFalse) # 解壓latent - K、V self.w_uk nn.Linear(latent_dim, n_heads * head_dim, biasFalse) self.w_uv nn.Linear(latent_dim, n_heads * head_dim, biasFalse) # Q 直接投影 self.w_q nn.Linear(d_model, n_heads * head_dim, biasFalse) # 門控分支從 latent 生成門控信號(hào) self.gate_proj nn.Linear(latent_dim, n_heads * head_dim, biasFalse) # 輸出投影 self.w_o nn.Linear(n_heads * head_dim, d_model, biasFalse) def forward(self, x): B, T, D x.shape # 1. 低秩壓縮 latent self.w_dkv(x) # [B, T, latent_dim] # 2. 解壓得到 K、V k self.w_uk(latent) # [B, T, n_heads * head_dim] v self.w_uv(latent) # [B, T, n_heads * head_dim] # 3. 門控調(diào)制 gate torch.sigmoid(self.gate_proj(latent)) # [B, T, n_heads * head_dim] k k * gate v v * gate # 4. Query 投影 q self.w_q(x) # [B, T, n_heads * head_dim] # 5. 多頭拆分 def reshape_to_heads(t): return t.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) q reshape_to_heads(q) # [B, n_heads, T, head_dim] k reshape_to_heads(k) v reshape_to_heads(v) # 6. 縮放點(diǎn)積注意力 attn_scores (q k.transpose(-2, -1)) / math.sqrt(self.head_dim) attn_weights F.softmax(attn_scores, dim-1) out attn_weights v # [B, n_heads, T, head_dim] # 7. 合并多頭并輸出 out out.transpose(1, 2).reshape(B, T, -1) return self.w_o(out) # 運(yùn)行測(cè)試 if __name__ __main__: torch.manual_seed(0) model SimplifiedGatedMLA( d_model128, n_heads4, latent_dim64, head_dim32 ) x torch.randn(2, 10, 128) # [batch2, seq_len10, d_model128] y model(x) print(Gated MLA output:, y.shape) # torch.Size([2, 10, 128])這個(gè)簡(jiǎn)化實(shí)現(xiàn)有幾個(gè)可以繼續(xù)深挖的點(diǎn)實(shí)際 MLA 中解壓后的 K、V 會(huì)先拆分出“內(nèi)容部分”和“RoPE 部分”分別處理再拼接。門控分支可以加在 latent 上也可以加在解壓后的 K、V 上。不同位置效果不同。真實(shí)實(shí)現(xiàn)里 (Q) 也會(huì)做低秩投影這里為了可讀性直接使用完整 Q 投影。4.4 門控分支對(duì)推理成本的影響Gated MLA 的工程價(jià)值在于門控分支在 latent 空間計(jì)算維度通常遠(yuǎn)小于完整的 K/V 維度。以latent_dim64、n_heads * head_dim128為例門控分支需要的計(jì)算量大約是 K/V 解壓后做門控的 1/2。在大模型推理場(chǎng)景下這部分額外計(jì)算成本是可以接受的因?yàn)楣?jié)省的 KV cache 顯存遠(yuǎn)大于新增門控的算力開銷。5. KDA 的門控分支注意力路徑的信息分流5.1 KDA 的基本思路KDA 通常被理解為 Key-Decoupled Attention鍵解耦注意力。社區(qū)中關(guān)于 Kimi K3 KDA 的討論重點(diǎn)在于“把 key 路徑拆成多個(gè)分支”。標(biāo)準(zhǔn)注意力中每個(gè) token 只有一個(gè) key 向量它同時(shí)承擔(dān)內(nèi)容語義匹配和位置關(guān)系匹配兩個(gè)職責(zé)。這實(shí)際上是一個(gè)隱式的多任務(wù)耦合content matching 和 position matching 共享同一個(gè)向量模型很難獨(dú)立調(diào)節(jié)兩者的權(quán)重。KDA 的思路是把 key 拆開內(nèi)容分支負(fù)責(zé) token 本身的語義相關(guān)性。位置分支負(fù)責(zé) token 之間的相對(duì)位置關(guān)系。門控分支學(xué)習(xí)一個(gè)權(quán)重決定最終 attention score 在多大程度上依賴內(nèi)容分支、多大程度上依賴位置分支。5.2 門控分支如何參與融合假設(shè)我們有(q)query 向量(k_{content})內(nèi)容分支的 key(k_{position})位置分支的 key傳統(tǒng)做法可能是直接拼接或者相加[ score q \cdot (k_{content} k_{position}) ]KDA 門控分支的做法是學(xué)習(xí)一個(gè)門控參數(shù) (g)讓模型自動(dòng)決定兩種信息的占比[ score g \cdot (q \cdot k_{content}) (1 - g) \cdot (q \cdot k_{position}) ]當(dāng) (g) 接近 1 時(shí)注意力主要由內(nèi)容語義驅(qū)動(dòng)當(dāng) (g) 接近 0 時(shí)注意力主要由位置關(guān)系驅(qū)動(dòng)。這種門控的優(yōu)勢(shì)是讓模型在不同層、不同注意力頭之間形成分工。有的頭可能更偏向內(nèi)容匹配有的頭更偏向位置匹配門控機(jī)制讓這種分工變得顯式可學(xué)。5.3 偽代碼與實(shí)現(xiàn)思路下面給出一個(gè)教學(xué)級(jí)別的偽代碼展示 KDA 門控分支的數(shù)據(jù)流def kda_attention(q, k_content, k_position, v, gate_logits): 教學(xué)簡(jiǎn)化版 KDA 門控分支 gate_logits: [B, n_heads, T, 1] 或 [B, n_heads, T, T] 這里簡(jiǎn)化為每個(gè) token 一個(gè)門控權(quán)重。 score_content q k_content.transpose(-2, -1) score_position q k_position.transpose(-2, -1) # 門控權(quán)重在 0~1 之間 gate torch.sigmoid(gate_logits) # 信息融合 total_score gate * score_content (1 - gate) * score_position attn_weight F.softmax(total_score, dim-1) out attn_weight v return out, gate實(shí)際實(shí)現(xiàn)中g(shù)ate_logits可以由 query 生成也可以由 query 和 key 的交互生成。前者計(jì)算量更小后者表達(dá)能力更強(qiáng)具體怎么選是一個(gè)效率和效果的權(quán)衡。6. MoE 里的門控路由與專家內(nèi)部門控6.1 MoE 路由門控的本質(zhì)MoEMixture of Experts混合專家是當(dāng)前超大模型的主流架構(gòu)。一個(gè) MoE 層通常包含一個(gè)路由網(wǎng)絡(luò)router。若干專家網(wǎng)絡(luò)experts。路由網(wǎng)絡(luò)本質(zhì)上就是一個(gè)門控分類器輸入 token 的表示輸出每個(gè)專家的選擇概率。通常的做法是用線性層把 token 映射到專家數(shù)量維度的 logits。softmax 得到概率分布。top-k 選出得分最高的 k 個(gè)專家。對(duì)選中的概率做歸一化作為加權(quán)系數(shù)。路由門控的數(shù)學(xué)形式如下[ p_i \frac{\exp((x \cdot W_r)i)}{\sum{j1}^{N}\exp((x \cdot W_r)_j)} ]然后取 top-k得到最終組合權(quán)重。6.2 路由門控的負(fù)載均衡挑戰(zhàn)路由門控最經(jīng)典的問題是負(fù)載不均衡如果某個(gè)專家總被選中其他專家?guī)缀醪槐皇褂谜麄€(gè) MoE 層就退化成單專家模型稀疏性的收益全部消失。解決辦法是加輔助均衡損失。常見的做法是計(jì)算每個(gè)專家的平均路由概率。計(jì)算每個(gè)專家的平均被選中次數(shù)。讓兩個(gè)分布盡量接近。公式可以表達(dá)為[ L_{balance} \alpha \cdot N \cdot \sum_{i1}^{N} f_i \cdot p_i ]其中 (f_i) 是第 (i) 個(gè)專家被選中的頻率(p_i) 是路由平均概率(N) 是專家數(shù)量(\alpha) 是平衡系數(shù)。6.3 專家內(nèi)部門控SiTU-GLU 與 SwiGLU 的定位MoE 里的門控其實(shí)存在兩層第一層是路由門控決定“選擇誰”。第二層是專家內(nèi)部的前饋門控決定“信息怎么變換”。SwiGLU 和 SiTU-GLU 屬于第二層。專家網(wǎng)絡(luò)本質(zhì)上是一個(gè) FFN而 FFN 的內(nèi)部結(jié)構(gòu)正好可以用 GLU 系列激活函數(shù)來強(qiáng)化。所以一個(gè)典型的 MoE 專家可以寫成Expert(x) OutputProj( GateActivation(InputProj(x)) * ValueProj(x) )其中GateActivation可以是 SiLU對(duì)應(yīng) SwiGLU或 SiTU對(duì)應(yīng) SiTU-GLU。6.4 完整示例Top-K Router SwiGLU 專家下面把路由門控和 SwiGLU 專家組合成一個(gè)可運(yùn)行的簡(jiǎn)化 MoE 層。import torch import torch.nn as nn import torch.nn.functional as F class TopKRouter(nn.Module): 簡(jiǎn)化版 Top-K 路由器 輸入: [B, T, d_model] 輸出: top_idx: [B, T, top_k] 每個(gè) token 選中的專家編號(hào) top_probs: [B, T, top_k] 歸一化后的路由權(quán)重 def __init__(self, d_model, n_experts, top_k2): super().__init__() self.top_k top_k self.router nn.Linear(d_model, n_experts, biasFalse) def forward(self, x): logits self.router(x) # [B, T, n_experts] probs F.softmax(logits, dim-1) top_probs, top_idx torch.topk(probs, self.top_k, dim-1) # 對(duì)選中的概率重新歸一化 top_probs top_probs / top_probs.sum(dim-1, keepdimTrue) return top_idx, top_probs class SimpleMoE(nn.Module): 簡(jiǎn)化版 MoE Layer 每個(gè)專家內(nèi)部使用 SwiGLU 風(fēng)格的前饋網(wǎng)絡(luò)。 def __init__(self, d_model, n_experts, top_k, hidden_dim): super().__init__() self.top_k top_k self.router nn.Linear(d_model, n_experts, biasFalse) # 每個(gè)專家: 兩個(gè)線性層 SiLU 門控 self.experts nn.ModuleList([ nn.Sequential( SwiGLU(d_model, hidden_dim), nn.Linear(hidden_dim, d_model) ) for _ in range(n_experts) ]) def forward(self, x): B, T, D x.shape x_flat x.reshape(-1, D) # [B*T, D] # 路由 logits self.router(x_flat) # [B*T, n_experts] probs F.softmax(logits, dim-1) top_probs, top_idx torch.topk(probs, self.top_k, dim-1) top_probs top_probs / top_probs.sum(dim-1, keepdimTrue) out torch.zeros_like(x_flat) # 逐專家計(jì)算貢獻(xiàn) for e, expert in enumerate(self.experts): # 哪些 token 選中了專家 e mask (top_idx e).any(dim-1) if not mask.any(): continue expert_out expert(x_flat[mask]) # 提取每個(gè) token 對(duì)專家 e 的權(quán)重 e_probs torch.where( top_idx[mask] e, top_probs[mask], torch.zeros_like(top_probs[mask]) ) weight e_probs.sum(dim-1, keepdimTrue) # [n_selected, 1] out[mask] expert_out * weight return out.view(B, T, D) # 測(cè)試 if __name__ __main__: torch.manual_seed(42) moe SimpleMoE( d_model128, n_experts8, top_k2, hidden_dim256 ) x torch.randn(2, 10, 128) y moe(x) print(MoE output:, y.shape) # torch.Size([2, 10, 128])如果你希望訓(xùn)練更穩(wěn)定可以把top_k換成可學(xué)習(xí)的 soft 門控權(quán)重但在稠密場(chǎng)景下 Top-K 的效果已經(jīng)足夠好。7. 常見誤區(qū)與排查清單門控相關(guān)概念多且相似下面整理幾個(gè)高頻誤區(qū)和排查思路。問題現(xiàn)象常見原因解決思路把 Gated MLA 理解成 MHA 后面加個(gè) sigmoidMLA 的核心是低秩 KV 壓縮門控是在 latent 空間調(diào)制 K/V先理解 MLA 的壓縮-解壓過程再看門控插入位置認(rèn)為 KDA 就是 GQAGQA 減少 KV 頭數(shù)量KDA 是把 key 路徑拆成多分支并用門控融合畫出 key 分支結(jié)構(gòu)對(duì)比把 SwiGLU 和 SiLU 當(dāng)成同一個(gè)東西SwiGLU 是門控線性單元SiLU 是激活函數(shù)二者不是同一層面SwiGLU 值分支 x SiLU(門控分支)訓(xùn)練 MoE 時(shí)某個(gè)專家始終不被選中路由初始化不當(dāng)或缺乏負(fù)載均衡損失調(diào)整 router 初始化加入 balance loss門控輸出經(jīng)常飽和梯度消失sigmoid 輸入絕對(duì)值過大檢查 gate_proj 的初始化或使用帶下界的門控變體推理時(shí)門控分支拖慢速度門控在完整 K/V 維度計(jì)算把門控放在 latent 維度或與 KV 解壓算子融合8. 工程實(shí)踐建議與學(xué)習(xí)路線8.1 架構(gòu)設(shè)計(jì)建議在實(shí)際設(shè)計(jì)或復(fù)現(xiàn)帶門控的 LLM 模塊時(shí)有幾點(diǎn)值得留意。首先門控分支的位置直接影響信息選擇效果。Gated MLA 中門控放在 latent 空間更高效但表達(dá)力可能不如放在解壓后的 K/V 空間如果顯存允許可以選擇在兩組位置同時(shí)加門控再用殘差連接做兜底。其次初始化很重要。sigmoid 在輸入為 0 時(shí)輸出 0.5如果門控權(quán)重初始化過大信號(hào)會(huì)接近飽和梯度難以流動(dòng)。實(shí)踐中可以讓gate_proj的輸出初始偏向 0 附近甚至給 bias 設(shè)置一個(gè)負(fù)初始值讓門控初始偏向“關(guān)閉”以穩(wěn)定早期訓(xùn)練。8.2 訓(xùn)練與推理注意事項(xiàng)訓(xùn)練階段要關(guān)注門控是否存在“退化成常數(shù)”的問題。如果訓(xùn)練后期門控輸出始終固定在某個(gè)值附近說明門控沒有學(xué)到有效的信息選擇邏輯可以檢查一下梯度或嘗試給門控分支添加正則。推理階段Gated MLA 和 KDA 門控分支的額外算子會(huì)帶來 kernel launch 開銷。真實(shí)部署時(shí)建議把gate_proj、w_uk、w_uv等線性層合并成一個(gè)大矩陣乘法減少訪存和 kernel 啟停次數(shù)。對(duì) MoE 來說路由門控的數(shù)值精度要留心。在低精度推理時(shí)router logits 的微小波動(dòng)可能導(dǎo)致 top-k 選擇結(jié)果變化進(jìn)而影響生成穩(wěn)定性。必要時(shí)用更高精度計(jì)算 router logits或者對(duì) router logits 做縮放防止 softmax 后概率過于集中。8.3 學(xué)習(xí)路線如果你想系統(tǒng)掌握門控機(jī)制可以按這個(gè)順序走從激活函數(shù)入手實(shí)現(xiàn)并對(duì)比 ReLU、GELU、SiLU、Mish。實(shí)現(xiàn) GLU 和 SwiGLU觀察門控分支對(duì)梯度流動(dòng)的影響。實(shí)現(xiàn)標(biāo)準(zhǔn) MHA再實(shí)現(xiàn) DeepSeek-V2 的 MLA最后加上門控分支。閱讀 LLaMA 系列和 Mistral 的代碼看 SwiGLU 在實(shí)際模型中如何落地。嘗試實(shí)現(xiàn)一個(gè)包含 router 和專家網(wǎng)絡(luò)的 MoE 層加入負(fù)載均衡 loss?；氐?Kimi K3 的公開架構(gòu)資料對(duì)照本文的模塊圖逐層印證。如果條件允許可以下載一個(gè)中等規(guī)模的開源 MoE 模型用推理框架觀察不同專家被激活的頻率配合日志分析路由門控的行為。這會(huì)比單純看論文更有體感。門控機(jī)制是理解現(xiàn)代大模型的一條重要線索。從 LSTM 的遺忘門到 SwiGLU 的門控前饋再到 Gated MLA 和 MoE 路由本質(zhì)上都在做同一件事讓模型自己學(xué)會(huì)信息該如何選擇、何時(shí)放行、以什么比例放行。建議你打開編輯器把文中的 SimplifiedGatedMLA 和 SimpleMoE 各跑一遍然后試著調(diào)整門控位置、初始化和 top-k 大小直觀感受這些改動(dòng)對(duì)訓(xùn)練和輸出的影響。動(dòng)手跑通之后再回去看 Kimi K3 的架構(gòu)分析會(huì)順手很多。