HASS(ICLR 2025)精读¶
核心定位¶
HASS 在 EAGLE-2 上只改训练、不加推理成本,修复两种不一致:全词表回归不直接服务候选 recall;训练时一直看到 target feature,而推理深层只能看到自己的有误 feature。Top-K ranking distillation + 多步 feature rollout context alignment 将 acceptance length 提高 8–16%,最终比 EAGLE-2 快 8–20%。
1. Harmonized objective¶
取 target probability 最高的 K 个 token 集 Ω̂,增加 L_topK=-Σ_{x∈Ω̂} q_target(x) log p_draft(x),让有限 drafter 容量优先覆盖 decoding 真正可能选择的 token,而不是长尾 vocabulary。它与 EAGLE feature regression/CE 联合训练。默认 K=10、weight=1;消融实际在 K≈5、w≈.5 峰值,说明 paper main setting未必是每模型最优。
比较 Top-P、normalized top-K、bidirectional、Recall@k surrogate、BiLD ranking:T=0 时 BiLD loss略好,T=1 时 Top-K 最好,综合 Top-K 最稳。用 target-generated data 还能把 T=1 三任务均值 4.92→5.19,说明 objective harmonization 不能完全替代 trajectory distillation。
2. Harmonized context alignment¶
EAGLE 训练第 j 个 future feature 时用真实 target previous feature,推理只能递归使用 draft feature,产生 feature-level exposure bias。HASS 做 n-step rollout training:step1 同 EAGLE;step j 的 query 来自上一步 draft feature,K/V 前部用 target features、最近 j−1 段用此前 draft rollouts。定制 attention mask 让整条序列并行构造这些混合 context。
2/3/4/5 对齐步的 acceptance 均值显示 3–4 最优,5 步反降:容量过多照顾远期会损害首 token。用 β^{j-1} 重权早期,Llama3-70B β=.5 略优。它揭示应优化 prefix survival 而非等权 future-token accuracy。
3. 实验¶
LLaMA2-Chat 7/13B、LLaMA3-Instruct 8/70B;MT-Bench/HumanEval/GSM8K;H800、batch=1。dynamic tree 60 nodes、depth 6,ShareGPT 68k training,与 EAGLE-2 其余设置一致。
T=0 三任务平均 HASS speed 为 3.24/3.65/3.09/4.05×,EAGLE-2 为 2.81/3.30/2.83/3.37×;T=1 为 2.90/3.48/2.81/3.85× vs 2.60/3.21/2.47/3.28×。acceptance length HASS 约 5.08–5.58,EAGLE-2 约 4.31–5.16。代码模板使 HumanEval 普遍最好。
4. 边界/复现¶
- target 冻结且 verification 不放宽,所以 correctness 与 EAGLE-2 相同;训练目标本身不承担 lossless 证明。
- n-step training 增训练 FLOPs/内存;附录显示开销,但标题“无额外开销”仅指 inference。
- 依赖 EAGLE target features/LM head,每个 target checkpoint 通常需专训。
- 60-node tree 的速度依赖 kernel;不同 node budget 比较不公平。
- top-K objective 可能忽略低概率 rare/safety token,虽最终 target 校正输出,仍会影响 tail acceptance 和调度。
审读导航¶
| 内容 | 页码 |
|---|---|
| 两类 mismatch | 1–3 |
| Top-K loss、多步 context mask | 3–5 |
| 完整结果/设置 | 5–7 |
| loss 与对齐步消融 | 7–9 |
| 训练成本、跨任务/低数据补充 | 13–22 |
原始来源¶
- https://openreview.net/forum?id=T9u56s7mbk
- https://github.com/HArmonizedSS/HASS