大型语言模型的推理效率与质量之间的平衡是当前研究的热点。现有的条件计算方法，如混合专家（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的有效性。