YichengWangCA/R1-Distill-Qwen-14B-AIME-DPO-LoRA

R1-Distill-Qwen-14B-AIME-DPO-LoRA

A LoRA adapter for deepseek-ai/DeepSeek-R1-Distill-Qwen-14B, trained with DPO ignition: DPO + NLL on self-generated preference pairs, used to recover specific AIME 2024 problems on which the base model has no correct signal at all.

📄 Code and full report: `wangyichengsh/dpo-ignition`

This is a targeted intervention experiment. It tests whether post-training can recover specific problems that the base model reliably fails. It is not an attempt to improve AIME performance in general.

What this adapter demonstrates

On-policy RL such as GRPO can only amplify a correct signal the model already has. If every sample on a problem is wrong, all rewards in the group are equal, the advantage is zero, and there is nothing to learn from. The usual fix is an SFT "ignition" step on solutions from another model. That transplants a foreign distribution into the target model and needs many distinct solution paths.

DPO ignition takes a different route. The positive examples are produced by the target model itself, conditioned on a hint that supplies the missing step. Training then happens on the hint-free prompt, so what transfers is the reasoning path, not the hint text. Rounds are repeated, with the hint weakened each time, until the model solves the problem with no hint at all.

The base model fails 6 of the 30 AIME 2024 problems in all 32 samples. Four were attempted:

Held-out originals

The original AIME problems were never trained on during ignition. All ignition pairs come from 8 programmatically generated isomorphic variants per problem (same structure, different parameters, answers computed exactly with Fraction). The original problem is kept as a held-out probe. So "solved" means the model now solves the real AIME problem after training only on its variants.

A note for anyone evaluating problem I-8

In simplescaling/aime24_nofigures, problem I-8 is under-determined as stored. The figure was stripped, and the text omits the constraint the figure carried: every circle in the chain is tangent to the same side, with the first and last circles tangent to the other two sides. Without it the chain can bend arbitrarily and the inradius is not unique. In my runs, simplescaling/aime24_figures did not fix this either.

I evaluated on a corrected AIME 2024 set: `YichengWangCA/aime24-official` The official wording of I-8 is:

Eight circles of radius 34 can be placed tangent to $\overline{BC}$ of $\triangle ABC$ so that the circles are sequentially tangent to each other, with the first circle being tangent to $\overline{AB}$ and the last circle being tangent to $\overline{AC}$. Similarly, 2024 circles of radius 1 can be placed tangent to $\overline{BC}$ in the same manner. […]

Usage

The adapter is published on its own; the base model is pulled automatically.

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

BASE = "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B"
ADAPTER = "YichengWangCA/R1-Distill-Qwen-14B-AIME-DPO-LoRA"

tok = AutoTokenizer.from_pretrained(BASE)
model = AutoModelForCausalLM.from_pretrained(
    BASE,
    quantization_config=BitsAndBytesConfig(
        load_in_4bit=True, bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True),
    device_map={"": 0}, dtype=torch.bfloat16)
model = PeftModel.from_pretrained(model, ADAPTER)
model.eval()

REASON = "Please reason step by step, and put your final answer within \\boxed{}."
prompt = tok.apply_chat_template(
    [{"role": "user", "content": f"{question}\n\n{REASON}"}],
    tokenize=False, add_generation_prompt=True)

ids = tok(prompt, return_tensors="pt", add_special_tokens=False).to(model.device)
out = model.generate(**ids, do_sample=True, temperature=0.6, top_p=0.95,
                     max_new_tokens=18500, pad_token_id=tok.pad_token_id)
print(tok.decode(out[0, ids["input_ids"].shape[1]:], skip_special_tokens=True))

Chains of thought on these problems routinely exceed 10k tokens. A small max_new_tokens truncates the reasoning before the final \boxed{}, which then looks like a wrong answer. Use temperature 0.6: higher temperatures severely hurt accuracy on long chains of thought.

Training data

All preference pairs are self-generated by the base model plus the in-progress adapter. No external teacher model and no human-written solutions were used.

The training loop

Each problem was worked one at a time, in rounds. A round is:

1. 1.Closed-book sampling. Sample the current model on each hint-free variant (32 samples). Any output with the correct answer becomes a self-generated positive, which is the target signal.
2. 2.Hint-conditioned sampling. Sample again with a hint appended that supplies the missing step. Correct outputs become positives. References to the hint are rewritten into first-person reasoning, and samples that still mention or quote the hint are dropped, since the stored prompt contains no hint. Hinted positives must also show the derivation, not just apply the formula.
3. 3.Pairing. Positives (self-generated first) are paired against closed-book failures from step 1. Both sides are stored against the hint-free prompt.
4. 4.Train, then repeat. Once step 1 yields correct answers on its own, the hint is dropped and a few consolidation rounds use self-generated positives vs self-generated negatives only. Pairing one positive with more than 4 negatives caused degradation.

