SPG: Sandwiched Policy Gradient for Masked Diffusion Language Models

1. 문제 정의 (Problem)

• 확산 언어 모델 (dLLM) 의 잠재력: 확산 모델은 여러 토큰을 병렬로 디코딩할 수 있어 추론 지연 시간을 크게 줄일 수 있는 잠재력을 가지고 있습니다.
• RL 적용의 난제: dLLM 을 인간의 선호도나 특정 작업 보상 (예: 추론 능력 향상) 에 맞춰 정렬하기 위해 강화 학습 (RL) 을 적용하려는 시도가 증가하고 있습니다. 그러나 dLLM 의 **로그 가능도 (log-likelihood) 가 계산적으로 다루기 어렵다 (intractable)**는 근본적인 문제가 있습니다.
• 기존 방법의 한계:
  • 기존 연구들은 로그 가능도를 근사하기 위해 **증하 하한 (Evidence Lower Bound, ELBO)**이나 1 단계 추정치를 사용했습니다.
  • ELBO 의 문제점: ELBO 는 실제 로그 가능도의 하한선일 뿐입니다. 보상 (Reward) 이 양수일 때는 하한선을 최대화하는 것이 유효하지만, 보상이 음수 (부정적 피드백) 일 때 하한선을 최소화한다고 해서 실제 로그 가능도가 반드시 감소하는 것은 아닙니다.
  • 이로 인해 정책 경사 (Policy Gradient) 에 편향 (Bias) 이 발생하여, 모델이 나쁜 출력을 효과적으로 처벌하지 못하거나 최적의 정책을 학습하지 못하는 문제가 발생합니다.

2. 방법론 (Methodology)

저자들은 **샌드위치 정책 경사 (SPG)**를 제안하여 로그 가능도의 불확실성을 해결하고 편향을 줄였습니다.

A. 샌드위치 정책 경사 (Sandwiched Policy Gradient)

SPG 는 생성된 시퀀스의 로그 가능도를 "샌드위치"처럼 하한선과 상한선으로 감싸는 전략을 사용합니다.

• 긍정적 보상 (Positive Advantage): 보상 Advantage 가 양수인 시퀀스에 대해서는 **하한선 (ELBO)**을 최대화하여 실제 로그 가능도를 증가시킵니다.
• 부정적 보상 (Negative Advantage): 보상 Advantage 가 음수인 시퀀스에 대해서는 **상한선 (Evidence Upper Bound, EUBO)**을 최소화하여 실제 로그 가능도를 감소시킵니다.
• 수식적 접근:
  • 기존 RL 목적함수 $J(\theta)$는 $\sum A_j \log \pi_\theta(x_j)$ 형태입니다.
  • SPG 는 이를 $J_{SPG} = \sum (1_{A_j \ge 0} \cdot A_j \cdot L_{ELBO} + 1_{A_j < 0} \cdot A_j \cdot L_{EUBO})$로 변환합니다.
  • 이를 통해 실제 목적함수의 하한선을 보장하며, 편향을 줄인 더 견고한 정책 경사를 유도합니다.

B. 상한선 (EUBO) 유도

• Rényi 변분 경계 (Rényi Variational Bound) 를 기반으로 마스크드 확산 모델에 적합한 **계산 가능한 상한선 (EUBO)**을 유도했습니다.
• 이 상한선은 로그 함수가 기대값 바깥에 위치하는 형태를 가지며, $\beta$ 파라미터를 통해 경계의 엄격함을 조절할 수 있습니다.

C. 블록 단위 마스킹 전략 (Block-wise Masking Strategy)

• Monte Carlo 추정을 위한 샘플링 과정에서 무작위 마스킹 대신 블록 단위 마스킹을 도입했습니다.
• 시퀀스를 블록으로 나누고, 현재 생성 블록을 선택하여 그 전의 블록은 깨끗하게 유지하고 그 이후의 블록은 완전히 마스킹하는 방식입니다.
• 이는 정책 롤아웃 (rollout) 시 모델이 접하는 데이터 분포와 최적화 시의 분포를 더 잘 정렬시켜 (Alignment) 추정의 안정성을 높이고 수렴 속도를 개선합니다.

