코드 생성 테스트‑타임 스케일링을 위한 데이터 재가중 LLM 심판 DAJ

초록

DAJ는 베스트‑오브‑N 방식에서 후보 코드를 평가하는 LLM 심판을 훈련할 때, 문제 난이도·도메인·생성 궤적의 세 가지 분포 이동을 자동으로 보정하는 데이터 재가중 bi‑level 최적화 프레임워크를 제안한다. 검증 가능한 보상을 이용해 선호 최적화 혹은 강화학습으로 학습하고, LiveCodeBench와 BigCodeBench에서 기존 방법과 상용 모델을 앞선 성능을 기록한다.

상세 분석

본 논문은 코드 생성 분야에서 테스트‑타임 스케일링(Test‑Time Scaling, TTS) 기법 중 베스트‑오브‑N 전략을 강화하기 위한 “LLM‑as‑a‑Judge” 모델의 학습 문제를 체계적으로 분석한다. 기존 심판 모델은 (1) 쉬운 문제에 과다 비중이 몰려 어려운 문제에 대한 일반화가 떨어지고, (2) 훈련에 사용된 과제와 실제 평가 벤치마크 간의 도메인 불일치, (3) 훈련 데이터가 저품질 모델이 생성한 후보를 기반으로 하는 궤적 불일치라는 세 가지 주요 분포 이동에 취약했다.

DAJ는 이러한 문제를 해결하기 위해 데이터 재가중(data‑reweighting) 을 핵심 설계 요소로 채택한다. 구체적으로, 각 훈련 샘플에 중요도 가중치를 할당하는데, 이 가중치는 (i) 난이도 가중치—문제의 해결 난이도를 자동 추정해 어려운 샘플에 높은 가중치를 부여, (ii) 도메인 유사도 가중치—훈련 과제가 목표 벤치마크와 얼마나 유사한지를 측정해 유사한 도메인에 가중치를 부여, (iii) 궤적 정렬 가중치—후보 코드가 테스트‑타임에 강력한 모델이 생성할 가능성이 높은지를 평가해 해당 샘플을 강조한다.

가중치 학습은 bi‑level 최적화 구조로 수행된다. 하위 레벨에서는 현재 가중치에 따라 재가중된 손실(선호 최적화 혹은 강화학습 기반 RL‑VR)을 최소화하며, 상위 레벨에서는 별도로 확보한 메타 데이터셋(고품질, 목표 도메인에 가깝게 설계)에서의 성능을 평가해 가중치를 업데이트한다. 이렇게 하면 메타셋의 성능 향상이 직접 가중치 조정에 반영되어, 분포 이동에 대한 적응이 자동화된다.

모델 아키텍처는 reasoning‑based LLM‑as‑Judge 방식을 채택한다. 입력으로 문제 설명과 n개의 후보 코드를 받고, “Let’s think step by step” 프롬프트를 통해 단계별 논리 전개를 생성한 뒤 최종 선택을 출력한다. 후보 코드는 실행 가능한 테스트 케이스를 통해 자동 검증 가능하므로, 검증 가능한 보상 함수 R을 정의한다: 올바른 선택이면 1, 오답이면 0.5, 포맷 오류이면 0. 이 보상은 인간 라벨이 필요 없는 RL‑VR(Verification‑Reward) 학습을 가능하게 하며, 선호 최적화에서는 올바른 선택을 ‘positive’, 나머지를 ‘negative’로 변환해 pairwise loss를 적용한다.

실험에서는 LiveCodeBench와 BigCodeBench 두 대규모 코드 평가 벤치마크에서 DAJ가 기존 베스트‑오브‑N 기반 심판(예: OpenAI gpt‑4, Claude 2) 및 최신 상용 TTS 시스템을 능가한다. 특히, 어려운 문제(Hard split)와 도메인 전이 상황에서 가중치가 자동으로 조정되는 효과가 두드러졌다. Ablation study는 (1) 도메인‑레벨 vs 인스턴스‑레벨 가중치, (2) 난이도 vs 궤적 가중치 각각의 기여도를 분석해, 세 가지 가중치가 모두 결합될 때 가장 큰 성능 향상이 발생함을 확인한다.

이 논문의 주요 기여는 (1) LLM‑as‑Judge 학습에 데이터 재가중 bi‑level 최적화를 최초로 도입해 분포 이동을 원칙적으로 보정한 점, (2) 인간 라벨 없이 검증 가능한 보상을 활용해 효율적인 강화학습/선호 최적화를 구현한 점, (3) 다양한 베이스 모델과의 호환성을 입증하고, 실제 코드 생성 시스템에 바로 적용 가능한 상용 수준의 성능을 달성한 점이다. 향후 연구는 (i) 가중치 학습에 메타‑러닝 기법을 결합해 더 빠른 수렴을 도모하고, (ii) 다중 후보를 동시에 평가하는 멀티‑플레이어 판정 메커니즘을 확장하며, (iii) 다른 생성 도메인(예: 수학, 자연어)에도 동일한 프레임워크를 적용하는 가능성을 탐색할 수 있다.