FunPRM: Function-as-Step Process Reward Model with Meta Reward Correction for Code Generation

기술 요약: FunPRM

문제 정의

코드 생성은 대규모 언어 모델(LLM)의 핵심적인 응용 분야로 남아 있으나, 복잡한 프로그래밍 작업에서는 다단계 추론 요구사항과 환각(hallucination) 현상으로 인해 최첨단 모델들조차 빈번하게 실패하곤 합니다. 수학적 추론에서는 중간 단계를 평가하는 프로세스 보상 모델(Process Reward Models, PRM)과 같은 테스트 시간 스케일링(test-time scaling) 접근 방식이 성공을 거두었지만, 코드 생성에 이를 적용하는 데에는 두 가지 중요한 장애물이 존재합니다:

1. 의미 있는 단계 분해의 부재: Chain-of-Thought(CoT) 프롬프팅 하에서 명시적인 추론 단계로 자연스럽게 분해되는 수학 문제와 달리, 코드는 "단계"에 대한 명확한 정의가 부족합니다. 기존 방식들은 종종 개별 코드 라인을 단계로 취급하며, 이는 프로그램당 수백 개의 단계를 생성하여 과도한 계산 비용을 초래합니다.
2. 노이즈가 섞인 부분 솔루션 보상: PRM을 학습시키려면 부분 솔루션에 대한 정답(ground-truth) 정확도 점수가 필요합니다. 대규모 사용을 위한 인간의 주석 작업은 비용이 너무 많이 들며, 자동화된 몬테카를로(Monte Carlo, MC) 추정 방식은 노이즈가 섞인 보상을 생성합니다. 부분 솔루션이 최종 정답으로 이어지는지를 샘플링하여 도출된 이러한 노이즈 섞인 신호는 PRM의 성능을 저하시킵니다.

방법론: FunPRM

저자들은 단계 정의와 보상 노이즈 문제를 해결하기 위해 코드 생성에 특화된 프로세스 보상 모델인 FunPRM을 제안합니다. 이 모델은 두 가지 주요 메커니즘을 통해 문제를 해결합니다.

1. Chain-of-Function (CoF) 프롬프팅

단계 분해 문제를 해결하기 위해, FunPRM은 LLM이 함수 단위로 조직된 모듈형 코드를 생성하도록 유도하는 새로운 프롬프팅 전략을 채택합니다.

• 함수를 단계로 정의: FunPRM은 코드 라인을 단계로 취급하는 대신, 각 함수를 하나의 추론 단계로 취급합니다.
• 프롬프팅 전략: 시스템 프롬프트는 LLM이 논리적으로 독립적인 연산을 별도의 함수로 추출하고, 이를 탑다운(고수준 함수 우선) 방식으로 구성하며, 설명적인 독스트링(docstring)을 포함하도록 안내합니다.
• 이점: 이를 통해 PRM이 평가할 수 있는 명확하고 의미론적으로 유의미한 추론 단계 시퀀스를 생성하는 동시에, 개발자를 위한 코드 가독성과 재사용성을 향상시킵니다.

2. 메타 학습 기반 보상 교정 (Meta-Learning-Based Reward Correction)

몬테카를로 추정된 부분 솔루션 보상의 노이즈를 해결하기 위해, FunPRM은 코딩 작업의 고유한 특성인 '유닛 테스트 평가 시스템을 통해 최종 솔루션에 대한 깨끗하고 신뢰할 수 있는 정확도 점수를 얻을 수 있다'는 점을 활용하는 이중 레벨 메타 학습 프레임워크를 도입합니다.

• 노이즈 섞인 신호 vs 깨끗한 신호: 이 프레임워크는 (MC 샘플링을 통해 추정된) 노이즈가 섞인 부분 솔루션 데이터셋과 (유닛 테스트를 통해 평가된) 깨끗한 최종 솔루션 메타 데이터셋을 모두 활용합니다.
• 보상 교정 테이블: 노이즈가 섞인 부분 솔루션 보상에 학습 가능한 잔차(residual)를 더하기 위해 경량화되고 학습 가능한 보상 교정 테이블($g_\theta$)을 도입합니다.
• 이중 레벨 최적화:
  1. 내부 루프(Inner Loop): 현재 교정된 보상을 사용하여 노이즈가 섞인 부분 솔루션에 대해 PRM($f_\phi$)을 학습시킵니다. PRM 파라미터는 단일 경사 하강 단계(gradient descent step)를 통해 업데이트됩니다.
  2. 외부 루프(Outer Loop): 업데이트된 PRM을 깨끗한 최종 솔루션 데이터셋에서 평가합니다. 결과로 나온 손실(meta-loss)은 보상 교정 파라미터 $\theta$에 대한 그래디언트를 계산하는 데 사용됩니다.
