aestudio/gemma-4-31B-it-GRPO-math-lora

gemma-4-31B-it + GRPO on DeepMath (LoRA adapter, generation 39)

A LoRA adapter over google/gemma-4-31B-it (instruct, thinking mode off), trained with GRPO and a verifiable correctness reward on DeepMath-103K.

On 1,000 held-out DeepMath problems it scores 0.758 pass@1 against the base model's 0.629, a paired difference of +0.129 [+0.108, +0.151], with instruction following (IFBench) unchanged. The held-out problems are ones the base model sometimes solves but often doesn't (see Results). Absolute scores are not representative of DeepMath as a whole, where the base model already solves about 88 %.

An earlier checkpoint of the same run, closer to the base, is also published: aestudio/gemma-4-31B-it-GRPO-math-lora-gen10. It has about half the gain (+0.063 [+0.044, +0.081] on the same test), with base-length answers.

This is an adapter, not a merged model. You need the base model above.

---

⚠️ Merge before serving with vLLM, and keep thinking mode off

vLLM's runtime LoRA under-applies this adapter. On vLLM 0.28 with Gemma 4, serving it through LoRARequest recovers only ~75–80 % of the adapter's effect, and the shortfall sits in the attention projections. Every number on this card was measured with the adapter merged into the base weights (bf16, merge_and_unload()), which matches the trainer's forward pass to bf16 rounding. With vLLM, merge first. With plain transformers + PEFT, no change is needed.

Thinking mode was never enabled, in training or in evaluation. Render every prompt with enable_thinking=False. The adapter has not been trained or measured with Gemma 4's thought channel and should not be expected to behave sensibly with it.

Stack this was trained and evaluated on: vllm 0.28.0 · transformers 5.17.0 · trl 1.13.0 · peft 0.20.0 · torch 2.13.0+cu130.

---

Usage

from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

base = AutoModelForCausalLM.from_pretrained("google/gemma-4-31B-it", dtype="bfloat16")
model = PeftModel.from_pretrained(base, "aestudio/gemma-4-31B-it-GRPO-math-lora").merge_and_unload()
tok = AutoTokenizer.from_pretrained("google/gemma-4-31B-it")

messages = [{"role": "user", "content": problem + "\n\nPut your final answer within \\boxed{}."}]
prompt = tok.apply_chat_template(messages, add_generation_prompt=True, tokenize=False, enable_thinking=False)

Prompt exactly as it was trained and evaluated: a single user turn containing the problem, followed by a blank line and

Put your final answer within \boxed{}.

No system prompt. The adapter touches only the text decoder (q/k/v/o/gate/up/down_proj under language_model), so the vision tower is the base model's.

Results

All reads are paired on identical problems with a bootstrap 95 % CI (10,000 resamples over problems). They use n = 4 samples per problem at Google's recommended sampling for this model (temperature 1.0, top_p 0.95, top_k 64), a 16,384-token cap, and thinking off. A completion is correct if its last \boxed{} matches the gold under math-verify, forgiving whitespace only. Truncated, unboxed or looping completions score 0.

Which problems. DeepMath problems were filtered by the base model's own success. Each was sampled 8 times at temperature 1.0, and only problems it solved 1 to 7 times out of 8 were kept. The rest it either already solves every time (the large majority) or never solves, so there is nothing for reinforcement learning to work with. The filtered set was split into a training pool (6,247), a 1,000-problem selection set and a 1,000-problem test set. The base model's scores below come from fresh samples, not from the samples used to filter.

Held-out test, 1,000 problems (read once, after the checkpoint was chosen):

	pass@1	paired Δ vs base
gemma-4-31B-it (base)	0.629	—
this adapter (generation 39)	0.758	+0.129 [+0.108, +0.151]

Checkpoint selection, a separate 1,000 problems, same protocol:

generation	base	10	20	30	39
pass@1	0.619	0.697	0.732	0.749	0.761
paired Δ vs base	—	+0.078	+0.113	+0.130	+0.142 [+0.119, +0.165]

Gains slowed with each read: generation 30 over 20 was +0.018 [+0.004, +0.032], and 39 over 30 was +0.012 [−0.001, +0.025]. Training was stopped at generation 39, and the best selection read was taken. Per problem, the adapter loses 28 test problems the base solved at least 3 of 4 times, and gains 117 it solved at most once.

Instruction following was not traded away

IFBench (300 prompts, 58 out-of-distribution verifiable constraints), graded programmatically (no LLM judge), n = 4 with the same sampling, paired against the base. The regression threshold was fixed before any checkpoint was scored: a regression counts only if the paired Δ's CI upper bound falls below −0.02.

	prompt-level strict	paired Δ	prompt-level loose	paired Δ
