论文
N-vium:用于加速精确生成的混合退出 Transformer
N-vium: Mixture-of-Exits Transformer for Accelerated Exact Generation
摘要
提高自回归 Transformer 的推理效率通常意味着减少每个词元的 FLOP,通常通过降低模型质量的近似来实现。我们引入了 N-vium,一种混合出口 Transformer,它在标准硬件上部分并行化深度计算,增加每秒有效的 FLOP 数,而不是最小化每个词元的计算量。 N-vium 在多个深度附加预测头,并将下一个词元分布定义为这些出口上的学习混合物,并具有词元自适应路由。该公式严格推广了标准 Transformer,当布线将零质量分配给所有中间头时,该公式将被准确恢复。从混合物中采样是准确的,并且通过推迟上层计算并将其与后面的词元进行批处理来恢复完整的 KV 缓存。我们以高达 1.5B 参数的规模预训练 N-vium。我们最大的模型比参数和数据匹配的标准 Transformer 实现了 57.9% 的挂钟加速,且没有任何困惑成本。