N-vium：加速精确生成的混合出口 Transformer
机器学习
2026-05-14 v1 人工智能
摘要
改进自回归 Transformer 推理效率通常意味着降低每个 token 的 FLOPs，通常通过近似方法实现，但会牺牲模型质量。我们引入 N-vium，一种混合出口 Transformer，通过在标准硬件上对深度进行部分并行化，提高每个秒钟的有效 FLOPs 而非最小化每个 token 的计算。N-vium 在多个深度位置附加预测头，并定义下一个 token 的分布为这些出口上学习的混合分布，具有 token 自适应路由。这种表述严格地概括了标准 Transformer，当路由将所有中间头的质量分配为零时，即可恢复标准 Transformer。从混合分布中抽样是精确的，完整的 KV 缓存通过延迟上层计算并与后续 token 批处理来恢复。我们在最高达 15 亿参数规模对 N-vium 进行预训练。我们最大的模型相较于参数和数据匹配的标准 Transformer 实现了 57.9% 的墙钟加速，同时保持原有的困惑度。
引用
@article{arxiv.2605.13190,
title = {N-vium: Mixture-of-Exits Transformer for Accelerated Exact Generation},
author = {Aleksander Lorenc and Frédéric Berdoz and Joël Mathys and Roger Wattenhofer},
journal= {arXiv preprint arXiv:2605.13190},
year = {2026}
}