CausalGate：基於因果重要性蒸餾的Transformer模組剪枝
現有的大語言模型自適應推理方法通常依賴隱藏狀態相似性或啟用幅度等啟發式規則來丟棄冗餘模組，但這些基於相關性的指標難以捕捉對語義準確性至關重要的非線性結構計算。CausalGate提出了一種干預引導框架，透過隔離子層並衡量輸出歸零對最終logit分佈的KL散度影響，建立結構重要性層級，再將其蒸餾為靜態輕量級標量門，實現零執行時開銷的高效推理。在多個模型和基準上的實驗表明，CausalGate顯著優於現有動態路由方法。
近年來，大語言模型（LLM）的規模持續增長，推動了自適應推理方法的發展，這些方法旨在透過動態跳過冗餘模組來降低計算成本。然而，現有方法大多依賴隱藏狀態相似性或啟用幅度等觀測性啟發式規則，這些基於相關性的指標往往無法捕捉到對語義準確性至關重要的非線性結構計算。針對這一侷限，研究團隊提出了CausalGate，一種基於因果乾預的框架。
在校準階段，CausalGate逐一隔離Transformer中的Attention和MLP子層，將其輸出置零，並透過最終logit分佈的Kullback-Leibler散度精確量化對語義的損害。具體來說，對於每個子層，CausalGate將其輸出置為零，然後計算修改後模型在驗證樣本上的最終logit分佈與原始分佈之間的KL散度。該散度值越大，說明該子層對語義輸出越重要。透過遍歷所有子層，得到重要性得分向量。這一過程建立了每個子層的結構性重要性層級。
為了消除執行時路由開銷，CausalGate採用指數移動平均（EMA）平滑目標結合可微分的成對排序損失，將重要性層級蒸餾為一組全域性靜態的輕量級標量門。EMA用於平滑訓練過程中的重要性估計，排序損失確保門控值反映子層之間的相對重要性順序。最終學習到的門控是標量值，位於0到1之間，推理時透過閾值決定是否跳過該子層。這樣，推理時無需任何動態決策，即可根據預先學習的門值選擇保留或跳過模組。
實驗在TinyLlama-1.1B、Qwen2.5-3B和Llama-3.1-8B上進行，涵蓋語言建模和常識推理基準。在TinyLlama-1.1B上，CausalGate在保持98%以上原始效能的同時，減少了50%的FLOPs；在Qwen2.5-3B和Llama-3.1-8B上，也展現出類似的效率提升。與SkipLayer、ActivationParing等基線相比，CausalGate在計算節省和硬體延遲降低方面一致優於主流動態路由和層跳過方法，且不增加執行時開銷。由於門控是靜態的，實際硬體延遲直接對應於計算減少，而不需要額外的路由決策開銷。
該工作已被會議接收，提供了高效Transformer推理的新思路。論文由Kiran Nair等人提交，於2026年7月21日釋出在arXiv上，目前正在會議審稿中。這項研究的價值在於，它從因果角度重新定義了模組重要性，避免了相關性的謬誤。未來，CausalGate可以擴充套件到其他架構和任務中，為高效推理提供通用框架。