• 목표: 이 과정은 교정 테이블을 최적화하여 노이즈가 섞인 부분 솔루션 보상을 "정화(purify)"함으로써, 해당 보상으로 학습된 PRM이 성공적인 최종 솔루션을 올바르게 식별하는 데 잘 일반화되도록 보장합니다.

주요 기여

1. FunPRM 프레임워크: Chain-of-Function 프롬프팅을 통해 추론 단계를 함수 수준으로 재정의함으로써, 라인 단위의 세밀함에 따른 계산 오버헤드 없이 효과적인 단계별 평가를 가능하게 하는 코드 생성 특화 PRM을 제안했습니다.
2. 메타 보상 교정: 유닛 테스트로부터 얻은 깨끗한 최종 솔루션 보상을 사용하여 몬테카를로 추정된 부분 솔루션 보상의 노이즈를 제거하는 새로운 메커니즘을 도입하였으며, 이를 통해 인간의 주석 없이도 학습 신호의 품질을 개선했습니다.
3. 경험적 우월성: 다양한 베이스 LLM과 벤치마크에 대해 FunPRM이 기존의 테스트 시간 스케일링 베이스라인(Self-Certainty, Outcome Reward Models, Skywork-PRM 포함)을 일관되게 능가함을 입증했습니다.

실험 결과

저자들은 LiveCodeBench 및 BigCodeBench에서 FunPRM을 평가하였으며, HumanEval+ 및 **MBPP+**를 통해 도메인 일반화 성능을 검증했습니다.

• 최첨단 성능: OpenAI의 **O4-mini (High)**와 결 조합되었을 때, FunPRM은 LiveCodeBench에서 **pass@1 80.9%**를 달성하며 O3 (High), DeepSeek-R1, Gemini-2.5-Pro를 포함한 리더보드의 모든 모델을 넘어섰습니다.
• 모델 간 일관성: FunPRM은 다섯 가지 서로 다른 베이스 LLM(Qwen3-30B, GPT-4o-mini, DeepSeek-Coder, Qwen2.5-7B, O4-mini)에 대해 LiveCodeBench와 BigCodeBench 모두에서 모든 베이스라인 테스트 시간 스케일링 방법을 능가했습니다.
• 절제 연구(Ablation Studies): Chain-of-Function 프롬프팅이나 메타 보상 교정 메커니즘 중 하나라도 제거할 경우 성능 저하가 발생하였으며, 이는 두 구성 요소의 필요성을 확인시켜 줍니다.
• 인간 평가: 3명의 연구자가 참여한 연구에서, FunPRM이 생성한 코드는 (정답 코드의 경우) 가독성 측면에서, (오답 코드의 경우) 재사용성 및 수정 용이성 측면에서 베이스라인 코드보다 일관되게 선호되었습니다.
• 도메인 일반화: LiveCodeBench와 BigCodeBench에서 학습된 FunPRM은 추가 학습 없이도 HumanEval+와 MBPP+에서 성능을 성공적으로 향상시켰으며, 이는 다양한 코딩 스타일과 난이도에 걸친 강력한 일반화 능력을 나타냅는 것입니다.

의의

본 논문은 FunPRM이 LLM 기반 코드 생성의 테스트 시간 스케일링을 개선하는 실질적이고 효과적인 접근 방식임을 주장합니다. 추론 단계의 정의를 소프트웨어 공학의 자연스러운 모듈성(함수)과 일치시키고, 유닛 테스트의 결정론적 특성을 활용하여 학습 신호를 정제함으로써, FunPRM은 그동안 PRM 도입을 저해했던 특정 한계점들을 극복했습니다. 이 결과는 본 접근 방식이 이미 높은 성능을 보이는 모델들조차도 코드 생성의 성능 한계를 실질적으로 밀어올릴 수 있으며, 동시에 개발자에게 더 읽기 쉽고 재사용 가능한 코드를 제공할 수 있음을 시사합니다.