🤖 machine learning

Predictive Divergence Masks for LLM RL

이 논문은 PPO의 비율 기반 방향 기준을 근접성 기준에 사용되는 확률 발산의 증가 또는 감소 여부에 대한 폐쇄형 예측으로 대체함으로써, 모델 규모 전반에 걸쳐 훈련 안정성과 성능을 향상시키는 LLM 강화 학습을 위한 새로운 방법인 예측 발산 마스크(Predictive Divergence Masks)를 제안한다.

원저자: Xiangxin Zhou, Jiarui Yao, Penghui Qi, Bowen Ping, Jiaqi Tang, Haonan Wang, Tianyu Pang

게시일 2026-07-14

📖 5 분 읽기🧠 심층 분석

CC BY 4.0

원저자: Xiangxin Zhou, Jiarui Yao, Penghui Qi, Bowen Ping, Jiaqi Tang, Haonan Wang, Tianyu Pang

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 초지능형 로봇에게 이야기를 쓰는 법을 가르치고 있다고 상상해 보세요. 당신은 로봇이 실수로부터 배우기를 원하지만, 한 가지 엄격한 규칙을 세웠습니다: "한 번에 너무 급격하게 글쓰기 스타일을 바꾸지 마라." 만약 로봇이 너무 흥분해서 완전히 다른 목소리로 글을 쓰기 시작한다면, 당신은 브레이크를 밟아야 합니다. 이것이 바로 거대 언어 모델(LLM)을 위한 **강화 학습(Reinforcement Learning, RL)**의 역할입니다.

오랫동안, 이 브레이크를 밟는 표준적인 방법은 PPO라고 불리는 방식이었습니다. 이것은 단 하나의 단어를 기준으로 작동하는 단순한 신호등과 같았습니다. 로봇이 단어를 선택할 때마다 시스템은 다음과 같이 확인했습니다: "이 단어가 로봇이 평소에 선택하는 것보다 훨씬 더 높은 확률인가? 그리고 보상이 이 단어를 더 높은 확률로 만들라고 요구하고 있는가?" 만약 두 질문에 모두 "예"라고 답한다면, 시스템은 "멈춰! 너무 멀리 갔어!"라고 말하며 그 단어의 교훈을 무시했습니다.

하지만 여기에 문제가 있습니다. 로봇은 단 하나의 단어만 고르는 것이 아니라, 그 단어를 선택하기 위해 자신의 전체적인 성격(확률 분포)을 재배열합니다. 기존 방식은 로봇이 선택한 그 단어 하나만을 보았습니다. 마치 오케스트라 전체를 평가하면서 단 하나의 바이올린 소리만 듣는 것과 같았습니다. 이 방식은 나머지 음악이 어떻게 변했는지는 놓치고 있었습니다.

새로운 아이디어: "예측적 발산 마스크(Predictive Divergence Mask)"

Tencent, UIUC, NUS의 연구진은 이러한 단일 단어 관점이 오해의 소지가 있다는 점을 깨달았습니다. 그들은 예측적 발산 마스크라고 불리는 새로운 규칙을 제안했습니다.

단순히 로봇이 선택한 하나의 단어를 보는 대신, 이 새로운 방식은 더 큰 질문을 던집니다: "만약 우리가 이 단계를 밟는다면, 로봇의 '전체적인 성격'이 원래의 모습으로부터 더 멀어지게 될 것인가?"

이를 다음과 같이 비유할 수 있습니다:

기존 방식 (PPO): 아이가 절벽을 향해 달려가는 것을 봅니다. 당신은 그 특정 발걸음이 위험한지 확인합니다. 만약 그렇다면, 당신은 "멈춰!"라고 외칩니다.
새로운 방식 (예측적 발산): 아이가 달리고 있는 것을 봅니다. 당신은 아이의 전체 경로를 봅니다. 설령 지금 이 특정 발걸음이 괜찮아 보이더라도, 아이의 관성이 아이를 회전시켜 다른 절벽으로 떨어뜨릴 수도 있습니다. 새로운 방식은 그 단계가 일어나기 전에 전체적인 회전을 예측합니다.

작동 원리 (마법 같은 수학)

연구진은 시뮬레이션을 실제로 돌려보지 않고도 이 "전체적인 회전"을 예측할 수 있는 영리한 수학적 트릭을 찾아냈습니다.

로컬 항 (The Local Term): 이것은 기존 방식이 이미 알고 있던 부분입니다. 로봇이 선택한 특정 단어의 변화를 나타냅니다.
글로벌 항 (The Global Term): 이것이 새로운 핵심 비법입니다. 한 단어를 바꾸는 것이 로봇이 선택하지 않은 다른 모든 단어의 확률을 어떻게 강제로 조정하게 만드는지를 설명합니다 (전체 확률은 항상 100%가 되어야 하기 때문입니다).

