大型語言模型的推理效率與質量之間的平衡是當前研究的熱點。現有的條件計算方法，如混合專家（MoE）、混合深度（MoD）和KV快取量化，各自獨立地最佳化了前饋網路、Transformer塊跳過和注意力記憶體壓縮。然而，這些決策實際上是相互關聯的：一個需要全注意力的罕見令牌可能也需要高精度快取，無論哪個專家處理它。針對這一問題，Andrii Balashov和Olena Ponomarova提出了TriRoute，一個輕量級的統一控制器，能夠為每個令牌在每一層聯合決定注意力模式（跳過、區域性或全域性）、一組稀疏的FFN專家（包括空專家以恢復MoD）以及KV快取位寬。
TriRoute的訓練採用了一種異構鬆弛方法：對於注意力模式和快取位寬的分類決策，使用Gumbel-Softmax與直通估計（straight-through estimator）來近似離散取樣；對於專家選擇，則使用負載均衡的top-k門控，確保每個專家接收的令牌數大致相等。整個訓練過程在拉格朗日預算約束下進行，將平均計算和記憶體成本轉化為可控的旋鈕——透過調整拉格朗日乘子，可以平滑地在效能與效率之間權衡。研究團隊發現，在聯合訓練中會出現跨軸路由崩潰級聯問題：即一個軸的決策崩潰（例如所有令牌選擇同一專家）會迅速傳播到其他軸，導致整個系統退化。為了解決這一問題，他們引入了逐軸歸一化和耦合感知平衡損失，前者確保每個軸的輸出分佈保持多樣性，後者則懲罰不同軸之間決策的過度相關性。
在實驗部分，研究人員在一系列解碼器-only模型上進行了測試，模型規模從1.6億引數到13億引數不等，所有模型均以計算最優的令牌數進行訓練。在相同的推理FLOPs和記憶體預算下，TriRoute在多個基準測試中Pareto優於最佳獨立MoD+MoE+KV量化組合。更重要的是，TriRoute更好地保留了尾部情況的魯棒性：在稀有實體、程式碼和算術任務上，純粹因困惑度最佳化而導致的效能下降得到了顯著緩解。事後分析揭示了可解釋的路由結構：控制器往往將全注意力和高精度快取分配給句首位置、稀有子詞和命名實體，而對功能詞則採用廉價的路由策略（如跳過注意力或使用低精度快取）。
這項工作表明，聯合最佳化注意力、專家和快取決策可以顯著提升語言模型在資源受限場景下的表現，併為未來更高效的自適應推理系統奠定了基礎。研究團隊已公開了論文和部分程式碼，期待在更大規模模型和多模態場景中進一步驗證TriRoute的有效性。