Stabilizing Reinforcement Learning for Diffusion Language Models

Each language version is independently generated for its own context, not a direct translation.

1. 배경: 새로운 학습 방법과 그 문제점

과거의 AI 는 글을 쓸 때 **한 글자씩 순서대로 **( autoregressive) 써나가는 방식이었습니다. 하지만 최근에는 **확산 모델 **(Diffusion Model)이라는 새로운 방식이 등장했습니다. 이는 마치 퍼즐을 맞추거나, 흐릿한 그림을 선명하게 만드는 과정처럼, 처음엔 막연한 상태에서 시작해 점차 단서를 찾아 정답을 완성해 나가는 방식입니다.

이 새로운 방식은 병렬 처리가 가능해서 훨씬 빠르고 유연한데, 문제는 **보상 학습 **(Reinforcement Learning)을 적용하려니 AI 가 엉망이 된다는 것입니다.

비유: 기존 AI 는 한 줄씩 글을 써가는 작가라면, 확산 모델 AI 는 한 번에 문장을 다 써놓고 수정하는 편집자입니다. 편집자가 자신의 글을 고칠 때, "어떤 수정이 좋은지"를 판단하는 기준 (확률) 을 계산하는 것이 매우 어렵고, 이 계산에 **오차 **(노이즈)가 생기기 쉽습니다.

2. 문제의 핵심: "폭주하는 학습" (Reward Collapse)

기존에 쓰이던 학습 알고리즘 (GRPO) 을 이 새로운 AI 에 그대로 적용하려니 재미있는 일이 벌어집니다.

상황: AI 가 학습할 때, "내 수정이 얼마나 좋은가?"를 계산하는 과정에서 계산 오차가 발생합니다.
문제: 이 오차가 **갑작스러운 폭주 **(Gradient Spike)를 일으킵니다. 마치 운전자가 브레이크를 밟아야 할 때, 오작동으로 인해 발이 페달에서 미끄러져서 차가 미친 듯이 가속하는 것과 같습니다.
악순환:
1. 오차로 인해 AI 가 엉뚱한 방향으로 크게 움직입니다 (폭주).
2. AI 의 방향이 틀어지면, 다음에 계산할 때 오차가 더 커집니다.
3. 오차가 커지면 다시 더 큰 폭주가 일어납니다.
4. 결국 AI 는 **학습을 포기하고 엉망이 되어버립니다 **(Reward Collapse).

이것을 **"불안정성 고리 **(Instability Loop)라고 부릅니다.

3. 해결책: 'StableDRL' (안정적인 학습 시스템)

저자들은 이 악순환을 끊기 위해 StableDRL이라는 새로운 방법을 제안했습니다. 두 가지 핵심 장치로 이루어져 있습니다.

① 무조건적인 브레이크 (Unconditional Clipping)

기존 방식: "수정이 너무 크면 멈추라"고 했지만, 오작동 (노이즈) 이 있으면 브레이크가 안 먹히는 경우가 있었습니다.
StableDRL 방식: 무조건 수정의 크기를 제한합니다. "너무 크면 무조건 잘라내라"는 원칙입니다.
비유: 차가 미친 듯이 가속하려 할 때, 브레이크가 고장 나더라도 자동으로 작동하는 비상 브레이크를 달아놓은 것입니다. 오작동으로 인해 차가 날아가지 못하게 막아줍니다.

② 자기 조절 (Self-Normalization)

기존 방식: 학습 속도를 조절할 때, 단순히 '평균'을 사용했습니다. 그런데 한 두 개의 엉뚱한 데이터 (오차) 가 평균을 왜곡하면 전체 학습이 흔들렸습니다.
StableDRL 방식: 엉뚱한 데이터의 영향을 줄이고, **모든 데이터가 합쳐진 '중심'**으로 학습 방향을 잡습니다.
비유: 배를 탈 때, 한 두 사람이 갑자기 배를 밀어내면 배가 뒤집힙니다. 하지만 모든 사람이 함께 밀어내는 힘의 중심을 찾아서 밀면, 아무리 한 두 사람이 미친 듯이 밀어도 배는 안정적으로 나아갑니다.

4. 추가 기술: '계단식 주의 (Staircase Attention)'

이 방법은 AI 가 긴 문장을 다룰 때도 적용됩니다.

문제: 긴 문장을 한 번에 다 볼 때, AI 가 **미래의 정답을 미리 훔쳐보는 **(Leakage) 문제가 생깁니다.
해결: **계단 **(Staircase)처럼 단계별로만 정보를 볼 수 있게 막아줍니다.
비유: 시험지를 풀 때, 앞쪽 문제만 보고 뒤쪽 문제의 정답은 볼 수 없게 가려주는 '계단식 커튼'을 치는 것입니다. 이렇게 하면 AI 는 정답을 훔치지 않고도 정확하게 학습할 수 있습니다.

