通过多出口机制实现推理加速，不牺牲生成质量。
N-vium: Mixture-of-Exits Transformer for Accelerated Exact Generation
- 在不同深度设置预测头，用可学习混合方式选择输出路径。
- 1.5B参数模型相比标准Transformer快57.9%，困惑度不变。
- 适合需要高效精确生成的场景，如长文本生成任务。
提升自回归Transformer的推理效率通常依赖降低每词的浮点运算量，但常以模型质量下降为代价。我们提出N-vium，一种基于多出口的Transformer结构，在标准硬件上部分并行化深度方向的计算，从而提升每秒有效浮点运算量，而非最小化每词计算量。N-vium在多个深度层附加预测头，并将下一个词的概率分布定义为这些出口的可学习混合，采用词级自适应路由。该形式严格推广了标准Transformer，当路由对所有中间头赋零权重时即完全还原标准结构。从混合分布采样为精确生成，通过延迟高层计算并与其他词批量处理，可完整恢复键值缓存。我们在最大达1.5B参数规模下预训练N-vium。最大模型在与标准Transformer参数和数据匹配条件下，实现了57.9%的墙钟速度提升，且困惑度无损失。
原文摘要 · Abstract (English)
Improving the inference efficiency of autoregressive transformers typically means reducing FLOPs per token, usually through approximations that degrade model quality. We introduce N-vium, a mixture-of-exits transformer that partially parallelizes computation across depth on standard hardware, increasing effective FLOPs per second rather than minimizing compute per token. N-vium attaches prediction heads at multiple depths and defines the next-token distribution as a learned mixture over these exits, with token-adaptive routing. This formulation strictly generalizes the standard transformer, which is recovered exactly when routing assigns zero mass to all intermediate heads. Sampling from the mixture is exact, and complete KV caches are recovered by deferring the upper-layer computation and batching it with later tokens. We pretrain N-vium at scales up to 1.5B parameters. Our largest model reaches 57.9% wall-clock speedup over a parameter- and data-matched standard transformer at no perplexity cost.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。