LayerSkip——新的训练与推理范式：层Dropout，早期退出机制与自投机解码结合
文摘
科技
2024-04-29 13:17
英国
![]()
![]()
LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding
LayerSkip——新的训练与推理范式：层Dropout，早期退出机制与自投机解码结合
本文提出了一个名为LayerSkip的方法，用于加速LLM推理。首先，在训练期间，他们使用层dropout，早期层的dropout率较低，后续层则较高，并且使用早期退出损失（early exit loss）。其次，在推理时，作者展示了这种训练方法如何提高早期层的早期退出精度，而无需为模型添加任何辅助层或模块。而且作者还提出一种新的自投机解码（self-speculative decoding）方案，在早期层退出后，使用剩余层进行验证和修正。该自投机解码方法比其他投机解码方法占用内存更小，并且在Draft（草拟？）和验证阶段共享了计算和激活而受益。作者在不同大小的Llama上实验，包括从头开始的预训练、继续预训练、针对特定数据域的微调以及针对特定任务的微调。这种推理方案，在CNN/DM文档的总结任务上提速2.16倍，在编码任务上提速1.82倍，在TOPv2语义解析任务上为2倍。![]()
![]()
方法的overview见figure1，为了看明白，先来看一个例子，下面figure2a是代码测试集HumanEval里的一个题，下面生成的地方每个token有不同程度的深浅，这些深浅代表着模型最先预测出这个词的那一层的层数，而figure2b是每个层在每次预测时所代表的预测。（这些token是将每层的输出嵌入投影到LM head上得到的，该head包括模型的最后一层归一化层和线性层），从figure2可以看到一些现象：1.早期层的token预测似乎不重要，因为它们对应于投影在模型嵌入层权重上的前一个token，这些权重与LM头部的权重不同，而在后续层，token预测趋于最终预测；2.并非总是需要所有层来预测正确token，大多数时候，最终的token预测是在结束前的较少层数内被预测出来的；3.中间层有时会犹豫并“改变主意”，如对于Token 05，模型在第07层就预测为“range”，但在第22层到第26层之间改变了主意，最终又重新定为“range”。从上一段的一些观察中来说，即使有一个完美的预测器知道合适预测达到饱和并及早退出，对于figure2的例子，平均而言一个token也需要模型32层中的23.45层，最多也只能节省26%的计算量。因此，需要设计方法使模型预测每个词使使用更少的层且减少由于犹豫带来的算力消耗，通常来说深度学习模型并不倾向于早期预测其最终输出，而是将计算分散到所有层（比如token02花了所有32层才预测出来for）。作者希望模型不那么依赖后期层，并且只对较难的token使用后期层，更多依赖早期层而非后期层，他们提出在训练中跳过某些层，既层dropout，并且为达上述目的他们对后期层用更高的dropout率，早期用更低的。此外还有一个改动，LM的head原先被训练于从最后的transfomer层unembed嵌入，他们为了让模型更好地理解早期层的embedding，在训练时加了一个损失函数并且对所有transformer层共享同一个head（好处是训练更快、训练与推理时内存消耗少且易于部署），模型的结构因此如figure3所示。![]()
还有一点，如果在推理中过早退出可能会导致准确率下降，因此需要一种方法验证早期预测的准确性，并通过执行剩余层来纠正它，作者提出的自投机解码先用早退机制自回归地生成token（self-drafting阶段），再用剩余的层并行验证一组token并修正（self-verification），图例见figure5。方法的各种公式细节都在论文里，很直观，我就不来贴了。![]()
![]()
![]()
![]()
撰文：戴剑波；编辑：戴剑波
未经本公众号授权不得转载，欢迎转发。
