기술 요약: 트레이스를 읽고, 경로를 조종하라 (CAPR)

문제 정의

확산 대규모 언어 모델(dLLM)은 병렬 단계를 통해 전체 마스킹된 시퀀스를 반복적으로 디노이징(denoising)하여 텍스트를 생성합니다. 자기회귀(autoregressive) 모델과 달리, dLLM은 모든 디노이징 단계에서 모든 마스킹된 위치에 대한 예측 분포를 노출하며, 이는 어떤 토큰이 확신을 갖고 안정화되는지, 혹은 불안정하거나 진동하는지를 보여주는 풍부한 "디노이징 트레이스(denoising trace)"를 생성합니다.

현재 dLLM을 위한 강화 학습(RL) 사후 학습 방법들은 이 트레이스를 효과적으로 활용하지 못하고 두 가지 극단에 빠져 있습니다:

1. 평면적 롤아웃(Flat Rollouts): 전체 궤적에 대해 단일 결과 보상을 할당합니다. 계산 비용은 저렴하지만, 어떤 특정 디노이징 단계나 토큰 블록이 최종 결과에 결정적이었는지에 대한 신호를 제공하지 못합니다.
2. 트리 롤아웃(Tree Rollouts): 대안적인 연속(continuations)을 분기하여 보상을 비교하고 세그먼트 수준의 어드밴티지(advantage)를 생성합니다. 이는 더 미세한 입도를 제공하지만, 브랜칭 팩터(branching factor)와 깊이에 따라 계산 비용이 급격히 증가하여 매우 비효리적일 수 있습니다.

본 연구가 다루는 핵심 질문은 다음과 같습니다: 디노이징 트레이스 자체가 전체 트리 확장의 계산 비용 없이도 트리와 유사한 감독 신호를 제공할 수 있는가?

방법론: CAPR

저자들은 각 블록으로부터 디노이징 트레이스에서 per-block 어드밴티지를 추출하도록 설계된 블록 기반 RL 알고리즘인 **CAPR (Cached-Amortized Path Refinement)**을 소개합니다. CAPR은 세 가지 주요 메커니즘을 통해 작동합니다.

1. 캐시 및 조종 (Cache & Steer)

CAPR은 초기 디노이징 단계를 기록하여 각 위치 $j$에 대한 압축된 경로 상태(path state) $h_j$를 구축합니다. 이 상태는 다음을 요약합니다:

• 현재 상위 토큰에 대한 확신도(Confidence).
• 확률 분포의 엔트로피(Entropy).
• 여러 단계에 걸친 상위 토큰의 안정성(Stability).

이 상태는 GRU 스타일의 메커니즘을 통해 업데이트됩니다. 높은 엔트로피는 메모리를 리셋하고(오래된 증거를 망각), 안정성은 이를 보존합니다. 경로 상태는 두 가지 기능을 수행합니다:

• 조종(Steering): 다음 역단계(reverse step)의 로짓(logits)에 편향을 줍니다. 안정적이고 확신이 있는 위치에는 양(+)의 푸시를 주고, 진동하는 위치는 억제합니다.
• 언마스킹 제어(Unmasking Control): 확신도에 따라 어떤 마스킹된 위치를 드러낼지(unmask) 결정하여, 모델이 안정적인 확약(commitment)에 집중하도록 합니다.

2. 분기 및 가지치기 (Branch & Prune)

전체 트리 확장 없이 학습 신호를 생성하기 위해, CAPR은 디노이징 궤적의 중간 지점($f_b = 0.5$)에서 단 한 번의 분기 이벤트를 수행합니다.

• 캐싱(Caching): 모델은 분기 지점까지의 접두사(prefix)(로짓, 예측된 토큰 및 경로 상태)를 캐싱합니다.
• 포킹(Forking): Gumbel 섭동(perturbation)을 어휘 확률에 적용하여 소수의 형제 연속형(sibling continuations, 예: $G_b = 2$)을 생성합니다. 이 형제들은 캐시된 접두사를 재사용하여 중복된 순방향 패스(forward pass)를 방지합니다.
• 가지치기(Pruning): 블록 경계에서, 경로 상태의 품질(평균 안정성)을 기준으로 브랜치를 가지치기합니다. 가장 유망한 브랜치만이 계속 진행되도록 하여 탐색 공간이 전체 트리로 폭발하는 것을 방지합니다합니다.

3. 블록 크리틱 (Block Critic)

경량화된 가치 헤드(Block Critic)가 각 블록 경계에서 경로 상태를 읽습니다.

• 보상 재분배: 롤아웃이 완료되면, 최종 결과 보상은 각 블록에서 드러난 토큰 수와 경로 상태 특징을 기반으로 블록 전체에 재분배됩니다.
• 어드밴티지 계산: TD/GAE(Generalized Advantage Estimation)를 사용하여, 크리틱은 희소한 최종 보상을 per-block 어드밴티지($\hat{A}_b$)로 변환합니다.
• PPO 업데이트: 이러한 블록 어드밴티지는 PPO에서 사용되는 표준 롤아웃 수준의 어드밴티지에 가중치를 다시 부여합니다. 이를 통해 모델은 어떤 특정 블록이 성공에 가장 크게 기여했는지 학습할 수 있으며, 결과적으로 트리 탐색의 정밀도를 계산 비용 없이 복구합니다.

