YichengWangCA

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

What this adapter demonstrates

A model's capability ceiling is largely set during pretraining. Post-training rarely extends that ceiling; what it can do is raise the floor — recovering problems the model has the latent components to solve but reliably fails to assemble on its own.

This adapter is a direct test of that claim. Three AIME 2024 problems were selected on a strict criterion: the base model produced zero correct answers in 64 sampled attempts. Not a low success rate — no correct signal at all. Under that baseline, ordinary rejection sampling or best-of-n has nothing to work with, since there is no positive example to select.

DPO + NLL post-training on those three problems recovered two of the three.

Table with columns: Problem, Base model, 64 samples, After post-training
Problem	Base model, 64 samples	After post-training
2024 AIME II-9	0 / 64	solved
2024 AIME I-11	0 / 64	not solved
2024 AIME I-8	0 / 64	solved

Positive examples were obtained by conditioning generation on a hint that supplies the missing step, then training on the hint-free prompt — so what transfers is the reasoning path, not the hint text.

How to read the AIME 2024 number

The two problems this adapter now solves are the two it was trained on. That is the experimental design, not a confound: the question being asked is whether targeted post-training can recover a specific unsolved problem, not whether it improves AIME performance in general.

A note for anyone evaluating problem I-8

If you evaluate on simplescaling/aime24_nofigures, problem I-8 as stored there is under-determined and cannot be solved as written. The figure was stripped, and the accompanying text omits the constraint the figure was carrying: that 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.

Restore the condition before testing, for example by using the official wording:

Eight circles of radius 34 can be placed tangent to $B C$ of $△ A B C$ so that the circles are sequentially tangent to each other, with the first circle being tangent to $A B$ and the last circle being tangent to $A C$ . Similarly, 2024 circles of radius 1 can be placed tangent to $B$ in the same manner. […]

simplescaling/aime24_figures retains the Asymptote source, which supplies the same constraint visually and leaves the problem solvable.

Usage

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

python
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=1, 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{} and will look like a wrong answer.

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:

Closed-book sampling. Sample the current model on the bare problem statement. Any output reaching the correct answer becomes a self-generated positive — the target signal.
Hint-conditioned sampling. When step 1 yields nothing (which it does by construction at round 0, since these problems start at 0/64), sample again with a hint appended that supplies the missing step. Correct outputs become positives. Outputs that explicitly reference the hint are discarded, since the stored prompt contains no hint.
Pairing. Positives are paired against closed-book failures from step 1. Both sides are stored against the hint-free prompt, so what the objective sees is a preference between two answers to the same plain question.
Train, then repeat. The stopping criterion is step 1 producing correct answers on its own — the hint scaffolding falling away.

Positives from hint-conditioned generation come from a slightly different policy than the one being trained, so self-generated positives were preferred whenever available.

Compute and the token budget

Table

Sampling / evaluation	RTX 5090, local
DPO training	RunPod: H200, or 2 × RTX PRO 6000
Max sequence length	18500 tokens

18500 is a memory limit, not a modelling choice, and it turned out to matter — see I-11 below. Reasoning traces on these problems routinely run past 10k tokens, so the budget is not generous.

2024 AIME II-9 — chips in a 5×5 grid (solved, 5 rounds)

Five rounds before closed-book sampling produced correct answers unaided. The slowest of the three; the problem's third condition (maximality — that no further chip could be added) is the part the model kept dropping.

2024 AIME I-8 — chain of tangent circles (solved, 2 rounds)

Fastest of the three, but only after correcting the problem statement. As distributed in simplescaling/aime24_nofigures, this problem is missing the constraint that all circles are tangent to the same side (see the note above); the model was failing a question that could not be answered. With the condition restored, two rounds were enough.

This is worth separating out: one of the three "unsolvable" problems was not a capability gap at all. It was a broken prompt.

2024 AIME I-11 — octagon two-colouring (abandoned)

Not solved. After several rounds the model's chains of thought grew past the 18500-token budget, so outputs were being cut off before reaching a final answer, and training could no longer tell "wrong" from "not finished."

Method

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

markdown
L = -log σ( β · [ (log π(y_c|x) - log π_ref(y_c|x)) - (log π(y_r|x) - log π_ref(y_r|x)) ] )
    + λ · ( -log π(y_c|x) / |y_c| )

The NLL term exists to keep the likelihood of the positive example from collapsing, a known failure mode when DPO is left to push both sides down.

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.
Optional length damping (--ld-alpha) in the style of LD-DPO: token log-probabilities past min(|y_c|, |y_r|) are scaled by α. α = 1 is standard DPO.

Hyperparameters

Table

LoRA rank / alpha / dropout	32 / 64 / 0.05
Target modules	`q_proj`, `k_proj`, `v_proj`, `o_proj`, `gate_proj`, `up_proj`, `down_proj`
Trainable parameters	137,625,600
Base precision	4-bit NF4, double quantisation, bf16 compute
β (DPO temperature)	0.1

Evaluation

Table with columns: Benchmark, Base + starting adapter, This adapter, Notes
Benchmark	Base + starting adapter	This adapter	Notes
AIME 2024 (30 problems, 8 samples)	80.00%	86.67%	Difference is the 2 solved problems

