Learning What to Imitate: Entropy-Aware Distribution Mixing

What it claims

Entropy-Aware Mixing replaces a fixed distillation target with a per-token mixture of student and teacher. At each token, the student's normalised entropy sets how far the target moves toward the teacher. The paper says the main benefit is fewer "confident conflicts", the tokens where a low-entropy student is pushed toward a token it considers unlikely. Its Qwen3-1.7B student / Qwen3-8B teacher analysis on MATH finds these conflicts are mostly boilerplate rather than reasoning: whitespace, newlines, LaTeX symbols such as $ and \frac, digits and short arithmetic fragments. They are common in teacher-generated traces and rare in student-generated ones, even within one model family.

How it works

• Uncertainty signal. Student entropy is computed over the top-K tokens, renormalised, and divided by log K, giving a value in [0,1]. K is chosen empirically to track full-vocabulary entropy. The value of K is not in the fetched text.
• Schedule. A monotonic f maps entropy to the mixing weight αt. A concave f gives the teacher more weight overall; a convex f concentrates teacher influence in high-entropy states. The defaults are:
• Offline: f(x)=x^(1/4), called "sqrt2".
• Online: f(x)=x², called "sq".
• Mixtures.
• ConMix: linear interpolation, (1−αt)·student + αt·teacher.
• GeoMix: a normalised geometric mean, which favours tokens both models rate likely.

flowchart TD
  S[Student top-K entropy] --> A[alpha = f entropy]
  A --> M{ConMix or GeoMix target}
  M -->|offline| O[Speculative decoding: student drafts, mixture verifies]
  O --> SFT[SFT on sampled traces]
  M -->|online| F[On-policy forward KL to mixture, stop-grad]

• Offline regime. Traces are sampled from the mixture using speculative decoding (vLLM), with the student as the draft model. The paper says this preserves the exact target distribution. For each question it samples 8 rollouts and keeps one, a correct one if any exist. It then runs SFT in verl.
• Online regime. Fully on-policy: one rollout per question per step. The mixture uses a renormalised top-K teacher distribution and is treated as fixed via stop-gradient.

Evidence

All results come from training on the MATH train split with three seeds, each run on one node of 4× GH200 120GB.

Offline, Qwen3-1.7B (non-thinking), teacher Qwen3-8B:

Training on teacher traces hurts per-sample accuracy and general capabilities. Training on student traces protects them but gives lower Pass@64. The mixtures roughly recover both.

Online, Qwen3-1.7B-Base, AIME24 Avg@64:

GeoMix-FKL also leads on AIME24 Pass@64 (39.72 vs FKL 35.56). On GSM8K Avg@64 it essentially ties OPD (81.62 vs 81.57).

How to use it

• Code: d23845jg/entropy-aware-mixing on GitHub. No install or run commands are in the fetched text.
• Stack: vLLM for trace generation, verl for training.
• Model setup: Offline students are non-thinking instruct models; online students are base models.
• Generation settings: temperature 1.0, top-p 1.0, max 4096 tokens.
• Schedule choice: Pick by regime, concave (x^(1/4)) offline and convex (x²) online. The paper says the optimum depends on the training source.

Limits

• Small gains. Many gains are within or near the reported standard deviations. For example, ConMix-SFT has AIME24 Pass@64 ±9.43 and IFEval ±3.88.
• Untuned baselines. Baselines were not hyperparameter-tuned; they follow Jin et al.'s setup.
• ConMix is not consistently better. Online, ConMix-FKL trails OPD on AIME24. Offline, student-generated traces still score best on MMLU.
• Narrow training setup. All training is on MATH, with students between 305M and 2B parameters.
• Results not in the fetched text. Code-generation results and the LFM2.5 and Granite students are in appendices that were not fetched. The page is cut off partway through Table 2, so online general-capability numbers are not shown.
• Related work. AMiD already covers both mixture forms; the authors discuss this overlap in Appendix B.