LT-Tuning-code — latent thoughts tuning on code execution, real public data

Latent-reasoning checkpoints for code output prediction (execution reasoning), trained with Latent-Thoughts-Tuning on a three-tier difficulty ladder built entirely from real public datasets. No synthetic data.

Status — read before using. All three stages are trained and complete. stage2/ is the one you want — it is the end of the curriculum. stage0 is plain SFT and has no latent reasoning; the latent mechanism starts at stage1 (hidden_state) and is completed at stage2 (soft_fusion). stage0 and stage1 are published for reproducibility and ablation, not for use.

Training data

Source	Citation	Rows	Tier role
CRUXEval	Gu et al. 2024, arXiv:2401.03065	800	shallow–medium
LiveCodeBench `execution-v2`	Jain et al. 2024, arXiv:2403.07974	479	deep (native `numsteps` 497–996)
MBPP	Austin et al. 2021, arXiv:2108.07732	374	shallow–medium

1488 train / 165 val, seed 0. Tier distribution: T0 562 · T1 612 · T2 479. Format: {question, steps, answer}. Identical underlying pool to the CoLaR and Latent-SFT runs, so the three platforms are directly comparable.

Difficulty ladder: depth_proxy = #lines + 8 × #loops (CRUXEval/MBPP, median split → T0/T1); LiveCodeBench uses its native numsteps and forms T2. Ordinal scale only.

Scope note on the reasoning chain. Questions and answers are 100% real public data. The intermediate steps are derived mechanically from the real source code (the function's own body lines, capped at 7, plus one templated concluding sentence) — not human- or model-written CoT.

Recipe


Driver	`NeosKnight233/Latent-Thoughts-Tuning` @ `c18aac6`
Base	`Qwen/Qwen2.5-1.5B-Instruct`
stage_epochs	`[1, 2, 7]` — upstream official (README "Key parameters")
Batch	4 × grad-accum 4 = effective 16 (upstream uses accum 1 × 4 GPUs; we are single-GPU, so this is an equivalence conversion, not a deviation)
LR	5e-5, cosine, warmup_ratio 0.05, weight_decay 0.01
`labels_per_stage`	`[0, 10, 16]`
`thinking_insertion_prob`	`[0.0, 0.85, 0.95]`
`fusion_alpha`	`[0.5, 0.5, 0.6]`
`thinking_strategy`	`confidence`
attn	`sdpa` (`use_flash_attention: false` — no flash-attn on the runtime)

Note: configs/example_config.yaml upstream ships [1,1,1], which conflicts with the README's [1,2,7]. We follow the README. Under [1,1,1] stage2 receives 1/7 of its official training and does not converge (observed: loss 19.1 → 18.7, essentially flat).

What the three stages actually do

They are not three epochs of the same thing — the forward computation graph differs. The only place stage_mode is read during training is model.py:512-529, and it decides what embedding occupies each <thinking> slot:

Stage	Mode	`<thinking>` slot receives
0	`common`	nothing — the dataset contains no `<thinking>` tokens at all → plain SFT
1	`hidden_state`	the previous position's hidden state, directly
2	`soft_fusion`	`α·hidden + (1−α)·(top-p filtered, temperature-scaled weighted average of vocabulary embeddings)`

The loss is identical in all three (plain shifted cross-entropy, model.py:543-552). Because the injected embeddings differ, loss values are not comparable across stages — stage0's 0.155 and stage1's ~19 are not on the same scale.

Checkpoints

`stage0/` ✅ — `stage0-cot`, plain SFT, 1 epoch

model.safetensors (3086.6 MB) and checkpoint-93/pytorch_model.bin (3086.7 MB; 93 steps = one epoch at effective batch 16). Training loss 5.132 → 0.155. checkpoint-93/pytorch_model.bin is the file to chain stage1 from — upstream run.py:232 calls torch.load(load_model_path) and therefore requires a file, not a directory.

`stage1/` ✅ — `stage1-hidden-state`, 2 epochs (official)

model.safetensors (3086.6 MB) plus per-epoch checkpoint-93/ and checkpoint-186/ (pytorch_model.bin, 3086.7 MB each). Training loss 38.878 → 0.915 over 19 logged points; train_runtime 3218 s, 0.925 samples/s. Chained from stage0/checkpoint-93/pytorch_model.bin. checkpoint-186/pytorch_model.bin is the file stage2 chains from.

At this stage the <thinking> slot is filled with the previous position's hidden state directly.

`stage2/` ✅ — `stage2-soft-fusion`, 7 epochs (official) — final checkpoint

model.safetensors (3086.6 MB) plus checkpoint-558/ (epoch 6) and checkpoint-651/ (epoch 7, final) at 3086.7 MB each; 651 steps = 7 × 93 at effective batch 16.

Training loss 1.5185 → 0.1042 over 66 logged points, converging smoothly (last eight points 0.128 / 0.108 / 0.113 / 0.138 / 0.116 / 0.143 / 0.105 / 0.104). train_runtime 14865.7 s (4.13 h), train_loss 0.6170, 0.701 samples/s. Chained from stage1/checkpoint-186/pytorch_model.bin.

For contrast, an earlier attempt using the example_config.yaml value stage_epochs: [1,1,1] left stage2 flat at 19.1 → 18.7 — one seventh of the official training, chained from an under-trained stage1. That run is not published; the README's [1,2,7] is what works.

Local patches applied

Upstream at c18aac6 needed four fixes to run to completion here. Recording them so results are reproducible:

Fix	Why
`save_safetensors: false`	embed/lm_head are tied; safetensors refuses shared-memory tensors (`RuntimeError: Some tensors share memory`)
`save_strategy: epoch`	default `steps`/1000 never fires on 1488 rows → no intermediate checkpoint ever written
`save_dataset: false`	`run.py:998` calls `train_dataset.save_to_disk()` unconditionally, but the `confidence` strategy takes the serial branch returning a plain `list` → `AttributeError`, after preprocessing completes
move model to GPU before preprocessing	`run.py` never moves the model to GPU (no `device_map`/`.to(cuda)`), and `build_thinking_strategy` (`run.py:967`) runs before the Trainer exists, so `dataset.py:243` captures `device=cpu` and the entire preprocessing forward runs on CPU (measured 48.90 s/it on an idle A100 — stage1's preprocessing took 2:31:35; after the patch the identical step on stage2 ran at 10.75 it/s, i.e. ~20 seconds)

Known deviation from upstream

Each stage is run as a separate job chained via load_model_path, for resumability. Upstream runs all epochs in one job under a single cosine schedule (run.py:113 default cosine, no optimizers= override); here each stage restarts its own cosine from 5e-5. This raises stage2's learning rate relative to the fused schedule.