Struggle to reproduce TaH+ at 0.6B #8

Thanks for the great work, and for releasing the code and the 1.7B checkpoint.

We're trying to reproduce the 0.6B TaH+ result and are getting essentially the Standard SFT baseline back. We'd appreciate a sanity check on our setup. Importantly, the gap is visible on GSM8K (which is low-variance), so this isn't just MATH500 sampling noise.

Setup

• Base: Qwen/Qwen3-0.6B-Base
• Method config: script/recipes/qwen3_0.6/sft_tah_step1.yaml → sft_tah_step2.yaml
  (max_iter=2, FixedLabelIterDecider label_type: mismatch, TrivialUpdater, AdditiveLogitsUpdater,
  adapter: lora r=16 α=32 dropout=0.1; 5 epochs step1 / 2 epochs step2, lr 4e-5, cosine_with_min_lr, max_length 8192)
• Data: the full OpenR1-Math-220k (59,824 samples / 238M tokens), processed with your pipeline;
  iteration (mismatch) labels generated with Qwen/Qwen3-0.6B as the reference model. Note we scaled the data up
  from the checked-in recipe default (openr1_math, a 7,502-sample / 30M-token subset) to approach the paper's
  ~300M-token budget; all other hyperparameters are the recipe defaults.
• Eval: your eval.py --backend tah, temperature 0.6, repeat_size: 1, max_new 4096, full GSM8K (1319) + MATH500 (500)

Results (our from-scratch 0.6B; Standard and TaH-plus trained on the same data)

For us TaH-plus lands right on top of Standard SFT on GSM8K (61.9 vs 60.9), whereas the paper's TaH-plus is
well above its Standard (67.6 vs 58.2). Since GSM8K is low-variance, this isn't a MATH500 sampling artifact.

Our eval harness is correct (so the gap is training-side)

Your released TaH-plus-1.7B checkpoint, evaluated on the exact same harness, reproduces its headline:
GSM8K 83.85 (paper 84.5) and MATH500 68.0. So the pipeline/eval is wired right.

We've ruled out the obvious causes

• The iter-decider fires: ~11–12% of tokens take a 2nd iteration on our 0.6B run (your released model shows ~7–9% on the same harness). Not a dead gate / mis-set threshold.
• The deep-iter LoRA is trained and non-trivial: ‖lora_B‖ ≈ 0.69 (well off the zero-init), applied on iter-2 via AdditiveLogitsUpdater.

Questions

1. Training data: the paper describes the math subset of Mixture-of-Thoughts, but the checked-in qwen3_0.6 recipe points at OpenR1-Math. Which produced the Table-2 numbers? (We used the full OpenR1-Math-220k, matching the recipe's source but scaled to ~238M tokens.)
2. 0.6B specifics: any difference from the checked-in recipe for the reported 0.6B run — the mismatch reference model used to generate iteration labels, epochs, or LoRA rank/α?
3. MATH500 evaluation: how many samples per question did the reported numbers use? On our side single-sample MATH500 (repeat_size 1, temp 0.6) is quite high-variance (we see multi-point swings across seeds on a fixed checkpoint), so we're trying to match your protocol.
4. Would you consider releasing the 0.6B TaH-plus checkpoint (only the 1.7B general checkpoint appears public) so we can diff a working reference against ours?

Thanks a lot!