跳转至

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