5. 결과: 무엇이 달라졌나요?

이 방법을 적용한 결과, AI 는 수천 번의 학습을 안정적으로 거칠 수 있게 되었습니다.

이전: 학습 300 번 정도 되면 엉망이 되어 멈췄습니다.
이제: 1,000 번 이상 학습해도 안정적으로, **수학 문제 **(MATH500) 등 복잡한 추론 능력을 비약적으로 향상시켰습니다.

요약

이 논문은 **"새로운 AI 모델 **(확산 모델)을 발견하고, **"무조건적인 브레이크 **(Unconditional Clipping)와 **"자기 조절 **(Self-Normalization)이라는 두 가지 장치를 넣어 학습을 안정화시켰다는 내용입니다.

마치 미친 듯이 가속하는 차에 강력한 비상 브레이크와 안정 장치를 달아주어, 이제 AI 가 복잡한 논리 문제도 안정적으로 해결할 수 있게 된 것입니다.

Each language version is independently generated for its own context, not a direct translation.

1. 문제 정의 (Problem)

최근 등장한 이산 확산 대규모 언어 모델 (Discrete Diffusion Large Language Models, dLLMs) 은 병렬 디코딩과 양방향 컨텍스트 모델링의 이점을 제공하지만, 기존 autoregressive (AR) 모델에 성공적으로 적용되던 그룹 상대적 정책 최적화 (Group Relative Policy Optimization, GRPO) 를 직접 적용할 때 심각한 불안정성을 보입니다.

주요 현상: dLLM 에 GRPO 를 적용하면 학습 초기 (약 300 스텝 부근) 에 보상 붕괴 (Reward Collapse) 가 발생하며 학습이 실패합니다.
근본 원인: GRPO 는 AR 모델의 시퀀스 확률 기반 중요도 비율 (Importance Ratio, $\rho$ ) 에 의존하지만, dLLM 에서는 이 확률이 계산적으로 다루기 어렵기 (intractable) 때문에 ELBO(Evidence Lower Bound) 나 평균장 (mean-field) 근사 등을 통해 추정 (Estimation) 해야 합니다.
불안정성 루프:
1. 노이즈 발생: 추정된 중요도 비율은 본질적으로 노이즈가 많고 분산이 큽니다.
2. 기울기 스파이크 (Gradient Spikes): 기존 GRPO 의 조건부 클리핑 (conditional clipping) 메커니즘은 AR 모델에서는 작동하지만, dLLM 의 추정 노이즈로 인해 특정 조건 (부정적 이득, $A < 0$ ) 에서 클리핑이 우회되어 거대한 기울기 스파이크가 발생합니다.
3. 정책 편향 (Policy Drift): 이 스파이크는 정책을 급격히 변화시켜 행동 정책 (behavior policy) 과의 거리를 벌립니다.
4. 악순환: 정책이 멀어질수록 추정의 분산이 더 커지고, 이는 다시 더 큰 노이즈와 기울기 스파이크를 유발하는 자기 강화 불안정성 루프 (Self-reinforcing Instability Loop) 를 형성합니다.

2. 방법론 (Methodology: StableDRL)

저자들은 이 불안정성 루프를 끊기 위해 StableDRL이라는 새로운 GRPO 변형 알고리즘을 제안합니다. 두 가지 핵심 기법을 통해 노이즈가 기울기 업데이트에 미치는 영향을 제어합니다.

가. 무조건적 클리핑 (Unconditional Clipping)

기존 GRPO: 이득 (Advantage) 의 부호에 따라 클리핑을 적용합니다. ( $A < 0$ 일 때 $\rho > 1+\epsilon$ 이면 클리핑하지 않음). 이는 dLLM 의 노이즈가 있는 $\rho$ 값이 클리핑 조건을 우회하여 무제한의 기울기 스파이크를 유발하게 합니다.
StableDRL: 이득의 부호와 관계없이 중요도 비율 $\hat{\rho}$ 를 항상 $[1-\epsilon, 1+\epsilon]$ 범위 내에서 무조건적으로 (Unconditionally) 클리핑합니다. 이를 통해 이상치 (outlier) 로 인한 기울기 스파이크를 원천 차단합니다.

나. 자기 정규화 (Self-Normalization)

