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

TL;DR

DAJ introduces a novel data-reweighted training framework for LLM judges that improves test-time code generation by emphasizing challenging and in-distribution samples, leading to state-of-the-art results.

Contribution

This work is the first to apply data reweighting to train LLM judges, enhancing their ability to select high-quality code solutions during test-time scaling.

Findings

DAJ outperforms existing baselines on LiveCodeBench and BigCodeBench.

It automatically emphasizes hard and in-distribution problems.

Achieves state-of-the-art performance with strong models.

Abstract

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.