Qwen3-4B-AbstractCoT-warmup

Qwen3-4B fine-tuned with the Abstract Chain-of-Thought (Abstract-CoT) warm-up procedure from "Thinking Without Words: Efficient Latent Reasoning with Abstract Chain-of-Thought" (Ramji, Naseem, Fernandez Astudillo, IBM Research AI, 2026). The model is taught to compress its reasoning into a short sequence (~16–22 tokens) drawn from a reserved 64-symbol abstract vocabulary V_abs = {<TOKEN_A>, …, <TOKEN_BL>}, used as a discrete latent scratchpad before emitting the answer.

prompt  ─►  <beginabstract>  z_1 ... z_m  <endabstract>  answer
            └─────── z̃ ∈ V_abs^m, m ≤ 128 ───────┘

This is the SFT half of the paper only — no RL stage. The comparison row is the paper's "Abstract-CoT (Warm-up)" line in Table 1.

Headline result

	MATH-500 acc	Mean tokens
Paper Baseline (Qwen3-4B verbal CoT)	83.2	1087
Our Baseline (Qwen3-4B verbal CoT, this hardware)	84.60	1045
Paper Abstract-CoT Warm-up	86.2	168
This model (T=3 PI, N=5k, 1 epoch, LoRA, seq 8k)	72.00	432

The accuracy gap to the paper's 86.2 is driven by reduced data scale (5k vs 600k), LoRA vs full fine-tuning, and 1 vs 3 epochs per phase. See docs/20260511_reader.md for a full discussion.

Repository layout

final/             ← end-of-round-3 merged model (THE warm-up checkpoint)
round2/            ← end-of-round-2 merged model
round1/            ← end-of-round-1 merged model
adapters/          ← all 6 LoRA adapters (pi{1,2,3}_phase{A,B})
results/           ← per-example eval JSONL (baseline + abstract)
teacher_traces/    ← on-policy V_abs traces used as Phase B/A teachers
train_logs/        ← per-phase loss + LR curves (verifies cosine fix)
docs/              ← run reports (technical + reader-oriented)

How it was trained

Three policy-iteration rounds, each with two phases:

• Phase A — Bottleneck SFT. Train on [prompt; verbal-CoT; z̃; answer] with the answer blocked from attending to the verbal CoT, forcing all CoT→answer signal through z̃.
  • Round 1: z̃ is random V_abs tokens.
  • Rounds 2+: z̃ is sampled on-policy from the previous round's model.
• Phase B — Self-distillation. Train on [prompt; z̃; answer] with standard causal attention, where z̃ is now generated from the prompt alone.

Training config:

• Base: Qwen/Qwen3-4B, extended with V_abs (M=64) + <beginabstract> + <endabstract> (151 669 → 151 735 tokens).
• LoRA r=32, α=64 on attention + MLP projections. Embedding table + LM head trained fully (so the new abstract-vocab rows can move freely). 842.9 M / 4.86 B trainable (17.3%).
• Data: 5 000 examples from allenai/Dolci-Think-SFT-7B, filtered to assistant messages with <think> blocks ≥ 200 chars.
• max_len 8192, batch 32, lr 1e-4, cosine schedule, 5% warmup.
• 2× A100-SXM4-80GB, ~11 hours wall.

Using the model

Inference (vLLM, recommended)

from vllm import LLM, SamplingParams
from transformers import AutoTokenizer
from huggingface_hub import snapshot_download

# Download the final checkpoint
model_path = snapshot_download(
    "leapeto/Qwen3-4B-AbstractCoT-warmup",
    allow_patterns=["final/*"],
)

tok = AutoTokenizer.from_pretrained(f"{model_path}/final", trust_remote_code=True)

# Abstract token ids
abs_tokens = []
for i in range(64):
    if i < 26:
        abs_tokens.append(f"<TOKEN_{chr(ord('A')+i)}>")
    else:
        j = i - 26
        abs_tokens.append(f"<TOKEN_{chr(ord('A')+j//26)}{chr(ord('A')+j%26)}>")
end_id = tok.convert_tokens_to_ids("<endabstract>")
abs_ids = tok.convert_tokens_to_ids(abs_tokens)
allowed = list(set(abs_ids + [end_id]))

llm = LLM(model=f"{model_path}/final", tensor_parallel_size=2,
          dtype="bfloat16", trust_remote_code=True)

# Two-stage decode: (1) constrained abstract trace, (2) unconstrained answer
prompt = "What is the integral of x^2 from 0 to 1? Put your final answer in \\boxed{}."
messages = [
    {"role": "system", "content": "Please reason step by step, and put your final answer within \\boxed{}."},
    {"role": "user", "content": prompt},
]
prefix = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)
prefix += "<beginabstract>"

# Stage 1: V_abs only, stop at <endabstract>
sp1 = SamplingParams(temperature=0.7, max_tokens=128,
                     allowed_token_ids=allowed, stop_token_ids=[end_id],
                     skip_special_tokens=False)
abstract = llm.generate([prefix], sp1)[0].outputs[0].text
prompt2 = prefix + abstract + "<endabstract>\n"

# Stage 2: unconstrained answer
sp2 = SamplingParams(temperature=0.0, max_tokens=2048)
answer = llm.generate([prompt2], sp2)[0].outputs[0].text
print(answer)

Loading the LoRA adapters (peft)

If you want to inspect individual round outputs without downloading the merged models:

from peft import PeftModel
from transformers import AutoModelForCausalLM
from huggingface_hub import snapshot_download

# You'll need the extended base model first — produce it locally via scripts/01_extend_model.sh
# OR start from one of our merged checkpoints and load a later adapter on top.

base = AutoModelForCausalLM.from_pretrained("path/to/extended/base", trust_remote_code=True)
adapter_path = snapshot_download(
    "leapeto/Qwen3-4B-AbstractCoT-warmup",
    allow_patterns=["adapters/pi3_phaseB/*"],
)
model = PeftModel.from_pretrained(base, f"{adapter_path}/adapters/pi3_phaseB")

Files of interest

File	What
final/	End-of-round-3 merged model. This is the main artifact.
round1/, round2/	Intermediate merged models for studying T=1 → T=2 → T=3 progression
adapters/pi{1,2,3}_phase{A,B}/	LoRA-only checkpoints from each phase
results/baseline_math500.jsonl	Qwen3-4B verbal-CoT eval (84.60% / 1045 tok)
results/abstract_math500_T3_N5000.jsonl	This model's eval (72.00% / 432 tok)
train_logs/*.json	Per-step loss + LR curves for each phase
docs/20260511.md	Technical report (full breakdown)
docs/20260511_reader.md	Reader-oriented report (concepts + reasoning)

Citation

@article{ramji2026thinking,
  title={Thinking Without Words: Efficient Latent Reasoning with Abstract Chain-of-Thought},
  author={Ramji, Keshav and Naseem, Tahira and Fernandez Astudillo, Ramón},
  journal={arXiv preprint arXiv:2604.22709},
  year={2026}
}