Meta等最新研究：多token预测，提升大模型推理效率

1,369 0 0

文章摘要

【关键词】 自然语言处理、大语言模型、多token预测、推理效率、模型优化

在自然语言处理领域，大语言模型如GPT-4、Gemini系列和Llama-3等，通常采用下一个token预测的损失函数进行预训练。然而，这种预训练方法需要大量数据且可能降低模型推理效率。为解决这些问题，Meta、巴黎理工大学和巴黎萨克雷大学的研究人员提出了多token预测的训练方法。

该方法的核心是在预训练时要求模型在每个位置上同时预测接下来的n个Token，而不增加预训练时间。研究人员在参数量不同的模型上评估了这一技术，发现130亿参数模型的解决问题能力在HumanEval上提高了12%，在MBPP上提高了17%，同时推理效率也有所提升。

为实现多Token预测，研究人员设计了一种包含共享Transformer主干网络和n个独立输出头网络的模型架构。在推理阶段，这种设计使得模型能够通过单个下一Token预测的输出头进行自回归生成，并利用其他输出头提高推理效率。

此外，为了降低多token预测可能导致的GPU内存使用增加问题，研究人员开发了一种新的前向和后向传播顺序，减少内存中同时存储的梯度数量。通过在每次计算完一个输出头后立即进行后向传播并释放中间数据，确保内存中只存在一个输出头的梯度。

进一步地，结合自推测解码与多token预测，研究人员优化了模型的推理效率。自推测解码允许模型一次性生成多个token，并通过额外输出头并行验证和优化预测，显著减少了生成文本所需的步骤，大幅提高了推理效率并减少了对算力的消耗。实验结果表明，这种优化方法比传统推理效率提升了大约3倍。