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

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Qin, Peijia, Zhang, Ruiyi, Cao, Qi, Xie, Pengtao
Formato: Preprint
Publicado: 2026
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866917233217765376
author Qin, Peijia
Zhang, Ruiyi
Cao, Qi
Xie, Pengtao
author_facet Qin, Peijia
Zhang, Ruiyi
Cao, Qi
Xie, Pengtao
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2601_22230
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle DAJ: Data-Reweighted LLM Judge for Test-Time Scaling in Code Generation
Qin, Peijia
Zhang, Ruiyi
Cao, Qi
Xie, Pengtao
Machine Learning
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.
title DAJ: Data-Reweighted LLM Judge for Test-Time Scaling in Code Generation
topic Machine Learning
url https://arxiv.org/abs/2601.22230