LayerSkip：加速大模型推理的端到端解决方案

 大模型（LLMs）在多种应用中表现出色，但其高昂的计算和内存需求导致部署成本昂贵，尤其是在GPU服务器上。现有加速方案在部署到普通GPU时往往会导致准确性显著下降，而将大模型（LLMs）进一步加速以部署到移动设备或边缘设备，仍然是自然语言处理领域的一个挑战。LayerSkip提出了一种创新的方法，通过在训练期间应用层dropout和提前退出损失，在推理时采用自我推测解码技术，有效地提高了早期层的退出准确率。这种方法不仅避免了在普通GPU上部署时的准确性损失，还为将LLMs部署到移动或边缘设备提供了可能性，实现了高达2.16倍的速度提升，同时保持了模型性能。

(a) 展示了一个来自 HumanEval 数据集的编程提示示例，以及 Llama1 7B 模型基于这个提示生成的文本。 (b) 展示了 Llama1 7B 模型中每个层级对生成 token 的预测过程，其中不同颜色表示模型在不同层级预测出特定 token 的时间点。在大模型（LLMs）中，并非总是需要所有的层来预测正确的 token。通过分析模型在不同层级的预测行为，研究者发现在较早的层级就可以做出准确的预测，而且模型在中间层级有时会犹豫不决或“改变主意”。这表明有机会通过在推理过程中提前退出（early exit）来提高效率，同时保持合理的准确性。

解决方案

LayerSkip的提出建立在dropout、层dropout、提前退出和推测解码等现有研究的基础上。通过结合这些技术，LayerSkip旨在提高LLMs的推理速度，同时保持或提高准确性。LayerSkip的解决方案包括三个阶段：使用层dropout和提前退出损失的训练、使用提前退出的推理，以及使用推测解码的验证和纠正。

 LayerSkip 的三个组成部分：训练时使用层dropout和早期退出损失、推理时使用早期退出、使用自我推测解码进行验证和校正。说明了如何通过这些组件来提高模型的推理速度和准确性。

•  使用层dropout和提前退出损失的训练

  训练阶段的一个关键创新是引入了层dropout和提前退出损失。这种结合的训练策略旨在提高模型在推理时早期层的退出准确性，从而使得模型能够更快地进行预测而不必总是利用所有层。

  层dropout是一种正则化技术，它随机地在训练过程中跳过网络中的某些层。LayerSkip在训练时不是均匀地应用dropout，而是采用一种分层的方法，即对于网络中较早的层采用较低的dropout率，而对于较晚的层采用较高的dropout率。这样的设计使得模型在训练时就逐渐习惯于在没有所有层的情况下进行预测，从而在推理时可以更早地退出。

  与层dropout相结合的是提前退出损失，这是一种在训练中增加的额外损失函数。在传统的语言模型训练中，通常只在模型的最后一层，也就是语言模型头，来计算预测的损失。然而，在LayerSkip中，每个transformer层的输出都会被传递到语言模型头，并且计算一个额外的损失。这样模型就被训练为在每个层都能够生成准确的预测，而不仅仅是依赖于最后一层。

  通过这种方式，模型学会了在不同层上进行有效的预测，并且能够在不需要完整网络参与的情况下生成高质量的输出。这种训练策略不仅提高了模型在推理时的效率，还保持了预测的准确性，因为它确保了模型的每个部分都能够独立地贡献于最终的预测结果。这种训练方法的结果是，模型能够在推理时更快地收敛到正确的预测，同时减少了计算资源的需求。

• 使用提前退出的推理

  在LayerSkip框架下，推理阶段的优化是通过提前退出机制实现的。这种方法的核心思想是在自回归解码过程中，不必要地运行模型的所有层，而是在达到一定的层数后直接跳转到语言模型（LM）头进行预测。这样做的目的是为了减少每次生成token所需的计算量，从而加快推理速度。

  具体来说，在推理时，模型会执行前E层的transformer操作，这里的E是一个小于模型总层数L的值。通过这种方式，模型可以在经过较少层的处理后，就尝试生成输出。一旦到达第E层，模型不会继续向更深层传递信息，而是直接将E层的输出传递给LM头，以生成下一个token的预测。

  这种提前退出的策略允许模型在较早的阶段就生成token，而不是等待所有层都完成处理。这样的设计显著减少了每次token生成所需的计算步骤，因为深层的transformer层往往需要更多的计算资源。通过优化推理路径，LayerSkip能够在保持预测准确性的同时，显著提高模型的运行效率。

  提前退出的策略还意味着模型的不同层可以独立地对输出做出贡献。这不仅提高了模型的灵活性，还使得模型能够更好地适应不同的任务和数据特性。例如，在某些任务中，可能较早的层已经能够捕捉到足够的信息来生成准确的预测，而在其他任务中，可能需要更深层的信息来做出精确的判断。

