零 阅读基础与可能帮助文档
一 论文摘要与加速背景详解
Training-Time Test (训练时测试)，顾名思义，指的是在模型训练过程中进行测试的技术。它与传统的训练后测试 (Post-Training Test) 相对应，后者是在模型训练完成后，使用独立的测试数据集评估模型的性能。
通过在训练过程中引入测试，能够更早地发现问题，从而更好地指导训练过程，提高模型性能和效率。它试图将测试环节融入到模型的学习循环中，而不是仅仅作为最后验证的手段。
EAGLE-1和EAGLE-2 在特征级别进行自回归解码，复用目标模型的最后一层特征，从而比其他普通的投机采样算法获得更好的结果，但这样的结构也存在一些不足，即数据量达到一定程度后，继续扩大数据规模对 EAGLE 的改进有限。
LLM 社区的一个日益增长的趋势是扩大训练数据规模，以提高模型智能，而无需增加推理成本。 由于EAGLE 的特征预测约束限制，扩大数据规模对EAGLE-1和EAGLE-2的增益减弱很多。
因此在EAGLE-3中放弃了特征预测，转而采用直接的 Token 预测，并用一种名为训练时测试 (training-time test) 的技术，通过多层特征融合取代了对顶层特征的依赖。 这些改进显著提高了性能，并使草稿模型能够充分受益于扩大训练数据规模。
在一系列实验上的结果表明，EAGLE-3 的加速比高达 6.5 倍，比 EAGLE-2 提高了约 1.4 倍。
二 论文引言与基本介绍详解
最近，测试时扩展（test-time scaling up）受到了极大的关注，像 ChatGPT o1 和 DeepSeek-R1 这样的模型在回复之前会进行深思熟虑的推理（即额外的深入思考过程），从而以更长的推理时间为代价，进一步提升了 LLM 的复杂推理能力。
然而，这些模型通常需要冗长的推理过程，并且响应时间的增加严重影响了用户满意度。这些推理模型显著增加了推理成本在整个 LLM 流程中的比例，从而促使研究人员探索更便宜、更快速的推理优化方法。
使用投机采样方法可以通过并行化生成过程来降低LLM的延迟，从而加速冗长的模型推理过程。应用投机采样时，一般会增加一个辅助模型，我们称这个辅助模型为草稿模型(Draft Model)，而需要推理加速的模型称为原始模型或基础模型(Base Model)，一般草稿模型远小于基础模型，草稿模型尺寸大概为基础模型的1/30-1/10。
当模型进行一次前向推理时，会先让草稿模型串行快速生成草稿token，然后再使用基础模型并行地验证这些token的正确性。这使得在单个前向传播中基础模型可以并行地生成多个token，从而显著降低推理延迟。
不同于一般的投机采样模型，EAGLE 使用了目标模型的最后一层特征（LM head之前的特征）进行辅助预测。它训练草稿模型以自回归方式预测下一个特征，然后使用目标模型的LM head来获得草稿token。通过利用来自目标模型的丰富信息，与原始的投机采样相比，EAGLE实现了显著更好的加速效果。这种方法的先进性也导致后续方法如HASS 和 Falcon均采用了使用当前特征序列预测下一个特征的方式来增强模型预测能力。
最近的LLM越来越多地依赖于更大的训练数据集来获得更好的性能。例如，LLaMA系列模型，大小为7B (8B)的模型，分别为LLaMA 1 、LLaMA 2 和LLaMA 3 使用了1T、2T和15T tokens的训练数据，从而在保持模型架构和推理成本基本不变的情况下，在各种指标上都取得了显著的改进。类似地，我们的目标是通过增加EAGLE的训练数据来提高其接受率和加速比。
遗憾的是，作者观察到EAGLE从额外训练数据中获得的收益是有限的。作者尝试分析这种现象背后的原因：如图所示，EAGLE在特征层面执行自回归预测，预测下一个特征，然后将该特征馈入目标模型的LM head以获得token分布。EAGLE的损失函数由两部分组成：特征预测损失loss_features和词元预测损失loss_tokens，其中loss_features是草稿模型多步预测能力大幅增强的关键。（看不懂的可以参考本文顶部的链接：EAGLE-1和EAGLE-2论文讲解）
然而，由于token预测是投机采样的最终目标，特征预测可以被看作是该过程的一个额外约束，这限制了草稿模型的表达能力，使其难以从增加的数据中获益。在移除特征约束并扩大训练数据后（如图所示），第一个草稿token的接受率0-α 显著提高;与此同时，第二个草稿token的接受率1-α却断崖式下降。换句话说，特征预测是增强了后续的token预测能力，但是同时也限制了数据规模增益的能力。
针对这个问题，作者通过将后续token预测步骤纳入训练过程来解决这个问题（效果如上图所示，过程如下图所示）。使用这种方法，增加训练数据的好处变得更加明显。作者将这种技术命名为训练时测试（training-time test）。简单来说，就是训练的时候，也考虑将第二个token，第三个token的预测损失也考虑在内，从而增强多步预测能力。
EAGLE和诸如Medusa之类的推测采样方法会重用目标模型的顶层特征，特别是紧邻LM head之前的特征。对于具有满秩权重矩阵的LM head，与下一个token的logits相对应的顶层特征是唯一的，确保了这些特征中包含的信息与下一个token的logits直接对齐。然而，仅仅基于顶层特征（其本质上仅限于下一个token）来预测下下个token构成了一个重大挑战。幸运的是，如上所述的训练时测试技术能够使用中间层的特征，而不是仅仅依赖顶层，因为在训练期间已经移除了特征预测损失lfea。
三 论文方法与原理细节详解
一段核心的模型修改代码：三个中间隐藏向量的代码实现，见modeling_llama_kv.py里面class LlamaModel(LlamaPreTrainedModel)的forward方法
for idx, decoder_layer in enumerate(self.layers):
# 隐藏向量关键：低层隐藏向量+中层隐藏向量+高层隐藏向量
if idx==len(self.layers)-3 or idx==len(self.layers)//2 or idx==2:
all_hidden_states += (hidden_states,)
3.1 EAGLE-3推理过程详解
跟其他投机采样方法一样，EAGLE-3的推理过程也大致分为两个阶段：草稿阶段(Drafting Stage)和验证阶段(Verification Stage)。EAGLE-3与之前的EAGLE版本不同之处在于草稿阶段的推理流程发生了一定变化，下面详细阐述草稿阶段的变化。
如下图的左边部分流程，首先会获取三个中间隐藏向量，核心代码如下：
for idx, decoder_layer in enumerate(self.layers):
# 隐藏向量关键：低层隐藏向量+中层隐藏向量+高层隐藏向量
if idx==len(self.layers)-3 or idx==len(self.layers)//2 or idx==2:
all_hidden_states += (hidden_states,)
这三个隐藏向量分别是低层隐藏向量l，中层隐藏向量m，高层隐藏向量h。获取到这三个隐藏向量后，见下图的中间部分，会进行拼接，然后经过FC_Layer层变为隐藏向量g，形状变化过程为：(3, hidden_size) --> (1, hidden_size*3) --> (1, hidden_size)
见下图的右边部分，之后这个隐藏向量g就和之前的EAGLE版本一样运算了。
3.2 EAGLE-3训练过程详解
在EAGLE-1和EAGLE-2中，草稿模型的输入是原始模型的最后一层特征： f 1 , f 2 , . . . , f t f_{1} , f_{2}, ..., f_{t} f1,f2,...,ft
但在EAGLE-3的草稿模型里面，存在两种输入： g 1 , g 2 , . . . , g t g_{1} , g_{2}, ..., g_{t} g1,g2,...,gt和 a t + 1 , a t + 2 , . . . , a t + j a_{t+1} , a_{t+2}, ..., a_{t+j} at+1,at+2,...,at+j。前者主要在启动阶段prefill输入，后者主要在解码阶段extends输入。要使模型适应这种输入模式，需要特地训练才行。
尽管模型输入的数据都是特征级的向量，但实际训练数据却是序列级别的词元，即输入prompt后得到token_ids，然后由token_ids得到vectors。
以输入"“How can I"为例，具体过程如图所示，由于因果关系，前面的词不能看到后面的词，因此attention mask得到的是一个下三角矩阵：
假如这三个位置的草稿模型输出分别是"are", “we”, “do”，那么注意力掩码会变成下面这个样子：
假如草稿模型进一步的输出为"you", “help”, “it”，那么注意力掩码会变成下面这个样子：
特征预测限制了EAGLE大规模数据受益的能力，因此EAGLE-3移除了特征预测（即只做token预测），转而用训练时多步测试来代替，如下图展示了区别：
该图上面部分展示的是之前EAGLE版本同时做特征预测和词元预测的训练和推理流程图，中间部分展示的是移除了特征预测后，第二步词元预测误差会比较大，下面部分展示了训练时多步测试可以解决token预测多步误差大的问题。
在训练上的细节可以参考HASS这篇论文：HASS
四 论文实验与效果对比详解