Per-problem log

2024 AIME II-9: chips in a 5×5 grid (solved, 5 rounds). Variants: 2×2, 2×3, 3×3, 3×4, 4×4, 3×6, 4×6, 4×7 grids. Rounds 1–2 used 50 hinted pairs each; hinted accuracy rose and total closed-book tokens fell. Round 3 used 200 pairs, after which self-generated positives appeared. Rounds 4–5 were consolidation rounds of 50 self-generated pairs each. The part the model kept dropping was the maximality condition (no further chip can be added).

2024 AIME I-8: chain of tangent circles (solved, 3 rounds). One round of ~200 hinted pairs produced self-generated positives, followed by two consolidation rounds of ~200 pairs each. The problem statement first had to be corrected (see the note above).

Global regression repair. After igniting II-9 and I-8, overall AIME 2024 pass@1 dropped sharply. Two rounds of closed-book self-generated pairs on the full AIME 2024 set (32 samples per problem, ~200 pairs per round) brought pass@1 back up to 64%.

2024 AIME II-8: torus–sphere tangency (not included). Three rounds of 200 pairs produced self-generated positives, but overall AIME 2024 pass@32 dropped too much. Rolling back to round 1 and mixing in AIME 2024 self-generated pairs did not fix it. A possible cause is the capacity limit of a rank-32 LoRA.

2024 AIME I-11: octagon two-colouring (abandoned). The problem requires exhaustive enumeration with deduplication, which is better done by a program than in a chain of thought. After several rounds the token count per solution did not converge. This adapter only targets problems the model can solve entirely within its CoT.

Compute

Method

A single-stage objective combining DPO with an NLL term on the positive example:

L = -log σ( β · [ (ℓ(y_c) - ℓ_ref(y_c)) - (ℓ(y_r) - ℓ_ref(y_r)) ] )  +  λ · ( -ℓ(y_c) / W_c )
ℓ(y) = Σ_t w_t · log π(y_t | x, y_<t)

• —NLL term. Keeps the likelihood of the positive from collapsing, a known failure mode when DPO is left to push both sides down. It is averaged per token.
• —DPO term uses the sum, not a per-token mean. The number of tokens spent is itself informative, and the sum naturally penalises overly long chains of thought.
• —Length damping (LD-DPO style, `--ld-alpha`). With K = min(|yc|, |yr|), tokens past K get weight α (w_t = 1 before K), and W = K + α(|y| − K). α = 1 is standard DPO. The rejected answer (usually the longer side) may still contain valid reasoning after an early mistake, so its tail is down-weighted instead of pushed down in full. Igniting II-9 with α = 1 caused visible regressions on other AIME problems, so α = 0.2 was used from I-8 onwards.

Implementation notes that affect reproduction:

• —Reference logprobs are precomputed once with the starting adapter's weights and cached, so no LoRA swapping is needed during training.
• —Chunked projection over the language-model head. Hidden states are taken once, then only the completion positions are projected in chunks, so the full [seq, vocab] logits tensor is never materialised. At 22k tokens this avoids a ~6.2 GiB allocation per sequence.
• —EOS is appended to positive examples only. Sampling decodes with skip_special_tokens=True, which strips the model's own EOS; appending it restores the real target. Appending it to negatives would train the model not to stop after a wrong answer.
• —Degenerate pairs are dropped at load time: empty positive, empty negative, or an identical pair. Each of these silently produces a zero or one-sided gradient while the training log still looks healthy.

Hyperparameters

Evaluation

AIME 2024 (corrected set), temperature 0.6, top-p 0.95, 32 samples per problem.

Limitations

• —Hyperparameters and sample sizes are tuned for this base model with a rank-32 LoRA.
• —Ignition only worked on problem types where isomorphic variants are easy to construct.
• —Igniting a third problem (II-8) caused large global degradation. Whether this is a LoRA capacity limit, a property of the dense architecture, or a conflict with the model's prior is not yet known.

Acknowledgements

Base model: `deepseek-ai/DeepSeek-R1-Distill-Qwen-14B`. Problems adapted from the 2024 American Invitational Mathematics Examination (MAA).