D. 혼합 목적함수 (Mixture Objective)

• 순수한 EUBO 는 Monte Carlo 추정이 편향될 수 있고 계산 비용이 높을 수 있습니다.
• 이를 해결하기 위해 ELBO 와 EUBO 를 혼합한 목적함수를 사용합니다.
  • $\tilde{L}_{Mix} = \omega \cdot \tilde{L}_{EUBO} + (1-\omega) \cdot L_{ELBO}$
• 이론적 분석 (Proposition 1) 을 통해 이 혼합 방식이 단일 경계 사용보다 분산 (Variance) 을 줄이고 더 안정적인 학습을 가능하게 함을 증명했습니다.

3. 주요 기여 (Key Contributions)

1. 새로운 RL 알고리즘 (SPG): 보상 기반의 샌드위치 변분 경계를 최적화하여 dLLM 의 로그 가능도 추정 편향을 줄인 새로운 정책 경사 알고리즘을 제안했습니다.
2. 블록 단위 마스킹 기법: Monte Carlo 추정 시 데이터 분포 정렬을 개선하여 학습 목적함수 추정의 안정성을 높이는 기법을 도입했습니다.
3. State-of-the-Art 성능: 수학 및 논리 추론 벤치마크 (GSM8K, MATH500, Countdown, Sudoku) 와 코딩 벤치마크 (HumanEval, MBPP) 에서 기존 dLLM RL 알고리즘 (D1, WD1, UniGRPO 등) 을 압도하는 성능을 달성했습니다.

4. 실험 결과 (Results)

SPG 는 다양한 벤치마크에서 기존 최첨단 (SOTA) 방법들보다 월등한 성능을 보였습니다. 특히 생성 길이 256 기준의 결과는 다음과 같습니다.

• 수학 및 논리 추론:
  • GSM8K: 86.1% (기존 SOTA 대비 +3.6% 향상)
  • MATH500: 40.0% (기존 SOTA 대비 +2.6% 향상)
  • Countdown: 70.7% (기존 SOTA 대비 +18.4% 향상)
  • Sudoku: 94.0% (기존 SOTA 대비 +27.0% 향상)
• 코딩 작업:
  • HumanEval: 41.5% (+1.9% 향상)
  • MBPP: 50.6% (+4.7% 향상)
• 학습 동향: SPG 는 학습 초기부터 빠른 수렴을 보이며, 다른 방법들보다 더 높은 보상 수준에 도달하는 것을 Fig 3 에서 확인할 수 있습니다.
• 추론 전략에 대한 강건성: 다양한 디코딩 전략 (Semi-AR, Random 등) 에서도 SPG 는 일관적으로 우수한 성능을 유지하여 모델의 일반화 능력을 입증했습니다.

5. 의의 및 결론 (Significance)

• 이론적 기여: dLLM 에 RL 을 적용할 때 발생하는 로그 가능도의 비가역성 (intractability) 문제를 해결하기 위해, 하한선과 상한선을 동시에 활용하는 이론적 프레임워크를 정립했습니다. 이는 단순한 ELBO 근사의 한계를 넘어설 수 있는 새로운 방향을 제시합니다.
• 실용적 가치: 확산 기반 언어 모델이 autoregressive 모델과 경쟁할 수 있는 추론 능력을 갖추기 위해 필수적인 **보상 기반 정렬 (Reward-based Alignment)**을 효과적으로 가능하게 합니다.
• 미래 전망: SPG 는 확산 모델의 병렬 처리 장점을 살리면서 RLHF(인간 피드백을 통한 강화 학습) 나 추론 능력 향상과 같은 고급 작업을 수행할 수 있는 기반을 마련했습니다.

요약하자면, 이 논문은 확산 언어 모델의 RL 학습을 위한 핵심적인 이론적 장벽 (편향된 경사) 을 "샌드위치" 전략으로 해결함으로써, 수학 및 논리 추론 분야에서 획기적인 성능 향상을 이끌어낸 중요한 연구입니다.