# DAJ: Data-Reweighted LLM Judge for Test-Time Scaling in Code Generation

arXiv: https://arxiv.org/abs/2601.22230 · HTML: https://mlsanity.com/paper/2601.22230/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Reasoning for agents, planning and code, one of 17 directions in LLM reasoning.
- Builds on 24 of its 54 references that are in the map.
- First posted on arXiv 29 Jan 2026.

## Placed in

- LLM reasoning / [Reasoning for agents, planning and code](https://mlsanity.com/lr/reasoning-for-agents-planning-and-code/)

## Builds on (most cited first, 24 of 24)

- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [Competition-Level Code Generation with AlphaCode](https://mlsanity.com/paper/2203.07814/) (2022)
- [OpenAI o1 System Card](https://mlsanity.com/paper/2412.16720/) (2024)
- [Qwen2.5-Coder Technical Report](https://mlsanity.com/paper/2409.12186/) (2024)
- [Large Language Monkeys: Scaling Inference Compute with Repeated Sampling](https://mlsanity.com/paper/2407.21787/) (2024)
- [Math-Shepherd: Verify and Reinforce LLMs Step-by-step without Human Annotations](https://mlsanity.com/paper/2312.08935/) (2023)
- [Show Your Work: Scratchpads for Intermediate Computation with Language Models](https://mlsanity.com/paper/2112.00114/) (2021)
- [ORPO: Monolithic Preference Optimization without Reference Model](https://mlsanity.com/paper/2403.07691/) (2024)
- [BigCodeBench: Benchmarking Code Generation with Diverse Function Calls and Complex Instructions](https://mlsanity.com/paper/2406.15877/) (2024)
- [Improve Mathematical Reasoning in Language Models by Automated Process Supervision](https://mlsanity.com/paper/2406.06592/) (2024)
- [The Lessons of Developing Process Reward Models in Mathematical Reasoning](https://mlsanity.com/paper/2501.07301/) (2025)
- [Self-critiquing models for assisting human evaluators](https://mlsanity.com/paper/2206.05802/) (2022)
- [KTO: Model Alignment as Prospect Theoretic Optimization](https://mlsanity.com/paper/2402.01306/) (2024)
- [Inference-Time Scaling for Generalist Reward Modeling](https://mlsanity.com/paper/2504.02495/) (2025)
- [A Survey on Test-Time Scaling in Large Language Models: What, How, Where, and How Well?](https://mlsanity.com/paper/2503.24235/) (2025)
- [Meta-Rewarding Language Models: Self-Improving Alignment with LLM-as-a-Meta-Judge](https://mlsanity.com/paper/2407.19594/) (2024)
- [Self-Taught Evaluators](https://mlsanity.com/paper/2408.02666/) (2024)
- [Rewarding Chatbots for Real-World Engagement with Millions of Users](https://mlsanity.com/paper/2303.06135/) (2023)
- [Evaluating Judges as Evaluators: The JETTS Benchmark of LLM-as-Judges as Test-Time Scaling Evaluators](https://mlsanity.com/paper/2504.15253/) (2025)
- [Process Supervision-Guided Policy Optimization for Code Generation](https://mlsanity.com/paper/2410.17621/) (2024)
- [PairJudge RM: Perform Best-of-N Sampling with Knockout Tournament](https://mlsanity.com/paper/2501.13007/) (2025)

## Built on by (most cited first, 1 of 1)

- [Verification Trap: Understanding Test-Time Selection Failures under False Premises in Code Generation](https://mlsanity.com/paper/2610.05170/) (2026)

## Abstract (arXiv)

Test-time scaling for code generation commonly relies on Best-of-N selection, in which multiple candidate solutions are sampled from a base model, and the best one is selected by an LLM judge. However, training reliable LLM judges is challenging due to severe distribution shifts, including imbalances between easy and hard problems, mismatches between training tasks and evaluation benchmarks, and trajectory mismatch arising from training data generated by cheaper models whose behavior differs from that of inference-time models. We propose DAJ, a reasoning-based LLM judge trained with verifiable rewards under a bi-level data-reweighted learning framework. The proposed framework learns data-importance weights (either domain-level or instance-level) to optimize generalization performance on a held-out meta set aligned with target benchmarks. To the best of our knowledge, this is the first application of data reweighting to LLM-as-a-Judge training for test-time scaling. Our approach automatically emphasizes hard problems, in-distribution samples, and trajectory-aligned data, without relying on hand-crafted heuristics. Empirically, DAJ achieves state-of-the-art performance on LiveCodeBench and BigCodeBench, outperforming strong test-time scaling baselines as well as leading proprietary models.