hbin0701/qwen3-1.7b-rewind-runs

Qwen3-1.7B rewind runs

Can a small reasoning model learn to undo its own mistakes with one new token? This repository holds every training run of that study: the training data, the trained models (LoRA adapters, every epoch), the starting checkpoint they all train from, and the start data everything was built from (the problems and the untrained model's attempts). Qwen3-1.7B, AIME problems.

Two methods were tried:

• Rewind. The token <rewind> (id 151669) may fire in the middle of the model's reasoning. The generation loop deletes the last 384 tokens, snaps back to a step boundary, and lets the model carry on. Two generations of training; neither beat the untrained model.
• Reset. The model writes its whole attempt, then makes one decision at the token after </think>: <rewind> deletes the whole attempt and it starts again. Letting the model reset itself hurt accuracy. But the same trained model used only as a judge of another model's finished attempts raised AIME24 accuracy from 48.4% to 58.3% in a reject-and-retry loop (up to two retries over 64 base-model attempts per problem). A judge that only reads length gets 56.2% of the way there, so most of the gain is the judge reading length.

Layout

One folder per training run. Each holds its training rows (data/), the model trained on them (model/, one folder per epoch), and a README saying what it tested and what came of it.

start-data/                  what everything was built from; see start-data/README.md
├── training-pool/           1,878 AIME/AMC problems, 4 base-model attempts each,
│                            741 located first errors                    0.3 GB
└── aime24/                  the 30 test problems, 64 base-model attempts each 0.1 GB
start-model/                 Qwen3-1.7B with the <rewind> row added (id 151669).
                             Every run trains from it.                     3.3 GB
reset/                       one decision after </think>
├── paired/                  one right + one wrong attempt per problem     0.5 GB
│   ├── README.md            what it tested, the result
│   ├── data/                train.jsonl (800 rows), val.jsonl (70), README.md
│   └── model/               epoch1/ epoch2/ epoch3/, train_log.jsonl, README.md
└── all/                     every attempt                                 0.8 GB
rewind/                      mid-trace rewind, generation 2
├── v1/                      one rewind example per error                  0.9 GB
├── v2/                      up to three examples per error                1.1 GB
├── v3/                      correct attempts only, no token (control)     0.8 GB
├── v1w1/                    v1 with the token's loss weight fixed at 1    0.9 GB
└── v2w1/                    v2 with the token's loss weight fixed at 1    1.1 GB

The start data

• start-data/training-pool/: the 1,878 training problems (918 AIME 1983-2023, 960 AMC 10/12 2000-2022), the untrained model's 4 attempts at each (7,512, 66.1% right), and the 741 wrong attempts where a larger critic model located the first mistake. Every training set below was built from these.
• start-data/aime24/: the 30 AIME 2024 test problems and the untrained model's 64 attempts at each (1,920; 48.4% right on one draw, 70.0% by majority vote). The judge-retry result was measured on this pool.

start-data/README.md explains every file and field.

The runs

run	training data	result
reset/paired	one right and one wrong rollout from each problem with both outcomes; 800 / 70 rows, one label each	held-out AUROC 0.669 / 0.723 / 0.734 (epochs 1-3). Resetting itself hurt. As a judge, epoch 3: 58.0% with up to two retries
reset/all	every rollout; 6,496 / 557 rows. The wrong rows carry weight 1.949, but it cancelled in training (each row's loss was divided by its own weight), so all trained unweighted	held-out AUROC 0.946 / 0.946 / 0.938, mostly from length (length alone: 0.818). Resetting itself collapsed into giving up. As a judge, epoch 2 is the headline: 58.3%, +9.9 points
rewind/v1	one <rewind> example per located error, plus correct attempts at 3:1 in graded tokens; token weight solved to 0.05% of the loss	40.6% on AIME24 vs the base model's 50.0%; fired on 98% of runs, 12x earlier than training placed it
rewind/v2	v1 with up to three examples per error	42.5%
rewind/v3	v1's correct attempts only (the control)	45.0%: the fine-tuning itself cost accuracy
rewind/v1w1	v1, token weight 1	47.5%
rewind/v2w1	v2, token weight 1	38.9%

Generation-2 accuracies are 12 greedy samples x 30 AIME24 problems, against the untrained model's 50.0% from one greedy sample; gaps under about 5 points are noise. The earlier pilot (generation 1) is in hbin0701/qwen3-1.7b-rewind-pilot.

Download

Everything (9.5 GB), into the folder the code expects:

huggingface-cli download hbin0701/qwen3-1.7b-rewind-runs --local-dir src/math/runs

Or one run plus the starting model:

huggingface-cli download hbin0701/qwen3-1.7b-rewind-runs --local-dir src/math/runs \
    --include "start-model/*" "reset/all/*"

From the code repository, python -m rewind.prepare.download_runs does the same (--only reset/all for one run).

Use a model

Each model/epochN/ is a LoRA adapter on top of start-model/, plus the trained <rewind> embedding row (embed_row.pt, which the adapter does not carry) and mode.json (whether generation should use rewind or reset rules).

To serve it, merge it into a full model first (CPU, about a minute):

python -m rewind.training.merge --dir runs/reset/all/model/epoch2

Without the code, the same steps in Python:

import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

base = AutoModelForCausalLM.from_pretrained("start-model", dtype=torch.float32)
row = torch.load("reset/all/model/epoch2/embed_row.pt", weights_only=False)
with torch.no_grad():
    base.get_input_embeddings().weight[row["token_id"]] = row["row"]   # the trained <rewind> row
model = PeftModel.from_pretrained(base, "reset/all/model/epoch2/adapter").merge_and_unload()
model.save_pretrained("all-epoch2-merged")
AutoTokenizer.from_pretrained("start-model").save_pretrained("all-epoch2-merged")  # keeps <rewind> one token

To use a reset model as a judge: build the prompt with the model's chat template, append the attempt up to and including </think>, run one forward pass, and read the probability of token 151669 at the last position. That is P(this attempt is wrong).

Never call resize_token_embeddings on these checkpoints: the token uses a spare row (the vocabulary has 151,936 rows for 151,669 tokens), and resizing corrupts Qwen3.

Data format

Every data/train.jsonl and data/val.jsonl row is one training sequence:

{"input_ids": [...], "labels": [[pos, token, weight, tag], ...], ...}

A label at position pos is predicted from the hidden state at pos - 1; positions without a label carry no loss. Tags: rewind (the <rewind> target), sft_prefix (the correct prefix of a wrong attempt), sft (a correct attempt). Reset rows have exactly one label, at the token after </think> (kind is reset or keep). Each data/README.md lists how that set was built; rewind/*/data/meta.jsonl records its counts and loss weight.

How they were trained

LoRA r 32, alpha 64, dropout 0.05 on all seven projections of all 28 layers; lr 1e-4 cosine with 5% warmup; 3 epochs; batch 8 per GPU x 2 GPUs; the <rewind> embedding row trained, gradient-masked to that row. Each model/README.md has the details and the validation loss per epoch.

Files not included

The trained models' own answers to AIME24 and the judges' scores, which the results above are computed from, are in the code repository (src/math/data/qwen3-1.7b/rewind/), not here. The optimizer states (needed only to resume training) and the abandoned Yes/No-readout runs were deleted before upload.