Tandem Transformers for Inference Efficient LLMs

📄 arXiv: 2402.08644v4 📥 PDF

作者: Aishwarya P S, Pranav Ajit Nair, Yashas Samaga, Toby Boyd, Sanjiv Kumar, Prateek Jain, Praneeth Netrapalli

分类: cs.AI, cs.CL

发布日期: 2024-02-13 (更新: 2024-10-20)

💡 一句话要点

提出串联变换器以解决大语言模型推理效率问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 推理效率 自回归模型 块模式 投机解码 模型组合 性能提升

📋 核心要点

  1. 现有的大语言模型在推理时由于自回归特性，导致生成速度较慢，影响了实际应用效率。
  2. 本文提出的串联变换器架构结合了小型自回归模型和大型块模式模型，利用大型模型的丰富表示来提升小模型的预测准确性。
  3. 实验结果表明，串联模型在下一个标记预测准确性上提高了3.3%，并在SPEED框架下实现了1.14倍的速度提升。

📝 摘要（中文）

传统的大语言模型（LLMs）由于其自回归特性，推理速度受到限制，因为生成的标记是顺序进行的。尽管存在投机性和并行解码技术来缓解这一问题，但它们面临着依赖较小模型生成或未能充分利用基础LLM表示的局限性。本文提出了一种新颖的架构——串联变换器，结合了小型自回归模型和以块模式运行的大型模型。通过让小模型关注大型模型的丰富表示，显著提高了预测准确性。在PaLM2预训练数据集上，PaLM2-Bison和PaLM2-Gecko的串联在下一个标记预测准确性上比单独的PaLM2-Gecko提高了3.3%，并且在下游性能相当的情况下，相比于PaLM2-Otter模型实现了1.16倍的速度提升。我们进一步将串联模型纳入投机解码（SPEED）框架，使得PaLM2-Bison和PaLM2-Gecko的串联在保持相同下游任务准确性的同时，速度提升约为1.14倍。

🔬 方法详解

问题定义：本文旨在解决传统大语言模型在推理过程中的速度瓶颈，现有方法往往依赖较小模型或未能充分利用大型模型的表示能力，导致效率低下。

核心思路：提出的串联变换器架构通过结合小型自回归模型与大型块模式模型，利用后者的丰富表示来增强前者的预测能力，从而提高推理速度和准确性。

技术框架：整体架构包括两个主要模块：小型自回归模型负责初步生成标记，而大型模型则以块模式处理多个标记，确保生成的标记具有更高的准确性。

关键创新：最重要的创新在于将小型和大型模型的优势结合，通过关注大型模型的表示来提升小模型的性能，这一设计与传统方法的单一模型生成方式形成了鲜明对比。

关键设计：在模型设计中，采用了特定的注意力机制，使小模型能够有效利用大型模型的输出，同时在损失函数的设置上进行了优化，以确保模型在训练过程中的稳定性和准确性。

📊 实验亮点

实验结果显示，串联模型在下一个标记预测准确性上比单独的PaLM2-Gecko提高了3.3%，并在下游性能相当的情况下，相比于PaLM2-Otter模型实现了1.16倍的速度提升。此外，在SPEED框架下，串联模型的速度提升约为1.14倍，且保持了相同的下游任务准确性。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和文本生成等。通过提高推理效率，串联变换器可以在实时应用中提供更快的响应时间，提升用户体验，未来可能在大规模语言模型的实际应用中发挥重要作用。

📄 摘要（原文）

The autoregressive nature of conventional large language models (LLMs) inherently limits inference speed, as tokens are generated sequentially. While speculative and parallel decoding techniques attempt to mitigate this, they face limitations: either relying on less accurate smaller models for generation or failing to fully leverage the base LLM's representations. We introduce a novel architecture, Tandem transformers, to address these issues. This architecture uniquely combines (1) a small autoregressive model and (2) a large model operating in block mode (processing multiple tokens simultaneously). The small model's predictive accuracy is substantially enhanced by granting it attention to the large model's richer representations. On the PaLM2 pretraining dataset, a tandem of PaLM2-Bison and PaLM2-Gecko demonstrates a 3.3% improvement in next-token prediction accuracy over a standalone PaLM2-Gecko, offering a 1.16x speedup compared to a PaLM2-Otter model with comparable downstream performance. We further incorporate the tandem model within the speculative decoding (SPEED) framework where the large model validates tokens from the small model. This ensures that the Tandem of PaLM2-Bison and PaLM2-Gecko achieves substantial speedup (around 1.14x faster than using vanilla PaLM2-Gecko in SPEED) while maintaining identical downstream task accuracy.