DAJ: Data-Reweighted LLM Judge for Test-Time Scaling in Code Generation
This addresses distribution shift issues in LLM-as-a-Judge training for code generation, offering an incremental improvement over existing methods.
The paper tackles the challenge of training reliable LLM judges for test-time scaling in code generation by proposing DAJ, a reasoning-based LLM judge trained with verifiable rewards under a bi-level data-reweighted learning framework, achieving state-of-the-art performance on LiveCodeBench and BigCodeBench.
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.