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

TL;DR

DAJ addresses distribution shifts in test-time scaling for code generation by introducing a bi-level data-reweighted training framework that learns domain- or instance-level sample weights to optimize generalization on a held-out meta set. It employs a reasoning-based LLM-as-a-Judge with verifiable rewards and supports both preference optimization and reinforcement learning objectives. The framework integrates three reweighting factors—instance difficulty, task similarity, and trajectory alignment—via end-to-end gradients, enabling automatic emphasis of informative data without hand-crafted heuristics. Empirically, DAJ achieves state-of-the-art performance on LiveCodeBench and BigCodeBench across multiple base models, demonstrating strong generalization and robustness to distribution shifts in code-generation tasks.

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.

Figures (6)

• Figure 1: DaJ achieves first place on both LiveCodeBench and BigCodeBench, outperforming other top models. Results are taken from the official leaderboard.
• Figure 2: Bi-level optimization-based data-reweighted training. Top left: Three reweighting designs. Domain reweighting assigns weights to sample groups (e.g., different coding domains); the instance table assigns explicit learnable weights to each sample, while the instance net parameterizes these weights via a lightweight MLP. Bottom left: The bi-level optimization framework. Lower-level optimization updates the judge's parameters using weighted training data, whereas upper-level optimization evaluates the judge on a held-out meta dataset and updates data weights to maximize generalization. Right: A detailed comparison between the lower-level low-quality dataset and the upper-level high-quality dataset.
• Figure 3: Overview of DaJ training and inference. (a)Top left: Given a coding problem, we sample $n$ candidate solutions from the policy model for parallel test-time scaling. Top right: During training, DaJ performs step-by-step reasoning ("Let's think step by step...") before outputting a selection. The model receives a verifiable reward (\ref{['eq:reward']}), enabling preference optimization without human-annotated reasoning traces. Bottom: At inference time, multi-round pairwise voting selects the final output (\ref{['sec:preliminary']}). (b) A simplified example of inference process, where the judge is asked to select the preferred solution from two candidates, based on a 5-step reasoning pattern (\ref{['appendix:prompts']}).
• Figure 4: Comparison with candidate solutions and oracle judge on LiveCodeBench.
• Figure 5: Evolution of learned domain weights during training for domain reweighting.