• 使用推测解码的验证和纠正

  LayerSkip中的自我推测解码算法是一种创新的机制，它旨在进一步提升模型的推理速度，同时确保生成的token保持高质量。这个算法的核心思想是分两个阶段进行token的生成和校验：首先，模型使用前E层生成初步的draft tokens；然后，利用模型剩余的L-E层对这些draft tokens进行验证和纠正。

  在自我推测解码的第一阶段，模型通过前E层的transformer生成draft tokens。这些tokens是模型对输入序列的初步预测，它们是通过前E层的信息积累和处理得到的。这个阶段的目的是快速生成一组候选的token，而不是等待整个模型的所有层都完成计算。

  算法进入第二阶段，即验证和纠正过程。在这个阶段，模型的剩余L-E层被用来对第一阶段生成的draft tokens进行细致的检查。这个过程相当于一个后置的审核机制，它能够利用模型更深层的知识来评估和改进初步预测的结果。如果发现draft tokens中有不准确或需要改进的地方，剩余的层将对其进行调整，以确保最终输出的token尽可能接近最优解。

  通过这种自我推测和后续纠正的机制，LayerSkip能够在减少计算量的同时，保持语言模型生成token的准确性。这种方法的优势在于，它允许模型在保持高效率的同时，还能够利用整个网络的深度来提升预测的质量。此外，由于验证和纠正过程是对draft tokens的集体评估，因此相比于逐个token的自回归生成，这种方法可以更高效地利用计算资源，进一步提高推理速度。

  

  自回归解码、推测性解码和自我推测解码之间的成本和效率，展示了自我推测解码如何通过重用草稿阶段的隐藏状态来减少内存和延迟。

实验

为了全面评估LayerSkip的性能，研究者设计了一系列实验，涵盖了不同规模的Llama模型，并采用了多样化的训练策略。实验包括了从零开始的预训练，这意味着研究者从随机初始化的模型参数开始，逐步训练模型以学习数据中的模式。实验还包括了持续预训练，即在已有预训练模型的基础上，继续用新数据进行训练以进一步提升模型性能。为了测试LayerSkip在特定领域的适用性，研究者还进行了特定数据领域的微调实验，这涉及到将预训练得到的模型在某一特定类型的数据上进行额外的训练，以期获得更好的领域适应性。研究者还对模型进行了特定任务的微调，这是为了让模型能够针对某一具体任务，如文本摘要或语义解析，达到最佳的性能。

研究者们首先在不同大小的Llama模型上进行了实验，包括7B和13B参数规模的模型，以验证LayerSkip在持续预训练和从头开始预训练情况下的有效性。实验任务涵盖了文本摘要、代码生成、阅读理解和常识推理等多个领域，使用了CNN/DM、XSUM、HumanEval、MBPP、BoolQ、PIQA、SIQA、HellaSwag、WinoGrande、ARC、RACE和MMLU等数据集。

结果显示，LayerSkip能够在保持模型最后层准确性的同时，显著提高早期层的退出准确率。例如，在Llama2 7B和Llama2 13B模型上进行的持续预训练实验中，LayerSkip配置（结合了层dropout和提前退出损失）在多个任务上都展现出了性能提升。具体来说，在Wikipedia测试集上，使用LayerSkip的模型困惑度从基线的4.32降低到4.3，表明模型生成文本的流畅性和准确性有所提高。在HumanEval编程任务上，执行正确率从基线的13.4%提高到了15.9%，显示了LayerSkip在特定任务上的有效性。

进一步的实验在从头开始的预训练设置中进行，研究者们发现，即使是在较小的数据集上，LayerSkip也能带来性能上的提升。例如，在Llama2 1.5B模型上，使用LayerSkip的模型在不同任务上的速度提升范围从1.34倍到2.16倍不等，这一结果证明了LayerSkip在加速模型推理方面的潜力。

在特定领域的微调和特定任务的微调实验中，LayerSkip同样展现出了其有效性。例如，在TOPv2语义解析任务上，经过LayerSkip训练的模型在保持准确性的同时，实现了高达2倍的速度提升。

研究者们还探讨了不同退出层对模型性能的影响。在Llama2 7B模型上，当在第16层退出时，模型在BoolQ任务上的准确率从基线的77.4%提高到了77.8%，这表明即使在模型的中间层退出，LayerSkip也能保持较高的性能。

实验结果表明LayerSkip是一个强大的工具，能够在不同规模的模型和多样化的任务中提高大型语言模型的推理速度和效率，同时保持或提升模型的性能。这些发现为未来在资源受限的环境中部署大型语言模型提供了新的思路和方法。

使用 LayerSkip 在持续预训练的 Llama2 7B 和 Llama2 13B 模型上的评估结果，包括不同层级的准确率和困惑度。

LayerSkip技术虽然在加速大型语言模型（LLMs）的推理方面取得了显著成果，但仍存在一些局限性。自我推测解码方案要求模型进行特定的微调或预训练，且需要对新引入的超参数进行细致调整，以避免影响模型最后层的准确性。这些要求可能会增加模型训练和调整的复杂性。尽管如此，LayerSkip的优势仍然明显。它通过结合层dropout、提前退出损失和自我推测解码，显著提升了模型的推理速度，同时保持了模型性能。这一成果不仅为大型语言模型的部署提供了新的解决方案，也为未来在资源受限的环境中应用LLMs开辟了新的可能性。

研究者们对LayerSkip的潜力持乐观态度，并期望它能够激发更广泛的研究和应用。通过进一步的研究和优化，LayerSkip有望克服现有局限，为自然语言处理领域带来更深远的影响。随着技术的进步和模型的迭代，LayerSkip有潜力成为推动大型语言模型在更广泛场景中高效部署的关键技术。未来，我们期待LayerSkip能够在更大规模的数据集上进行训练，支持更复杂的任务，并在各种应用中实现更高的性能和效率。这不仅将推动自然语言处理技术的发展，也将为人工智能领域带来新的突破。

论文链接：https://arxiv.org/abs/2404.16710