기존 방식은 이 글로벌 항을 무시했습니다. 새로운 방식은 이를 합산합니다. 만약 수학적으로 로봇의 전체적인 성격이 너무 멀리 벗어나고 있다면, 마스크가 업데이트를 차단합니다. 만약 수학적으로 로봇이 실제로 안전한 방향으로 돌아오고 있다면, 업데이트가 허용됩니다.

"Top-K" 문제

현실 세계에는 한 가지 걸림돌이 있습니다. 로봇의 엔진(롤아웃 엔진)은 시간과 메모리를 절약하기 위해 가장 확률이 높은 상위 20개 단어("Top-K")만을 보여줍니다. 나머지 수천 개의 단어는 "꼬리(tail)" 버킷에 숨겨져 있습니다.

이 제한된 시야 속에서 예측을 가능하게 하기 위해, 연구진은 두 가지 가벼운 추정기를 발명했습니다:

집계된 꼬리 추정기 (Aggregated-Tail Estimator): 숨겨진 모든 단어를 하나의 거대한 덩어리로 간주합니다.
균등 꼬리 추정기 (Uniform-Tail Estimator): 숨겨진 단어들이 모두 동일한 확률을 가진다고 가정합니다.

연구진은 상위 20개 단어가 보통 거의 모든 확률(99% 이상)을 차지하기 때문에, 두 방법 모두 거의 똑같이 작동한다는 것을 발견했습니다. 이는 하늘을 보고 날씨를 예측하는 것과 같습니다. 비가 올지 알기 위해 구름 하나하나를 다 셀 필요는 없습니다.

실험 결과

연구진은 다양한 크기의 로봇(40억에서 300억 개의 파라미터를 가진 모델)을 대상으로 테스트했으며, 심지어 "FP8" 정밀도(학습 및 작성 엔진이 다르게 작동하게 만드는, 매우 빠르지만 약간 덜 정밀한 수학 모드) 환경에서도 테스트했습니다.

결과: 새로운 마스크는 학습을 더 안정적으로 만들었습니다.
증거: 그들은 61가지의 서로 다른 실험(seed)을 실행했습니다. 기존 방식은 실제로 로봇을 더 멀리 벗어나게 만드는 업데이트를 유지하는 경우(안전하지 않은 유지율, "unsafe keep" rate)가 36.9%에 달했습니다. 새로운 방식은 이를 **34.2%**로 줄였습니다.
개선점: 숫자가 작아 보일 수 있지만, 거대 AI 모델을 훈련하는 세계에서는 단 몇 번의 나쁜 단계를 방지하는 것만으로도 로봇이 더 잘, 더 빠르게 배울 수 있게 도와줍니다. 새로운 방식은 모든 모델 크기와 정밀도 설정에서 일관되게 작동했습니다.

명시적으로 거부한 내용

논문은 무엇이 효과가 없는지에 대해서도 매우 명확하게 밝히고 있습니다:

연구진은 업데이트의 방향을 결정하기 위해 샘플링된 토큰 중요도 비율(기존의 "단일 바이올린" 체크 방식)을 사용하는 것에 반대합니다. 이 단일 단어 체크 방식이 로봇의 전체적인 성격 변화와 자주 어긋난다는 것을 보여주었습니다.
또한, 단순히 "안전 구역"(신뢰 영역)을 더 좁게 만드는 것(임계값을 0.15에서 0.05로 줄이는 것)이 오히려 성능을 저하시킨다는 점을 보여주었습니다. 이는 단순히 더 엄격해지는 것이 답이 아니라, 방향에 대해 더 똑똑해지는 것이 중요하다는 것을 시사합니다.

얼마나 확신하는가?

저자들은 실험에서 얻은 측정된 데이터를 바탕으로 자신들의 발견에 확신을 가지고 있습니다. 그들은 진공 상태에서 시뮬레이션만 한 것이 아닙니다. 실제 수학 문제(DAPO-Math-17k 데이터셋 사용)로 실제 모델을 훈련시켰으며, AIME24 및 AIME25와 같은 벤치마크에서 정확도를 측정했습니다.

그들은 이 접근 방식이 더 나은 규칙의 정렬을 제공한다고 제안합니다. 만약 당신이 전체적인 성격(발산)을 보고 안전 구역을 정의한다면, 업데이트의 방향 또한 단 하나의 단어가 아닌 전체적인 성격을 보고 확인해야 한다는 것입니다. 결과는 이 방식이 더 안정적인 학습을 이끈다는 것을 시사하지만, 이것이 여전히 "로컬 1차 근사(local first-order approximation)"라는 점, 즉 미래 전체에 대한 완벽한 수정구슬이 아니라 바로 다음 단계에 기반한 매우 좋은 추측이라는 점을 인정하고 있습니다.

