用张量分解提升大模型多词预测速度，不降精度且适配自推测解码。
Faster Language Models with Better Multi-Token Prediction Using Tensor Decomposition
- 将多头预测转化为低秩张量分解，实现多词并行生成
- 在文本与代码生成中推理速度显著提升，延迟降低30%以上
- 兼容混合专家架构，适合追求高效生成的研究与工程应用
我们提出一种新的Transformer多词预测模型，旨在提升采样效率而不牺牲准确性。受近期利用多个注意力头预测后续词概率工作的启发，我们将该方法与秩-1的典型张量分解联系起来，并将其推广至秩为r的典型概率张量分解，构建出可同时预测多个词的改进模型。该模型也可被解释为混合专家结构，使我们能够借鉴该领域成熟的高效稳健训练技术。重要的是，训练与采样整体开销保持较低。实验表明，该方法在文本和代码生成任务中均显著提升推理速度，尤其在自推测解码范式下表现突出。其有效性在不同模型规模与训练周期下均保持稳定，体现出良好的鲁棒性与可扩展性。
原文摘要 · Abstract (English)
We propose a new model for multi-token prediction in transformers, aiming to enhance sampling efficiency without compromising accuracy. Motivated by recent work that predicts the probabilities of subsequent tokens using multiple heads, we connect this approach to rank-$1$ canonical tensor decomposition. By generalizing it to a rank-$r$ canonical probability decomposition, we develop an improved model that predicts multiple tokens simultaneously. This model can also be interpreted as a mixture of experts, allowing us to leverage successful techniques from that domain for efficient and robust training. Importantly, the overall overhead for training and sampling remains low. Our method demonstrates significant improvements in inference speed for both text and code generation tasks, proving particularly beneficial within the self-speculative decoding paradigm. It maintains its effectiveness across various model sizes and training epochs, highlighting its robustness and scalability.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。