base	0.499	—	0.546	—
generation 10	0.505	+0.006 [−0.007, +0.020]
generation 20	0.510	+0.011 [−0.004, +0.028]
generation 30	0.503	+0.003 [−0.013, +0.020]
this adapter (generation 39)	0.512	+0.013 [−0.004, +0.029]	0.554	+0.008 [−0.008, +0.026]

No checkpoint regresses. IFBench responses are the same length as the base's (paired mean difference +2 tokens [−51, +56], median 127 vs 126 tokens), and no \boxed{} leaks into them.

These are n = 4 at temperature 1.0, not IFBench's usual protocol. They are not comparable to published leaderboard scores; they compare these checkpoints to each other.

Response length on math

Math answers are longer than the base's: mean 1,511 vs 1,261 tokens on the test set (+20 %), median 1,131 vs 1,005. Most of the increase is in answers that end up wrong (median 1,236 vs 948 tokens). Correct answers grew only modestly (1,110 vs 1,033). Truncation stays rare (0.15 % vs 0.05 %), and unboxed answers fall (0.02 % vs 0.10 %).

Checking your setup

Score the base model and this adapter on any shared math set, with thinking off and the prompt above. If the adapter is not clearly ahead, check first that it was merged, and second that the chat template was rendered with enable_thinking=False. A direct check: an HF forward pass of the merged model should agree with vLLM ≥ 0.28's sampler to about 0.01 nats/token on the same tokens. Trainer and sampler agreed to 0.005–0.009 throughout training.

Training

algorithm	GRPO (TRL 1.13 with a colocated vLLM 0.28 engine, tensor-parallel 2, sleep mode), DAPO token-level loss, β = 0, ε 0.2 / 0.28, truncated importance-sampling correction
adapter	LoRA r = 128, α = 256, dropout 0, text-decoder projections only, bf16 (~3.9 GB)
optimiser	constant lr 1e-5, unscaled advantages
rollouts	64 prompts × 8 samples = 512 per generation, temperature 1.0, top_p 1.0, no top_k, 8,192-token cap, thinking off
run	39 generations = 19,968 rollouts, 156 optimiser steps, ~14 h on 8×H200
reward	1.0 if the final \boxed{} answer matches the gold under math-verify (whitespace forgiven), else 0.0. No format shaping, no length term. Truncated completions are unmasked and score 0
data	6,247-problem pool drawn from a decontaminated DeepMath-103K subset, filtered to problems the base solves 1–7 of 8 times. 2,496 distinct problems seen, 64 per generation, each once

Rollouts are sampled at temperature 1.0 with no top-p or top-k truncation. That is deliberate, even though Google recommends top_p 0.95 / top_k 64 for generation: GRPO differentiates the full-vocabulary distribution, and truncated sampling would make the rollouts off-policy for it. The learning rate was chosen from a 10-generation comparison of 1e-5 and 5e-6. 1e-5 was ahead by +0.028 [+0.013, +0.044] at generation 10 (that comparison was read at temperature 0.6), and this run continues that arm. Trainer-and-sampler agreement was gated throughout and held at ≤ 0.009 nats/token. Gemma 4's thought-channel tokens were also suppressed in the sampler (logit bias −100) and checked for in every completion.

Exactly which problems it trained on

Every problem the adapter trained on is listed in this repo as training_problems.jsonl. It has 2,496 rows, one per problem: generation (1–39), deepmath_row (the row index in DeepMath-103K's train split), problem, answer, difficulty, topic, and base_correct_of_8 (how many of 8 base-model samples at temperature 1.0 were correct). None of them appear in the selection or test sets, so seen/unseen splits of the pool against this model can be drawn exactly. The list was rebuilt after training by replaying the trainer's seeded data order, because the run did not log it directly. The replay is checked against the run's own logs. With the recovered problems, the base model's recorded success rates predict the per-generation training reward (detrended correlation 0.74), and none of 2,000 alternative orderings comes close (95th percentile 0.26). A rerun of the trainer that logged the problems directly matched the replay exactly, all 64 problems in each generation checked. It covered generations 1 and 2 from scratch, and generations 21 and 22 resumed from a saved checkpoint.

Limitations

• Filtered difficulty. The held-out sets contain only problems the base model gets right sometimes but not always. The gain measures how much RL helps where there is room to learn, not DeepMath accuracy overall.
• One seed, one recipe, 39 generations. Gains were still positive but slowing when training stopped.
• Retention covers one axis. Instruction following (IFBench) was measured; nothing else was.
• Math only. No open-ended quality, multilingual, coding, multimodal, or safety evaluation.
• Thinking mode off only. Behaviour with Gemma 4's thinking mode is unknown.
• Adapter, not a model. Requires the base weights. With vLLM, merge it first (see above).

License

The adapter follows the base model's license (Apache-2.0, google/gemma-4-31B-it). DeepMath-103K is MIT.