요약하자면, 새로운 방식은 로봇에게 현재 딛고 있는 발걸음만 보는 것이 아니라, 전체 지도를 보는 GPS를 주는 것과 같습니다. 이는 로봇이 세상의 끝으로 사라지지 않으면서도 더 빠르게 배울 수 있도록 돕습니다.

기술 요약: LLM RL을 위한 예측적 발산 마스크 (Predictive Divergence Masks)

1. 문제 정의

대규모 언어 모델(LLM)을 위한 강화 학습(RL)은 롤아웃(rollout)에 사용되는 추론 엔진과 학습 스택 간의 불일치, 그리고 정책의 노후화(staleness)로 인해 본질적으로 오프-폴리시(off-policy) 특성을 가집나다. 이러한 오프-폴리시 업데이트를 안정화하기 위해, 근사 정책 최적화(PPO)와 같은 표준 방법들은 신뢰 영역 마스크(trust-region masks)를 채택합니다. 이 마스크들은 토큰의 그래디언트를 마스킹(0으로 처리)할지 여부를 결정하기 위해 두 가지 기준을 사용합니다:

근접성 기준 (Proximity Criterion): 정책이 행동 정책(behavior policy)으로부터 너무 멀어졌는지 판단합니다.
방향 기준 (Direction Criterion): 업데이트가 정책을 행동 정책으로부터 더 멀어지게 만드는지(외향적), 아니면 다시 끌어당기는지(교정적)를 판단합니다.

최근 연구인 DPPO는 단일 샘플 중요도 비율(importance ratio)을 대신하여 top-K 어휘 집합에 대해 추정된 분포적 발산(예: KL 발산)을 사용하는 방식으로 근접성 기준을 개선했습니다. 그러나 DPPO는 PPO의 원래 방향 기준인 단일 샘플링된 토큰의 중요도 비율( $r_t - 1$ )에 의존하는 방식을 그대로 유지하고 있습니다.

저자들은 이러한 하이브리드 접근 방식에서 근본적인 불일치를 식별했습니다. 즉, 근접성 기준은 이제 전역적인 분포적 발산에 의해 정의되지만, 방향 기준은 여전히 단일 샘플 대리 지표(proxy)에 머물러 있다는 점입니다. 저자들은 샘플링된 비율의 변화 방향(sign)이 실제 발산의 변화 방향과 일치하지 않을 수 있음을 관찰했습니다. 결과적으로, 업데이트가 실제로는 전역적 발산을 증가시키더라도 (비록 비율은 교정적 업데이트를 시사하더라도) 해당 토큰이 유지될 수 있고, 그 반대의 경우도 발생할 수 있습니다.

2. 방법론

본 논문은 방향 기준을 근접성 기준의 분포적 성격과 일치시키는 방법인 **예측적 발산 마스크(Predictive Divergence Mask)**를 제안합니다. 이 방법은 "샘플링된 비율이 하나로부터 멀어질 것인가?"라고 묻는 대신, "다음 그래디언트 단계가 발산을 증가시킬 것인가?"라고 묻습니다.

이론적 유도

자기회귀(autoregressive) LLM에서 사용되는 소프트맥스(softmax) 정책에 대하여, 저자들은 정책 그래디언트 단계를 따른 쿨백-라이블러(KL) 발산의 방향 도함수 $\dot{D}$ 에 대한 폐쇄형(closed-form) 표현식을 유도했습니다.
$\dot{D} = \underbrace{(\pi_k - \mu_k)}_{\text{국소 항 (Local Term)}} + \underbrace{\sum_{i} \pi_i (\mu_i - \pi_i)}_{\text{전역 항 (Global Term)}}$
여기서:

$\pi_k$ 와 $\mu_k$ 는 각각 훈련 정책과 행동 정책 하에서의 샘플링된 토큰 $k$ 의 확률입니다.
**국소 항(Local Term)**은 표준 비율 기반 방향 기준에서 사용되는 양( $r - 1$ )과 정확히 일치합니다.
**전역 항(Global Term)**은 전체 어휘 집합에 걸친 정규화 결합(normalization coupling)을 포착합니다. 이 항은 단일 샘플링된 비율에서는 보이지 않지만, 실제 발산의 변화를 결정하는 데 매우 중요합니다.

예측 마스크는 이미 신뢰 영역을 벗어난 경우( $D_t > \delta$ ), $\hat{A}_t \cdot \dot{D}_t$ 의 부호가 비양수(non-positive)일 때(즉, 교정적이거나 중립적인 움직임을 나타낼 때)만 업데이트를 유지합니다.

실용적 추정