Acknowledgements

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

Available on FriendliAI

Dedicated Endpoints

Run this model inference on single tenant GPU with unmatched speed and reliability at scale.

Learn more

Container

Run this model inference with full control and performance in your environment.

Learn more

Model Details

Model Provider

YichengWangCA

Model Tree

Base

deepseek-ai/DeepSeek-R1-Distill-Qwen-14B

Adapter

this model

Input Modalities

Text

Output Modalities

Text

Supported Functionality

Dedicated EndpointsContainer

Explore FriendliAI today

Get started Talk to an engineer

What this adapter demonstrates

DPO + NLL post-training on those three problems recovered two of the three.

Table with columns: Problem, Base model, 64 samples, After post-training
Problem	Base model, 64 samples	After post-training
2024 AIME II-9	0 / 64	solved
2024 AIME I-11	0 / 64	not solved
2024 AIME I-8	0 / 64	solved

How to read the AIME 2024 number

A note for anyone evaluating problem I-8

Restore the condition before testing, for example by using the official wording:

Eight circles of radius 34 can be placed tangent to $B C$ of $△ A B C$ so that the circles are sequentially tangent to each other, with the first circle being tangent to $A B$ and the last circle being tangent to $A C$ . Similarly, 2024 circles of radius 1 can be placed tangent to $B$ in the same manner. […]

simplescaling/aime24_figures retains the Asymptote source, which supplies the same constraint visually and leaves the problem solvable.

Usage

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

python
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=1, 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{} and will look like a wrong answer.

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:

Closed-book sampling. Sample the current model on the bare problem statement. Any output reaching the correct answer becomes a self-generated positive — the target signal.
Hint-conditioned sampling. When step 1 yields nothing (which it does by construction at round 0, since these problems start at 0/64), sample again with a hint appended that supplies the missing step. Correct outputs become positives. Outputs that explicitly reference the hint are discarded, since the stored prompt contains no hint.
Pairing. Positives are paired against closed-book failures from step 1. Both sides are stored against the hint-free prompt, so what the objective sees is a preference between two answers to the same plain question.
Train, then repeat. The stopping criterion is step 1 producing correct answers on its own — the hint scaffolding falling away.

Positives from hint-conditioned generation come from a slightly different policy than the one being trained, so self-generated positives were preferred whenever available.

Compute and the token budget

Table

Sampling / evaluation	RTX 5090, local
DPO training	RunPod: H200, or 2 × RTX PRO 6000
Max sequence length	18500 tokens

18500 is a memory limit, not a modelling choice, and it turned out to matter — see I-11 below. Reasoning traces on these problems routinely run past 10k tokens, so the budget is not generous.

2024 AIME II-9 — chips in a 5×5 grid (solved, 5 rounds)

2024 AIME I-8 — chain of tangent circles (solved, 2 rounds)

This is worth separating out: one of the three "unsolvable" problems was not a capability gap at all. It was a broken prompt.

2024 AIME I-11 — octagon two-colouring (abandoned)

Method

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

markdown
L = -log σ( β · [ (log π(y_c|x) - log π_ref(y_c|x)) - (log π(y_r|x) - log π_ref(y_r|x)) ] )
    + λ · ( -log π(y_c|x) / |y_c| )

The NLL term exists to keep the likelihood of the positive example from collapsing, a known failure mode when DPO is left to push both sides down.

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.
Optional length damping (--ld-alpha) in the style of LD-DPO: token log-probabilities past min(|y_c|, |y_r|) are scaled by α. α = 1 is standard DPO.

Hyperparameters

Table

LoRA rank / alpha / dropout	32 / 64 / 0.05
Target modules	`q_proj`, `k_proj`, `v_proj`, `o_proj`, `gate_proj`, `up_proj`, `down_proj`
Trainable parameters	137,625,600
Base precision	4-bit NF4, double quantisation, bf16 compute
β (DPO temperature)	0.1

Evaluation

Table with columns: Benchmark, Base + starting adapter, This adapter, Notes
Benchmark	Base + starting adapter	This adapter	Notes
AIME 2024 (30 problems, 8 samples)	80.00%	86.67%	Difference is the 2 solved problems

Acknowledgements

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

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

README

What this adapter demonstrates

How to read the AIME 2024 number

A note for anyone evaluating problem I-8

Usage

Training data

The training loop

Compute and the token budget

2024 AIME II-9 — chips in a 5×5 grid (solved, 5 rounds)

2024 AIME I-8 — chain of tangent circles (solved, 2 rounds)

2024 AIME I-11 — octagon two-colouring (abandoned)

Method

Hyperparameters

Evaluation

Acknowledgements

Explore FriendliAI today

README

What this adapter demonstrates

How to read the AIME 2024 number

A note for anyone evaluating problem I-8

Usage

Training data

The training loop

Compute and the token budget

2024 AIME II-9 — chips in a 5×5 grid (solved, 5 rounds)

2024 AIME I-8 — chain of tangent circles (solved, 2 rounds)

2024 AIME I-11 — octagon two-colouring (abandoned)

Method

Hyperparameters

Evaluation

Acknowledgements