PAPO: 확산 LLM 추론을 위한 보상과 상태 정렬 강화학습
목차
개요
이 글은 논문 “Back on Track: Aligning Rewards and States for Reasoning in Diffusion Large Language Models”를 정리한 것이다. 확산 대규모 언어 모델(diffusion LLM, 이하 dLLM)은 토큰을 하나씩 예측하는 자기회귀(AR) 모델과 달리 반복적 디노이징(denoising)과 리마스킹(remasking) 루프로 여러 토큰을 병렬 생성한다. 이 병렬 생성 덕분에 추론 속도가 빠르지만, AR 모델 최적화의 핵심 기술인 강화학습(RL)을 비자기회귀 구조에 적용하는 방법은 아직 충분히 연구되지 않았다.
논문은 기존 dLLM용 RL 프레임워크가 두 가지 근본적인 정렬 불일치(misalignment)에 의해 제약된다고 지적한다. 첫째는 과정 보상 불일치(process-reward misalignment)다. 최종 결과에서 나온 희소한 보상(sparse terminal reward)이 생성 과정의 모든 중간 단계에 무차별적으로 할당되어, 단계별로 차별화된 크레딧 할당(credit assignment)을 제공하지 못한다. 둘째는 상태-궤적 불일치(state-trajectory misalignment)다. 정책 업데이트가 최종 완성본을 무작위로 마스킹해 만든 인위적인 상태(artificial state)에서 수행되어, 디노이징 과정의 실제 중간 맥락을 충실히 반영하지 못한다.
이 두 불일치를 해결하기 위해 저자들은 Process Aligned Policy Optimization(PAPO)을 제안한다. PAPO는 두 가지 모듈로 구성된다. 하나는 희소한 최종 보상을 단계별 밀집 보상으로 바꾸는 Step-Aware Process Rewards(SPR)이고, 다른 하나는 불확실성이 높은 단계에서 실제 궤적을 재현하는 Entropy-Guided Historical Re-enactment(EHR)다. 네 가지 벤치마크 실험에서 PAPO는 GSM8K에서 최대 4.5%, MATH500에서 4.8%, Countdown에서 42.2%, Sudoku에서 16.1%의 절대 성능 향상을 달성했다.
방법론
PAPO는 RL 업데이트 과정을 dLLM의 실제 생성 궤적과 정렬하는 것을 목표로 한다. 먼저 마스크 확산 dLLM의 전제를 정리한다. 순방향 과정은 깨끗한 토큰 시퀀스를 시간 t에 따라 확률적으로 마스크 토큰으로 치환한다. 모델은 마스킹된 토큰을 예측하도록 학습하며, LLaDA처럼 노이즈 스케줄 alpha_t = 1 - t를 사용하는 경우 손실은 마스킹된 토큰에 대한 항으로 축소된다.
정책 최적화에는 critic이 필요 없는 GRPO(Group Relative Policy Optimization)를 사용한다. 프롬프트 q에 대해 G개의 응답 그룹을 샘플링하고, 각 응답의 보상으로부터 그룹 상대 어드밴티지(advantage)를 계산한다. 어드밴티지는 해당 응답의 보상에서 그룹 평균 보상을 뺀 값으로 정의된다. 다만 dLLM은 비인과적 다단계 디노이징 때문에 AR 모델처럼 연쇄 법칙으로 시퀀스 로그 우도를 구할 수 없어, 평균장 근사(mean-field approximation)로 토큰별 주변 확률의 합으로 근사한다.
Step-Aware Process Rewards (SPR)
기존 RL 방법의 핵심 한계는 과정 보상 불일치다. 최종 결과 하나에 기반한 희소 보상은 중간 추론 경로의 품질에 대한 감독 신호를 전혀 주지 못한다. SPR은 디노이징 궤적의 각 중간 단계를 평가하는 밀집(dense) 보상을 생성해 이 문제를 해결한다.
dLLM은 완전히 마스킹된 상태에서 시작해 각 단계 t에서 마스킹된 토큰을 예측하고 남은 토큰을 다시 마스킹하며 최종 완성본까지 정제한다. SPR은 각 롤아웃의 전체 생성 궤적을 기록한다. 각 단계 t에서 리마스킹 직전의 맥락으로부터 한 단계 디노이징 예측(one-step denoised prediction)을 만든다. 이 예측은 현재 맥락에서 추정한 완전한 해답이며, 여기에 보상 함수를 적용해 단계별 과정 보상을 계산한다.
이 즉각적 보상을 장기 목표와 연결하기 위해, 단계별 과정 보상을 최종 결과 보상과 합쳐 종합적인 단계별 총 보상을 만든다. 이 종합 보상으로부터 더 정확한 단계별 어드밴티지를 계산해 궤적과 함께 캐시에 저장한다. 이렇게 하면 정밀하고 단계별로 차별화된 크레딧 할당이 가능해진다.
Entropy-Guided Historical Re-enactment (EHR)
두 번째 불일치는 정책 업데이트를 위한 비효율적이고 불충실한 상태 선택이다. diffu-GRPO는 무작위로 마스킹된 프롬프트와 완전히 마스킹된 완성본으로 구성된 합성 상태에서 업데이트를 수행한다. 이는 모델이 초기화 이후 결코 마주치지 않는 무정보(zero-information) 맥락이다. UniGRPO는 최종 완성본에 무작위 마스크 비율을 적용하지만, 이 역시 실제 생성 궤적과 일치할 가능성이 낮은 비현실적 맥락을 만든다. 두 경우 모두 학습에 쓰이는 합성 맥락과 추론 시 실제로 마주치는 상태 사이에 분포 불일치가 발생한다.
EHR은 이 문제를 해결하기 위해 자연스러운 롤아웃 과정에서 캐시된 실제 궤적을 활용한다. 실제 중간 상태에서 직접 정책 업데이트를 수행해 학습-추론 분포 불일치를 제거한다. 다만 균등 샘플링은 이미 잘 학습된 저엔트로피 단계를 반복적으로 다시 방문하기 때문에 비효율적이다. EHR은 대신 정책의 불확실성이 가장 큰 고엔트로피 상태를 우선시한다.
구체적으로 EHR은 각 시점 t에 대해 마스킹된 위치들의 토큰 수준 엔트로피를 평균한 단계 엔트로피를 계산한다. 그리고 단계 tn을 선택할 확률이 그 엔트로피에 비례하도록 시점 샘플링 분포를 구성한다. 이때 sharpness를 조절하는 하이퍼파라미터를 사용하는데, 값이 0이면 균등 샘플링이 되고 값이 커질수록 고불확실성 단계를 더 선호한다. 각 업데이트 반복마다 이 분포에서 tn을 샘플링하고, 캐시된 궤적으로부터 실제 맥락을 재구성해 정책을 업데이트한다.
PAPO 학습 목적함수
PAPO는 SPR이 유도한 어드밴티지와 EHR의 상태 선택을 통합한다. 전체 목적함수는 엔트로피 기반 샘플링 분포에서 추출한 단계 tn에 대해, 마스킹된 토큰 집합에 걸쳐 중요도 비율(importance ratio)과 캐시된 단계별 어드밴티지의 곱을 클리핑(clipping)한 값을 최소화하는 형태다. 여기에 참조 모델에 대한 KL 정규화 항이 더해진다. 중요도 비율은 새 정책과 이전 정책에서 평가한 토큰 확률의 비로 정의되며, clipping 폭과 KL 정규화 강도는 각각 별도의 하이퍼파라미터로 제어한다.
전체 학습 알고리즘은 다음과 같은 흐름으로 진행된다.
| 단계 | 내용 |
|---|---|
| 롤아웃과 궤적 기록 | 이전 정책으로 G개 응답을 샘플링하고 전체 궤적을 캐시 |
| SPR 계산 | 각 단계에서 한 단계 디노이징으로 과정 보상과 어드밴티지 산출 |
| EHR 업데이트 | 엔트로피에 비례해 단계를 샘플링하고 실제 맥락을 재구성해 업데이트 |
주요 결과
네 가지 벤치마크 성능
실험은 네 가지 추론 과제에서 진행되었다. GSM8K는 다단계 논리 추론이 필요한 초등 수학 데이터셋, MATH500은 경시대회 수준 고교 수학 500문제, Countdown은 목표 값을 만드는 산술 조합 게임, 4x4 Sudoku는 제약 충족 계획 과제다. 모든 실험은 LLaDA-8B-Instruct 모델을 기반으로 하며, d1과 달리 PAPO는 지도 미세조정(SFT) 없이 베이스 모델에 직접 RL을 적용한다. 온라인 궤적 생성 시퀀스 길이는 256 토큰으로 고정하고, 평가는 128/256/512 토큰에서 수행해 일반화를 측정했다. 모든 실험은 8장의 NVIDIA A100 GPU에서 진행되었다.
SFT 없이도 PAPO는 LLaDA-8B-Instruct 대비 평균 16.9%의 절대 성능 향상을 달성했다. 이득은 계획 과제에서 특히 두드러져 Countdown 42.2%, Sudoku 16.1%에 달했다. 주요 벤치마크별 베이스라인 대비 PAPO의 성능은 다음 표와 같다.
| 모델 / 시퀀스 길이 | GSM8K 256 | MATH500 256 | Countdown 256 | Sudoku 256 |
|---|---|---|---|---|
| LLaDA-8B-Instruct | 76.7 | 32.4 | 19.5 | 6.7 |
| diffu-GRPO | 79.8 | 37.2 | - | - |
| UniGRPO | 75.8 | 34.2 | 23.4 | - |
| SAPO | 79.8 | 37.0 | - | - |
| PAPO (Ours) | 82.4 | 35.6 | 65.6 | 25.0 |
저자들은 GRPO 계열 SFT 없는 방법 중 대부분 과제에서 최첨단(SOTA) 결과를 세웠다고 밝힌다. 또한 학습 보상 곡선에서 PAPO는 diffu-GRPO보다 일관되게 더 높은 보상 궤적을 더 가파른 상승과 낮은 변동성으로 달성해, 샘플 효율이 우수함을 보였다.
Ablation과 분석
각 구성 요소의 기여를 검증하기 위해 ablation 연구를 수행했다. SPR을 제거하고 희소한 최종 보상만 사용하면 성능이 크게 떨어졌고, 이는 밀집 과정 인식 피드백의 필요성을 확인한다. EHR을 엔트로피 안내 없이 균등 샘플링하는 단순 Historical Re-enactment(HR)로 대체해도 성능이 눈에 띄게 하락했다. 실제 상태에서 학습하는 것만으로는 부족하며, 정책이 가장 불확실한 고엔트로피 상태를 우선하는 것이 학습 효율 극대화에 핵심임을 보여준다.
| 구성 | GSM8K 256 | MATH500 256 | Countdown 256 | Sudoku 256 |
|---|---|---|---|---|
| LLaDA-8B-Instruct | 76.7 | 32.4 | 19.5 | 6.7 |
| HR | 80.1 | 36.8 | 56.6 | 22.6 |
| HR + SPR | 81.3 | 37.6 | 60.6 | 23.6 |
| EHR + SPR (PAPO) | 82.4 | 35.6 | 65.6 | 25.0 |
SPR의 충실도-효율 트레이드오프 분석에서, 룩어헤드를 1단계에서 16단계로 늘려도 정확도는 1.6%만 개선되는 반면 GPU 시간은 44% 증가했다. 즉 1단계 SPR이 최적의 충실도-효율 균형을 제공한다. 또한 초기 단계(t가 32 미만) 보상을 클리핑하면 성능이 뚜렷하게 저하되어, 초기 과정 보상이 안정적 수렴에 중요함을 확인했다.
EHR 분석에서 생성 과정의 엔트로피는 디노이징이 진행될수록 감소하는 비균등 분포를 보였다. 엔트로피 가중치가 0일 때(균등 샘플링) 보상이 가장 낮고 후반 불안정성이 나타난 반면, 가중치가 양수일 때 보상과 수렴 안정성이 모두 개선되었다. 또한 PAPO는 diffu-GRPO보다 더 짧은 완성문으로 수렴하면서 더 높은 성능을 달성해 토큰 효율도 향상되었다.
교차 도메인 일반화 실험에서 PAPO는 LLaDA-1.5 백본에서도 모든 벤치마크에서 diffu-GRPO를 능가했고, 특히 계획 과제에서 큰 이득을 보였다. 코드 생성(HumanEval, MBPP)처럼 실행 기반 보상을 쓰는 복잡한 도메인에서도 경쟁력 있거나 우수한 결과를 얻어, 수학과 계획 과제에 국한되지 않음을 확인했다. 효율 측면에서 정책 업데이트 횟수 뮤를 2에서 24로 늘리면 수렴이 빨라지지만 24에서는 후반 불안정성이 나타났고, 뮤를 12로 두면 안정적으로 더 높은 최종 보상에 빠르게 수렴했다.
한계와 주의사항
논문은 두 가지 확장 과제를 한계로 제시한다. 첫째, 과정 보상은 다단계 마스크 확산 롤아웃에서 한 단계 디노이징 예측에 의존한다. T3D 같은 소수 단계 증류나 Fast-dLLM 같은 학습 불필요 고속 디코더가 롤아웃 단계 수를 줄이면, 블록 수준이나 제한된 다단계 룩어헤드로 과정 보상을 정의해 과도한 비용 없이 보상 충실도를 높이는 연구가 필요하다. 둘째, 현재 실험은 텍스트 기반 추론 벤치마크에 한정된다. 모달리티별 과정 보상과 도메인 구조가 텍스트와 크게 다를 수 있는 멀티모달 dLLM에서의 효과는 추가 연구가 필요하다.
계산과 메모리 효율 측면에서, PAPO는 모든 디노이징 단계가 아니라 엔트로피 안내로 선택한 가장 정보량 높은 뮤개 단계에 대해서만 과정 보상을 계산하고 상태를 캐시한다. 과정 보상에 필요한 한 단계 예측 상태는 표준 롤아웃 과정에서 디노이즈와 리마스크 사이에 생성되는 일시적 산물이므로 추가 생성 단계가 필요 없다. 또한 diffu-GRPO에서 매 업데이트마다 인위적 학습 입력을 만드는 비용이 제거되어 단계당 벽시계 시간(wall time)이 오히려 줄어든다고 보고한다.
결론
PAPO는 dLLM에 RL을 적용할 때 나타나는 두 가지 근본적 불일치, 즉 희소한 최종 보상에 의존하는 과정 보상 불일치와 불충실한 맥락에서의 비효율적 상태 선택을 해결한다. 핵심은 RL 업데이트 과정을 모델의 실제 생성 궤적과 전체적으로 정렬하는 것이다. 이를 위해 밀집하고 단계별로 차별화된 크레딧 할당을 제공하는 SPR과, 효율적이고 실제적인 상태 선택을 수행하는 EHR이라는 두 시너지 모듈을 사용한다. 네 가지 추론 벤치마크 실험에서 PAPO는 강력한 베이스라인을 큰 폭으로 능가했으며, 과정 정렬 접근이 dLLM의 복잡한 추론 능력을 끌어올리는 더 안정적이고 효율적인 경로임을 보였다.