You need to agree to share your Hugging Face information to access this model
This repository is publicly accessible, but you have to accept the conditions to access its files and content.
Research checkpoint shared on request. Access is approved manually.
Log in or Sign Up to review the conditions and access this model content.
CODI on Ministral-3-3B-Instruct (research checkpoints)
Plain CODI (no step grounding) on mistralai/Ministral-3-3B-Instruct-2512-BF16 @
b6d637bef2393152b3da2b2fde72eecdee30557e, lr 4e-4. LoRA + projection checkpoints for interpretability research, not a
merged model, not for production use. Epoch 2 is the peak of latent use; epoch 3 is the final checkpoint.
| epoch (step) | files | latent / removed / direct / teacher | latent β removed | transplant follow vs random | mean-fill |
|---|---|---|---|---|---|
| 1 (12,175) | β | 30.4 / 25.3 / 25.8 / 44.1 | +5.1 [+3.3, +6.9] | 36/82 vs 7/82 | 23.3 |
| 2 (24,350) | light + full state | 38.7 / 31.5 / 32.4 / 49.8 | +7.3 [+5.4, +9.2] | 64/109 vs 4/109 | 28.7 |
| 3 (36,525) | light + full state | 44.4 / 39.7 / 40.1 / 52.2 | +4.7 [+3.0, +6.4] | 55/139 vs 2/139 | 36.3 |
Use falls from epoch 2 to 3 as the direct path memorises training items. A linear probe finds the first intermediate result better encoded in the latents than in the question state (+7.9 points at epoch 3).
lr 2e-4 variant (lr2e-4/ subfolder)
Same model, recipe, data order and seed as above, only the learning rate halved (2e-4). Files live under lr2e-4/
(ckpt_<step>.pt light checkpoints and checkpoint-<step>/ full trainer state incl. optimizer). Trained to epoch 3; all five
looks are here. Epoch 2 is the peak of latent use (latent β direct +10.3 [+8.2, +12.4]); epoch 3 is the final checkpoint
(latent β direct +7.5 [+5.5, +9.5]). The lr 4e-4 files above stay at the repo root.
| epoch (step) | files | latent / removed / direct / teacher | latent β removed | transplant follow vs random | mean-fill |
|---|---|---|---|---|---|
| 0.5 (6,088) | light + full state | 27.5 / 28.8 / 28.7 / 51.9 | β1.3 [β2.5, +0.0] | 0/62 vs 0/62 | β |
| 1 (12,175) | light + full state | 41.2 / 34.8 / 35.3 / 58.4 | +6.4 [+4.5, +8.5] | 51/123 vs 10/123 | 30.9 |
| 1.5 (18,263) | light + full state | 44.4 / 36.2 / 36.5 / 60.4 | +8.2 [+6.1, +10.3] | 71/135 vs 3/135 | 31.8 |
| 2 (24,350) | light + full state | 50.8 / 39.1 / 40.5 / 59.7 | +11.7 [+9.6, +13.8] | 84/172 vs 6/172 | 36.6 |
| 3 (36,525) | light + full state | 53.4 / 45.0 / 45.9 / 62.5 | +8.4 [+6.4, +10.4] | 83/177 vs 9/177 | 41.0 |
At 2e-4 the model is ~10 points more accurate than at 4e-4 at matched epochs, with the same or slightly larger latent use.
KL-to-base variant (kl_base/ subfolder; training in progress)
The lr 2e-4 recipe above plus Ξ² Β· KL(p_ΞΈ β p_base) toward the original Ministral-3-3B-Instruct, applied on the
written-CoT (teacher) sequence: question + <|bocot|> + CoT + <|eocot|> + answer, all positions, full vocabulary.
p_base is the same network with the LoRA disabled. Two arms, kl_base/beta0.5/ and kl_base/beta1.0/, same data
order and seed as lr2e-4/. Looks so far: 0.5, 1, 1.5 and 2 epochs (epoch 3 will be added when training ends).
Each look has a light ckpt_<step>.pt and full trainer state checkpoint-<step>/; md5 in kl_base/md5.txt.
| epoch (step) | arm | latent / removed / direct | latent β direct | written CoT (strict) | general-text KL to base |
|---|---|---|---|---|---|
| 1 (12,175) | lr 2e-4 (Ξ² = 0), for reference | 41.0 / 34.6 / 35.1 | +5.9 [+3.8, +8.0] | 58.4 | 0.49 |
| 1 (12,175) | Ξ² = 0.5 | 36.6 / 31.8 / 32.8 | +3.8 [+2.0, +5.6] | 25.6 | 0.26 |
| 1 (12,175) | Ξ² = 1.0 | 39.0 / 30.9 / 33.0 | +6.0 [+4.1, +7.9] | 4.7 | 0.20 |
| 1.5 (18,263) | lr 2e-4 (Ξ² = 0), for reference | 44.4 / 36.1 / 36.3 | +8.0 [+5.9, +10.1] | 60.6 | 0.50 |
| 1.5 (18,263) | Ξ² = 0.5 | 40.3 / 35.9 / 35.9 | +4.4 [+2.6, +6.2] | 48.4 | 0.26 |
| 1.5 (18,263) | Ξ² = 1.0 | 42.9 / 31.3 / 34.9 | +8.0 [+6.1, +10.0] | 7.0 | 0.21 |
| 2 (24,350) | lr 2e-4 (Ξ² = 0), for reference | 50.6 / 39.0 / 40.5 | +10.2 [+8.1, +12.4] | 59.7 | 0.54 |
| 2 (24,350) | Ξ² = 0.5 | 46.9 / 37.8 / 38.7 | +8.2 | 35.3 | 0.27 |
| 2 (24,350) | Ξ² = 1.0 | 46.6 / 35.6 / 39.1 | +7.4 | 6.3 | 0.21 |
At 2 epochs the KL arms trail the Ξ² = 0 run on latent β direct (paired difference Ξ² 0.5 β2.0 [β4.6, +0.6], Ξ² 1.0
β2.7 [β5.4, β0.2]) and on latent accuracy (β3.7 / β4.1). The 0.5-epoch looks are included but were not evaluated. Known side effect: the KL is also applied at the position
that predicts <|eocot|>, a token the base model never emits, so the KL-trained teachers often do not close their
written CoT (hence the low strict written-CoT score; latent and direct modes are unaffected because the evaluation
inserts <|eocot|> itself). MMLU (500-item MMLU-Redux subset) is within Β±3 points of the Ξ² = 0 run at every look (CIs span 0). One seed
per arm; preliminary.
KL on question + written CoT only (kl_textonly/ subfolder; training in progress)
Same as kl_base/beta1.0/ (lr 2e-4 recipe + Ξ² = 1.0 KL to the original model, same data order and seed), with one
change: the KL covers only the question and written-CoT positions. It drops every position whose target is a CODI
special token and every position at or after <|eocot|>; <|bocot|> / <|eocot|> logits are also excluded from the
vocabulary. The goal is to stop the KL from fighting the <|eocot|> token and the answer, which broke kl_base's
written CoT. Files in kl_textonly/beta1.0/; md5 in kl_textonly/md5.txt. Looks so far: 0.5 epoch.
| epoch (step) | latent / removed / direct | latent β direct | written CoT strict / lenient | general-text KL to base | MMLU (CODI-direct prompt) |
|---|---|---|---|---|---|
| 0.5 (6,088) | 25.9 / 26.7 / 27.4 | β1.5 [β2.7, β0.4] | 37.8 / 38.6 | 0.194 | 64.2 |
At 0.5 epoch no arm uses its latents yet (the Ξ² = 0 run was also below zero here). The written CoT closes on 1,066 of 1,319 items. The Ξ² = 0 run skips the CoT about as often as this one, but this one also gives 166 empty outputs and 87 unclosed CoTs, which the Ξ² = 0 run does not. Under investigation. One seed; preliminary.
KL on question + written CoT, also masking the position after <|bocot|> (kl_textonly_opt2/; training in progress)
Same as kl_textonly/beta1.0/ with one more masked position: the one whose input token is <|bocot|> (more generally,
every position whose input is a CODI special token). The original model has never seen <|bocot|> and predicts </s>
there with p β 0.11, which kl_textonly was pulling the student towards (empty outputs). Files in
kl_textonly_opt2/beta1.0/; md5 in kl_textonly_opt2/md5.txt. Looks so far: 0.5 epoch.
| epoch (step) | latent / removed / direct | latent β direct | written CoT strict / lenient | empty outputs | general-text KL to base | MMLU (CODI-direct prompt) |
|---|---|---|---|---|---|---|
| 0.5 (6,088) | 26.2 / 25.5 / 25.5 | +0.7 [β0.6, +1.9] | 56.5 / 56.8 | 0 | 0.185 | 62.6 |
At 0.5 epoch the written CoT is fixed: strict +4.5 [+2.0, +7.1] over the Ξ² = 0 run at the same step (51.9) and +18.7
over kl_textonly (37.8); 0 empty outputs (kl_textonly: 166); p(</s>) after <|bocot|> β 2e-7. No arm uses its
latents yet at this step (transplant 0/58). One seed; preliminary.
Chat-mix variant (chat_mix/M/; training in progress)
Trained from scratch (fresh LoRA, lr 2e-4, seed 11) on Mira's "150" data: GSM8k-Aug-NL + CommonsenseQA CODI rows
with ordinary chat rows interleaved, (3 CODI + 1 chat) Γ 8 per step. The chat rows are the original model's own
greedy replies ([INST]β¦[/INST] template, plain LM loss) and are there to keep its normal chat and step-by-step
reasoning. 16,233 steps = one epoch of CODI rows; warm-up 1,461, cosine over 48,700; batch 32 as micro-batch 16 Γ 2. Files: light ckpt_<step>.pt and full trainer state hf_keep/checkpoint-<step>/ at steps
4,000 / 8,117 / 12,175 and the final step 16,233; md5 in chat_mix/md5.txt.
| step (β epoch) | latent / removed / direct | latent β direct | teacher (written CoT) |
|---|---|---|---|
| 4,000 (0.25) | 22.3 / 21.1 / 21.8 | +0.5 | 61.6 |
| 8,117 (0.5) | 34.4 / 27.3 / 28.9 | +5.5 | 65.7 |
| 12,175 (0.75) | 41.4 / 31.9 / 34.6 | +6.8 | 68.0 |
| 16,233 (1, final) | 46.7 / 35.1 / 37.6 | +9.1 [+6.8, +11.3] | 68.5 |
Chat retention (natural chat, accuracy %; GSM8k-100 / order tracking / card swaps / web of lies):
| model | "think step by step" | plain question |
|---|---|---|
| this arm, final | 73 / 53 / 78 / 89 | 16 / 14 / 85 / 81 |
| same recipe without chat rows, final | 24 / 10 / 23 / 61 (answers with a bare number) | 32 / 14 / 31 / 53 |
| original model | 96 / 83 / 91 / 97 | 12 / 16 / 30 / 75 |
At the final step the chat rows keep most of the step-by-step chat ability that plain CODI loses, with latent use at least as large (latent β direct +9.1 vs +8.3 for the no-chat twin; transplant follow minus random +42.6 points on 277 pairs vs +59.3). Numbers are from the run's final evaluation log before its write-up; one seed; preliminary.
Recipe (all checkpoints in this repo)
- 3-mode CODI (latent student, written-CoT teacher, direct answer), batch 32, LoRA r128 / alpha 32 on q/k/v/o/gate/up/down, 6 latent thoughts, projection + LayerNorm, answer-only distillation.
- Data: GSM8k-Aug-NL + CommonsenseQA. 3% linear warmup then cosine over 10 epochs (121,750 steps); training stopped
at epoch 3. Text-only language model (vision tower dropped). Special tokens
<|bocot|>= 131072,<|eocot|>= 131073. - Prompt (template-free): latent mode
<s> {question} Output only the answer and nothing else.</s><|bocot|>; direct mode... </s><|eocot|>.
Metrics: GSM8k test, n = 1,319, greedy. "latent β removed" = accuracy with the 6 latent thoughts minus the same checkpoint with the thoughts removed (paired bootstrap 95% CI). Transplant = latents taken from a copy of the problem with one number changed; "follow" = the answer becomes the copy's answer, on pairs both versions solve, vs norm-matched random latents. Mean-fill = each latent replaced by its across-problem mean (no problem information).
Files
ckpt_<step>.pt: light eval checkpoint (trainable LoRA + projection weights).checkpoint-<step>/: full Hugging Face Trainer state for exact resume (trainable.safetensors,optimizer.pt,scheduler.pt,rng_state.pth,trainer_state.json).md5.txt: md5 of every file; all match the training pod's records.
Loading and resuming
The training and evaluation scripts (train62.py, eval29.py / eval62.sh, make_text_base.py) are in the authors'
research repository and are shared together with access. To resume exactly, pass checkpoint-<step>/ as --resume
with the same flags as the original job and the same data order; the Trainer restores optimizer, scheduler and RNG and
skips the data already seen.
Model tree for mjkenney/codi-ministral3-3b-instruct
Base model
mistralai/Ministral-3-3B-Base-2512