hbin0701/Qwen3-1.7B-rewind-sftv3-wait

Qwen3-1.7B · <rewind> SFT v3 ("prior mode", wait init)

Qwen3-1.7B fine-tuned to use a single-token <rewind> action while it thinks. Writing <rewind> asks the environment to erase the last 4 sections of the reasoning, insert [rewound], and let the model continue from there (at most once per solution, never after </think>).

This is the SFT start for RL. It was trained to give <rewind> a small, fixed probability (P0 = 2.5%) at legal spots and 0 elsewhere, and was never shown a <rewind> to copy. Deciding when to fire is left to RL.

Code: hbin0701/sclm, branch rl/paired-rloo-stability-hb (rewind/sft_train.py, launched by sft_prior.sh).

The tag and the environment

• <rewind> is one added token (id 151669, not special, so it decodes as the text <rewind> and works as a stop string). It sits in a spare embedding row. Its embedding was initialised to the mean of the embeddings of "Wait", "Alternatively" and "Hmm". Embeddings are tied, so this sets its output logit row too.

• Paragraphs are blocks split at blank lines (\n\n); pieces under 40 characters merge forward.

• Sections: a paragraph starts a new section when fewer than 43% of its content words appeared in the previous 3 paragraphs, once the current section has at least 300 tokens. The count is causal: the section starts of a prefix are the section starts of the whole text up to that point.

• Rewind: delete the last 4 sections of the thinking (the section in progress counts as one), append [rewound]\n\n, and continue. Fired with 4 or fewer sections written, it erases everything (a restart).

• System prompt used in training and in every eval:

  You may write while reasoning. It erases the last part of your reasoning (about the last four steps) so you can continue from the earlier point where things were still on track. Use it only when you realize your reasoning has gone wrong and continuing will not lead to a correct answer; do not use it for routine double-checking. You can use it at most once per problem, and never after . After a rewind, [rewound] marks where the erased part was.

  The user turn is the problem + \n\nPlease reason step by step, and put your final answer within \boxed{}., and the assistant turn opens with <think>\n.

Training data

From Sangsang/rewind-run1-sft-data (146 rows of each type, built from Qwen3-1.7B's own DAPO-Math solutions):

row type	prompt → completion	used here
trigger	reasoning up to a drop → <rewind>	dropped (no tag to copy)
post_rewind	reasoning 4 sections before the drop + [rewound] → a correct continuation	101 (the 45 whose rewind kept nothing, i.e. restarts, are dropped)
negative	an on-track prefix → a correct continuation	146
clean	the problem → a full correct solution	146

393 rows in total, with 3,091 positions targeted at P0 (10.6 per non-rewound row) and 80,795 targeted at 0.

Loss (exact)

For one row with completion tokens x₁…x_T (prompt = system + problem + prefix, masked):

L_row = −(1/T) · Σ_{t=1..T} log π(x_t | ctx_t)                         (completion-only next-token loss)
        + λ · (1/|S|) · Σ_{t∈S} BCE(p_t, y_t)                           (prior term)

p_t      = π(<rewind> | ctx_t)            the softmax probability of the tag token at completion position t
BCE(p,y) = −[ y·log p + (1−y)·log(1−p) ]
λ        = 1  (--prior_weight)

Batch loss = (1/8) · Σ over the 8 rows of the global batch of L_row.

S and the targets y_t (prior_targets in sft_train.py):

• y_t = P0 = 0.025 at legal spots: completion positions that are section starts of the thinking with at least 5 sections before them (starts[5:], --prior_min_sections 5). A rewind fired there always keeps something, so it is never a restart.
• y_t = 0 at every other paragraph start (the previous token ends with \n\n), and at the first thinking token of rows whose prefix is empty.
• y_t = 0 at every paragraph start of post_rewind rows (the one rewind is already used).
• All other positions are not in S.

The BCE gradient on the tag's logit is p_t − y_t, so the term is minimised exactly at p_t = y_t (its floor at P0 = 0.025 is the entropy H(0.025) ≈ 0.117 nats). At a P0 position the next-token term also pushes p_t down (gradient p_t / T), which moves the optimum to P0 · T / (T + |S|) ≈ 0.975 · P0 on this data.

Optimisation: full fine-tuning, AdamW (β = 0.9 / 0.999, weight decay 0.01), lr 2e-6 with 8 warmup steps then cosine to 0, 2 epochs = 98 steps of 8 rows, grad clip 1.0, seed 42, fp32 master weights with bf16 compute, sequences up to 32,768 tokens, 2 × A100 (≈ 27 min, 10.1M tokens).

torchrun --nproc_per_node 2 sft_train.py --model Qwen/Qwen3-1.7B --data sft.jsonl --out ckpt_wait --lr 2e-6 \
  --epochs 2 --batch_size 8 --warmup_steps 8 --max_rewinds 1 --system_prompt on --rewind_unit section \
  --rewind_sections 4 --rewind_mark --tag_token --tag_init Wait+Alternatively+Hmm --rewind_prior 0.025 \
  --prior_on all --prior_min_sections 5 --prior_weight 1 --drop_restarts

Results

Calibration at the end of training (mean P(<rewind>) over the last batch):

positions	target	reached
legal spots	0.025	0.0020
other paragraph starts	0	0.00055
first thinking token	0	0.0

The target at legal spots was not reached. Over the 98 steps, P at legal spots wandered between 0.001 and 0.005. The model learned roughly one flat rate at every paragraph start and little about which spots are legal.

Fire test: 64 DAPO-Math problems × 4 samples, temperature 1.0, top-p 1.0, 16,384 new tokens, system prompt on, thinking on. Each sample stops at its first <rewind>, and nothing is rewound. A second run blocks the tag (logit ban).

measure	value
solutions that write <rewind> (before </think>)	21.5%
fires per paragraph break generated	0.0010 (the best flat rate for this data would be 0.00092)
fires at a paragraph break	100%
fire position, tokens (p25 / median / p75)	1,989 / 4,230 / 7,535 (median 9 sections)
fires that would be restarts (≤ 4 sections written)	30.9%
fires after </think>	0%
accuracy of samples that never fired	69.2%
accuracy with the tag blocked	64.1%
lone < lines with the tag blocked	0%
paragraphs starting with "Wait" (sampled / tag blocked)	8.4% / 9.0%
fire rate on hard / easy problems (0–1 vs 3–4 of 4 right with the tag blocked)	28.6% / 16.7%

Use with a guard. Because the rate is flat across paragraph starts, about 31% of fires would be restarts. In the accompanying RL code, --guard_tokens 1500 --guard_sections 5 blocks the tag until the solution has 1,500 tokens and 5 sections. That makes a restart impossible: 0 restarts in 15 RL steps and in eval.

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer
tok = AutoTokenizer.from_pretrained("hbin0701/Qwen3-1.7B-rewind-sftv3-wait")
model = AutoModelForCausalLM.from_pretrained("hbin0701/Qwen3-1.7B-rewind-sftv3-wait", torch_dtype="bfloat16")

Generate with <rewind> as a stop string. On a stop, delete the last 4 sections of the thinking, append [rewound]\n\n, and continue with the tag banned. rewind_common.apply_section_rewind and eval_rewind.py in the repo implement exactly this. Without the system prompt, the model behaves close to the base model.