기존 GRPO: 업데이트 크기를 고정된 그룹 크기 ( $G$ ) 로 나누어 정규화합니다. dLLM 의 중요도 비율 추정 분산이 높을 경우, 이 고정된 정규화는 기울기 크기의 급격한 요동을 유발합니다.
StableDRL: 고정된 $G$ $G$ 대신 클리핑된 중요도 비율의 합 ( $\sum \text{clip}(\hat{\rho}_i)$ ) 으로 나누어 업데이트를 정규화합니다.
- 효과: 이는 업데이트 벡터를 개별 샘플 기울기들의 볼록 껍질 (Convex Hull) 내부로 제한합니다. 즉, 그룹 수준의 노이즈로 인한 크기 변동을 제거하고, 업데이트 크기를 샘플 간 가중치 변동과 해리 (decouple) 시킵니다.

다. 블록 확산 모델 확장 (Staircase Attention)

블록 확산 (Block Diffusion) 모델에 적용하기 위해 계단식 어텐션 (Staircase Attention) 메커니즘을 도입했습니다.
이 구조는 청크 (block) 단위로 과거의 깨끗한 컨텍스트는 참조하되, 현재 블록의 정답 (Ground Truth) 은 누출되지 않도록 가중치를 설계하여, 단일 패스 ( $O(1)$ ) 로 유효한 ELBO 추정치를 계산할 수 있게 합니다.

3. 주요 기여 (Key Contributions)

불안정성 메커니즘 규명: dLLM 에 GRPO 를 적용할 때 발생하는 보상 붕괴가 '추정 노이즈 $\rightarrow$ 기울기 스파이크 $\rightarrow$ 정책 편향 $\rightarrow$ 노이즈 증폭'의 자기 강화 루프에서 비롯됨을 이론적 및 실험적으로 증명했습니다.
StableDRL 프레임워크 제안: 무조건적 클리핑과 자기 정규화를 결합하여 dLLM 의 전 파라미터 (Full-parameter) RL 학습을 안정화하는 새로운 알고리즘을 개발했습니다.
SOTA 성능 달성: 풀 어텐션 (Full-attention) 및 블록 확산 (Block diffusion) 아키텍처 모두에서 1,000 스텝 이상의 안정적인 학습을 가능하게 하여, 기존 방법론들 (ESPO, SPG 등) 보다 뛰어난 추론 능력을 입증했습니다.

4. 실험 결과 (Results)

학습 안정성: 기존 GRPO 기반 방법들은 300 스텝 내외에서 보상 붕괴가 발생하지만, StableDRL 은 1,000 스텝 이상 안정적으로 수렴하며 보상이 지속적으로 향상됩니다.
추론 성능 (Full-Attention, LLaDA-8B):
- GSM8K: 84.2% (기존 SOTA 대비 향상)
- MATH500: 41.8% (SPG 대비 +5.2% 향상)
- Countdown: 83.5% (SPG 대비 +13.7% 향상)
- Sudoku: 91.5%
- 다양한 시퀀스 길이 (128~512 토큰) 에서 일관된 성능을 보이며, 길이 일반화 (Length Generalization) 능력이 뛰어납니다.
블록 확산 모델 (SDAR-8B):
- AIME 2024: 16.7% (기존 AR 모델인 Qwen3-8B 의 10.0% 를 상회)
- 동적 샘플링 환경에서도 Trado 등 기존 방법보다 뛰어난 견고성 (Robustness) 을 보입니다.
스트레스 테스트: 중요도 비율의 분산을 인위적으로 폭발시키는 "Exploding Weight" 테스트에서도 StableDRL 은 안정적인 성능을 유지한 반면, ESPO 와 SPG 는 즉각적인 붕괴를 겪었습니다.

5. 의의 및 결론 (Significance)

이 논문은 확산 기반 언어 모델 (dLLMs) 이 가진 잠재력을 RL 을 통해 완전히 unlocking 하는 데 있어 핵심적인 장벽인 학습 불안정성 문제를 해결했습니다.

이론적 통찰: 중요도 비율 추정의 노이즈가 어떻게 RL 학습을 붕괴시키는지 그 메커니즘을 명확히 규명했습니다.
실용적 가치: StableDRL 은 dLLM 의 전 파라미터 미세조정 (Full-parameter fine-tuning) 을 가능하게 하여, 파라미터 효율적인 방법 (LoRA 등) 이나 제한된 학습을 통해 얻기 힘든 고도화된 추론 능력을 실현했습니다.
미래 전망: 이 연구는 확산 모델과 강화 학습의 결합을 위한 새로운 표준을 제시하며, 복잡한 추론 작업 (수학, 논리 퍼즐 등) 에서 dLLM 이 AR 모델을 능가할 수 있는 가능성을 열었습니다.

요약하자면, StableDRL은 dLLM 의 고유한 특성 (추정 노이즈) 을 고려하여 GRPO 의 구조적 결함을 수정함으로써, 확산 모델의 RL 학습을 안정화하고 최첨단 추론 성능을 달성한 획기적인 연구입니다.