Show Your Work: Gemma2 Process Rewards

I submitted this for Google Tunix Hack (Train a model to show its work, Jan 2026). The brief was not just “get the answer right.” The model should show how it got there.

What I built#

A 4-tier hybrid reward system for Gemma2 2B-IT, trained with GRPO and LoRA on a TPU v5e-8, all inside one ~9-hour session.

The bet: score the reasoning process with structured LLM-as-judge rubrics, not answer correctness alone.

Reward tiers#

1. Format compliance (+0.5). Binary check for <reasoning>…</reasoning> / <answer>…</answer>. Kept weak so the model cannot farm tags.
2. Reasoning quality (±1.0). Gemini 2.0 Flash as judge under a Rubrics as Rewards setup:
   • Logical coherence (30%)
   • Completeness (25%)
   • Groundedness (20%)
   • Non-circularity (15%)
   • Clarity (10%)
3. Answer correctness (±1.0). Hybrid:
   • math/science: string match or numeric epsilon
   • open-ended: Gemini on relevance, completeness, clarity, helpfulness, tone
4. Conciseness (±1.0). Word-count bands that favor focused traces under a ~1K token budget.

Scores land in roughly [-1, 1] for GRPO so weak process actually hurts.

Data#

5k prompts, about 70% non-verifiable / 30% verifiable. That matched the competition focus and what a 2B model is good at:

Training stack#

• Model: Gemma2 2B-IT + LoRA (rank 64, alpha 128) on attention and MLP paths
• Algo: GRPO, 4 generations per prompt, KL beta 0.04, advantage clip ε 0.2
• Opt: AdamW, warmup cosine (peak lr 3e-6), grad clip 0.1
• Hardware: TPU v5e-8, FSDP + tensor parallel mesh
• Framework: Tunix + JAX / Flax NNX on a Kaggle TPU session

Infrastructure that mattered#

LLM-as-judge is expensive and flaky if you are careless. The eval path used:

• ThreadPoolExecutor with semaphore-limited concurrency
• MD5 cache (high hit rate mid-run, big API cost cut)
• exponential backoff retries
• temperature 0.0 for low judge variance

Early sanity runs showed classic reward gaming: perfect tags, shallow reasoning. Fixes were a weaker format weight, flexible regex (.search instead of anchored .match), a minimum reasoning quality before positive format payoff, and tighter judge prompts.

Takeaways#

Process rewards work when format is a soft scaffold and quality and correctness carry the real weight. Caching and batching make Gemini-as-judge usable inside a fixed TPU budget. Hybrid deterministic plus LLM correctness covers math and open domains without forcing one evaluator everywhere.