실제 운영 환경의 롤아웃 엔진은 일반적으로 전체 어휘 분포가 아닌 top- $K$ 토큰 확률만을 노출합니다. 저자들은 이 절단된 뷰(truncated view) 하에서 $\dot{D}$ 를 계산하기 위해, 분포의 보이지 않는 꼬리(tail) 부분에 대한 두 가지 경량 추정기를 제적합니다:

집계 꼬리 추정기 (Aggregated-Tail Estimator): 잔여 확률 질량을 하나의 집계된 토큰으로 통합합니다.
균등 꼬리 추정기 (Uniform-Tail Estimator): 잔여 질량을 보이지 않는 어휘 토큰들에 균등하게 분배합니다.

두 추정기 모두 유지된 top- $K$ 확률과 꼬리 질량에 대한 간단한 축약(reduction)만을 필요로 하므로, 계산 오버헤드가 매우 적습니다.

3. 주요 기여

불일치 식별: 본 논문은 분포적 근접성 기준(발산)과 단일 샘플 방향 기준(비율)을 혼합하는 것이, 방향 마스크가 실제 발산 변화를 추적하는 데 실패할 수 있는 정렬 불량(misalignment)을 초래함을 강조합니다.
예측적 발산 마스크: 발산 자체의 1차 변화를 사용하여 업데이트의 방향을 결정하는 새로운 마스킹 규칙을 제시합니다.
폐쇄형 유도: 국소(샘플링된) 성분과 전역(정규화) 성분을 명시적으로 분리하여 소프트맥스 정책에 대한 방향 도함수를 유도했습니다.
절단된 어휘 추정기: 실용적이고 오버헤드가 낮은 두 가지 추정기(집계 및 균등 꼬리)를 통해, 표준 프로덕션 롤아웃 제약 조건(top- $K$ 뷰) 하에서도 작동할 수 있도록 했습니다.

4. 실험 결과

저자들은 Qwen3 모델(4B, 8B, 30B)을 사용하여 다양한 정밀도 설정(FP8 포함)으로 DAPO-Math-17k 데이터셋에서 성능을 평가했습니다.

학습 안정성: 모든 모델 규모와 정밀도 설정(FP8 스트레스 테스트 포함)에서, 예측적 발산 마스크는 베이스라인인 DPPO-TopK-KL보다 학습 안정성을 향상시켰습니다. 반면, clip-higher 기법을 사용한 표준 GRPO는 대규모 모델 설정에서 불안정성과 붕괴를 보였습니다.
방향 기준 정렬: 토큰 수준 분석 결과, 발산 기반 방향 기준이 비율 기반 기준보다 실제 실현된 발산 변화와 더 잘 정렬됨을 확인했습니다.
- 두 기준이 상충하는 토큰들 중에서, 발산 기반 기준은 더 낮은 "위험한 유지(unsafe-keep)"율(34.2% vs 36.9%)을 보였으며, 이는 발산을 실제로 증가시키는 업데이트를 유지할 가능성이 더 낮음을 의미합니다.
- 반대로, 발산을 성공적으로 수축시킨 업데이트는 더 높은 비율로 유지했습니다(65.1% vs 62.9%).
강건성: 이 방법은 신뢰 영역 임계값( $\delta$ ) 하이퍼파라미터에 대해 개선된 강건성을 보였으며, 더 엄격한 제약 조건( $\delta = 0.05$ ) 하에서도 베이스라인보다 우수한 성능을 보였습니다.
추정기 동등성: 집계 꼬리 추정기와 균등 꼬리 추정기는 거의 동일한 결과를 생성하였으며, 이는 본 방법이 특정 꼬리 모델링 가정에 민감하지 않음을 시사합니다.

5. 의의 및 주장

본 논문은 현재 LLM을 위한 발산 기반 RL 방법론의 구조적 불일치를 해결한다고 주장합니다. 근접성과 방향 기준 모두를 동일한 분포적 발산으로부터 도출함으로써, 이 방법은 더 충실한 신뢰 영역 제어를 제공합니다.

저자들은 비록 "위험한 유지" 비율의 개선 폭은 완만하지만(두 기준이 상충하는 토큰 집합이 작기 때문), 이러한 일관성이 이론적으로 타당하며 규모에 관계없이 경험적으로 유익하다고 강조합니다. 결론적으로, 발산 기반 신뢰 영역의 이점을 완전히 활용하기 위해서는 방향 기준 또한 근접성 기준과 동일한 분포적 처리를 받아야 한다고 주장합니다. 이 방법은 기존의 발산 임계값 외에 새로운 하이퍼파라미터를 요구하지 않는, 기존 DPPO 구현에 바로 적용 가능한 드롭인(drop-in) 교체 모델로 제시됩니다.