大型語言模型的推理效率與質量之間的平衡是當前研究的熱點。現有的條件計算方法，如混合專家（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的有效性。