학습 목적 함수(Training Objective): 총 손실(loss)은 클립된 PPO 손실(블록 가중치 어드밴티지 사용), KL 발산(참조 모델 근처 유지용), $\chi^2$ 정규화 항(희귀한 로그 비율 스파이크 억제용), 자기 증류(self-distillation) 손실(해결된 롤아웃을 지도 타겟으로 사용), 그리고 가치 헤드 손실을 결합합니다.

주요 기여

1. 디노이징 트레이스 기반 RL: 본 논문은 확신도와 안정성을 블록 단위 어드밴티지로 변환함으로써, 전체 트리 확장의 필요성을 제거하는 새로운 관점을 제안합니다.
2. CAPR 알고리즘: 접두사를 캐싱하고, 형제 롤아웃을 분할 상환(amortize)하며, 경로 상태를 통해 디노이징 단계를 조종하고, 블록 크리틱을 사용하여 정밀한 PPO 가중치를 생성하는 실용적인 구현을 제시합니다.
3. 최첨단 성능(SOTA): CAPR은 256 및 512 토큰 예산 환경에서 밀집(dense) 및 혼합 전문가(MoE) LLaDA 백본 모두에 대해 네 가지 추론 벤치마크(4×4 Sudoku, Countdown, GSM8K, Math500)에서 새로운 SOTA 결과를 달성했습니다.

실험 결과

저자들은 Diffu-GRPO, VRPO, wd1, SAPO, d2-stepMerge, TraceRL, 그리고 트리 구조의 d-TreeRPO를 포함한 베이스라인들과 CAPR을 비교 평가했습니다.

• 정확도: 256 토큰 예산에서, CAPR은 LLaDA-8B-Instruct 및 LLaDA-MoE-7B-A1B-Instruct 모두에 대해 모든 태스크에서 베이스라인들을 능가하거나 동등한 성능을 보였습니다.
  • Sudoku: CAPR은 LLaDA-8B에서 93.0%, LLaDA-MoE에서 **84.8%**를 달성하여 가장 강력한 트리 구조 베이스라인(d-TreeRPO: 92.9% 및 78.3%)을 능가하거나 대등한 성적을 냈습니다.
  • Countdown: 상당한 이득이 관찰되었으며, CAPR은 72.7%(8B) 및 68.8%(MoE)에 도달하여 d-TreeRPO를 앞질렀습니다.
• 효율성: CAPR은 트리 롤아웃에 비해 계산 비용을 크게 줄였습니다.
  • 생성 비용: CAPR은 표준 설정에서 평면적 롤아웃의 약 0.75배, 트리 롤아웃의 약 0.6배의 비용만 필요합니다.
  • 실행 시간(Wall-Clock Time): Sudoku의 경우, CAPR은 최적화 단계당 165초가 소요된 반면 d-TreeRPO는 596초가 소요되었습니다. 즉, CAPR은 1/3 미만의 계산량으로 대등한 정확도를 달성했습니다.
• 절제 연구(Ablation Studies):
  • 블ock Critic이 가장 큰 중간 이득을 제공했으며, 학습 곡선을 실질적으로 변화시켰습니다.
  • Branch & Prune은 유용한 다양성을 추가했지만, 단독으로는 효과가 덜했습니다.
  • Self-Distillation은 제로 분산 그룹(해결된 롤아웃)에서의 학습을 안정화하고 그래디언트 스파이크를 줄이는 데 결정적이었습니다.

의의 및 주장

본 논문은 CAPR이 dLLM-RL을 위한 실용적인 "중간 지점"을 제공한다고 주장합니다. CAPR은 (대안들의 국소적 비교라는) 트리 탐색의 정밀도를 성공적으로 복구하면서도, (전체 트리 확장의) 계산 폭발은 피합니다.

• 효율성 vs 정밀도: CAPR은 디노이징 트레이스 자체가 블록 수준에서 보상을 할당하기에 충분한 정보를 포함하고 있음을 입증하여, 값비싼 트리 기반 롤아웃의 필요성을 제거했습니다.
• 일반화: 이 방법은 태스크별 재조정 없이 밀집(dense) 백본에서 희소(sparse, MoE) 백본으로 효과적으로 전이됩니다.
• 실용성: 트리 롤아웃의 약 0.6배 수준으로 per-step 계산량을 줄이면서도 그 이상의 정확도를 달성함으로써, CAPR은 더 큰 모델과 긴 컨텍스트에 대한 고정밀 RL 학습을 가능하게 만듭니다.

저자들은 한계점 또한 언급했습니다. 현재 이 방법은 자동으로 검증 가능한 추론 태스크에서 평가되었으며, 개방형 지시 이행(open-ended instruction following)이나 노이즈가 있는 비지역적(non-local) 보상이 존재하는 도메인에는 적응이 필요할 수 있습니다. 또한, 블록 크리틱이 온라인으로 학습되므로 매우 긴 추론 체인에서는 신뢰성 문제가 발생할 수 있습니다.