Mask-Aware Policy Gradients for Diffusion Language Models
이 논문은 마스크드 확산 언어 모델(Masked Diffusion Language Model) 생성을 토큰 선택과 위치 마스킹을 공동으로 최적화하는 2단계 의사결정 과정으로 정형화하여, 로그 가능도(log-likelihood)의 난해함을 극복하고 수학적 추론 및 코딩 벤치마크에서 최첨단 성능을 달성하는 강화 학습 프레임워크인 마스크 인식 정책 경사(Mask-Aware Policy Gradients)를 소개한다.
원저자:Haran Raajesh, Kulin Shah, Adam Klivans, Philipp Krähenbühl
당신이 로봇에게 이야기를 쓰거나 수학 문제를 푸는 법을 가르치려 한다고 상상해 보십시오. 오랫동안 이 작업을 수행하는 가장 좋은 방법은 로보트가 사람이 키보드로 타이핑하는 것처럼 한 번에 한 단어씩 쓰게 만드는 것이었습니다. 이것을 "자기회귀적(autoregressive)" 생성이라고 부릅니다. 이는 신뢰할 수 있지만, 고속도로를 가로지르는 달팽이처럼 느립니다. 최근 과학자들은 "확산(Diffusion)"이라 불리는 더 빠른 방법을 발견했습니다. 처음부터 새로 쓰는 대신, 확산 방식의 로봇은 빈 마스크(예: [MASK])로 가득 찬 페이지에서 시작하여 하나씩 빈칸을 채워나가며 문장이 완성될 때까지 점진적으로 내용을 드러냅니다. 이것은 마치 가려진 그림에서 덮개를 하나씩 들어 올려 이미지를 보여주는 퍼즐과 같습니다.
하지만 이 빠른 "확산" 로봇에게 보상을 통해 더 똑똑해지는 법을 가르치는 것(강화 학습이라는 기술)은 악몽과 같았습니다. 기존의 "단어 하나씩" 방식에서는 로봇이 특정 선택을 할 확률이 정확히 얼마였는지 계산하기 쉽습니다. 하지만 "확산" 방식에서 로봇은 매 단계마다 두 가지 선택을 합니다. 즉, 빈 공간에 어떤 단어를 넣을지 결정해야 하고, 동시에 어떤 빈 공간을 다음에 드러낼지도 결정해야 합니다. 이전의 방법들은 로봇에게 가르침을 줄 때 오직 첫 번째 선택(단어)만을 보고 두 번째 선택(드러내는 순서)은 무시하려고 했습니다. 이는 마치 요리사에게 완벽한 식사를 만드는 법을 가르치면서, 재료를 고른 것만 비평하고 재료를 어떤 순서로 썰고 저었는지는 완전히 무시하는 것과 같습니다.
COLM 2026 컨퍼런스에서 발표된 이 논문은 우리가 퍼즐의 거대한 부분을 놓치고 있었다고 제안합니다. 저자인 하란 라제쉬(Haran Raajesh), 쿨린 샤(Kulin Shah) 및 텍사스 대학교 오스틴 캠퍼스의 연구팀은 확산 모델을 진정으로 숙달하려면 로봇이 올바른 단어를 고르는 것뿐만 아니라, 올바른 순서로 내용을 드러내는 것에 대해서도 보상을 주어야 한다고 주장합니다. 그들은 "마스크 인식 정책 경사(Mask-Aware Policy Gradients)"라는 새로운 방법을 개발했습니다. 이 시스템은 단순히 다음 단어를 추측하는 대신, "다음에 어떤 마스크를 들어 올릴 것인가"라는 결정을 로봇 전략의 핵심적인 부분으로 취급합니다. 학습 과정을 이 두 가지 뚜렷한 부분으로 수학적으로 분해함으로써, 그들은 로봇이 훨씬 더 빠르게 학습하고 실수를 덜 한다는 것을 발견했습니다.
결과는 매우 인상적입니다. 그들이 이 새로운 방법을 어려운 수학 문제와 코딩 과제에 테스트했을 때, 로봇은 문제를 해결하는 능력이 눈에 띄게 향상되었습니다. 유명한 수학 테스트인 GSM8K에서 이들의 방법은 87.1%의 정확도에 도달했으며, 코딩 테스트인 MBPP에서는 53.4%를 기록했습니다. 이 수치들은 보상을 사용하여 확산 모델을 가르치려 했던 이전의 그 어떤 방법보다 높습니다. 이 논문은 로봇이 생각을 드러내는 순서, 즉 "마스킹" 결정에 주의를 기울임으로써, 우리가 이러한 빠르고 유연한 AI 모델의 새로운 지능 수준을 끌어올릴 수 있음을 보여줍니다. 이를 통해 모델을 단순히 더 빠르게 만들 뿐만 아니라 더 똑똑하게 만들 수 있습니다.
기술 요약: 확산 언어 모델을 위한 마스크 인식 정책 경사 (Mask-Aware Policy Gradients for Diffusion Language Models)
문제 정의
강화 학습(RL)은 다루기 쉬운 로그 가능도(log-likelihood)를 기반으로 정책 경사(policy gradients)를 최적화함으로써 자기회귀(autoregressive) 대규모 언어 모델(LLM)의 추론 능력을 성공적으로 향러해 왔습니다. 그러나 **마스크 확산 언어 모델(Masked Diffusion Language Models, MDLMs)**에 이러한 방법들을 확장하는 것은 여전히 도전적인 과제입니다. 텍스트를 왼쪽에서 오른쪽으로 순차적으로 생성하는 자기회귀 모델과 달리, MDLM은 완전히 마스킹된 시퀀스로부터 반복적으로 토큰을 언마스킹(unmasking)하며 텍스트를 생성합니다.
핵심적인 어려움은 MDLM의 **로그 가능도 추정의 난해함(intractability)**에 있습니다. 기존 방식들은 이 가능도를 ELBO 변형 모델을 사용하거나 생성 궤적(trajectory)의 토큰 예측 구성 요소만을 모델링하여 근사합니다. 이러한 방법들은 MDLM 생성의 핵심적인 측면인 **위치 언마스킹 순서(order in which positions are unmasked)**를 간과합니다. MDLM의 생성 과정은 각 단계에서 두 가지 별개의 결정(1) 마스킹된 위치에 어떤 토큰을 배치할 것인가, (2) 해당 위치 중 어떤 것을 언마스킹 상태로 유지하고 어떤 것을 다시 마스킹([MASK] 상태로 복구)할 것인가를 포함합니다. 두 번째 결정을 무시하는 기존 방식들은 정책 최적화를 위해 생성 과정에서 활용 가능한 전체 신호를 충분히 활용하지 못합니다.
방법론
저자들은 MDLM 생성을 **2단계 액션 마르코프 결정 과정(two-stage action MDP)**으로 정식화합니다. 각 디노이징(denoising) 단계 t에서 모델은 다음을 수행합니다:
토큰 예측 (Token Prediction): 현재 마스킹된 모든 위치에 대한 토큰을 예측합니다.
위치 선택 (Position Selection/Remasking): 예측된 위치 중 일부를 드러내고 나머지를 [MASK] 상태로 되돌릴 부분 집합을 결정합니다.
이 과정을 통해 정책 경사 최적화를 가능하게 하기 위해, 저자들은 **확률적 리마스킹(Probabilistic Remasking)**을 도입합니다. 표준 MDLM 추론은 종종 확신도 점수(confidence scores)를 기반으로 언마스킹할 위치를 선택할 때 결정론적인 그리디(greedy) top-K 전략을 사용합니다. 이러한 그리디 선택은 미분 불가능하므로, 언마스킹 순서에 대한 경사 기반 최적화를 방해합니다. 제안된 방법은 이를 다음과 같은 확률적 변형으로 대체합니다:
결정론적으로 top-K 위치를 선택하는 대신, 모델은 플랙-루스(Plackett–Luce) 분포로부터 위치의 부분 집합 Ut를 샘플링합니다.
이 분포의 확률은 모델 자체의 토큰 로그 가능도(logits)로부터 직접 유도되며, 온도 파라미터 τ에 의해 스케일링됩니다.
이 공식화는 위치 부분 집합에 대한 미분 가능한 분포를 생성하여, 언마스킹 결정의 로그 확률이 궤적 가능도에 포함될 수 있도록 합니다.
정책 경사 분해 (Policy Gradient Decomposition)
전체 궤적 가능도 πθ(z^∣c)를 토큰 예측 확률과 언마스킹 확률의 곱으로 모델링함으로써, 정책 경사는 자연스럽게 두 개의 항으로 분해됩니다: ∇θJ(θ)=Ez^R(c,zT)t=1∑T토큰경사 (Token Gradient)∇θlogπθ(z^t∣c,zt−1)+마스킹경사 (Masking Gradient)∇θlogpunmask(Ut∣z^t,zt−1,c)
토큰 항 (Token Term): 예측된 토큰의 정확성을 강화합니다.
마스킹 항 (Masking Term): 언마스킹 순서의 선택을 강화합니다.
효율성: 마스킹 항은 모델의 기존 로짓(logits)으로부터 계산되므로, 추가적인 파라미터, 구조적 변경, 또는 추가적인 순방향 패스(forward pass)가 필요하지 않습니다.
저자들은 이를 그룹 시퀀스 정책 최적화(Group Sequence Policy Optimization, GSPO) 알고리즘을 사용하여 구현하였으며, 토큰 및 위치 항에 대해 각각 중요도 비율(importance ratios)을 클리핑(clipping)합니다. 전체 궤적을 평가하는 데 드는 계산 비용을 처리하기 위해, 저자들은 디노이징 단계를 세그먼트로 그룹화하여 궤적 구조를 보존하면서 순방향 패스의 수를 줄이는 StepMerge 근사법을 채택했습니다.
주요 기여
MDLM을 2단계 MDP로 정식화: 본 논문은 MDLM 생성이 토큰 선택과 위치 선택을 모두 포함하며, 정책 경사가 그에 따라 분해된다는 이론적 프레임워크를 제공합니다.
확률적 리마스킹: 모델의 기존 로짓을 기반으로 한 미분 가능한 언마스킹 분포(Plackett–Luce)를 도출하여, 구조적 변경 없이 언마스킹 순서를 최적화할 수 있게 했습니다.
이론적 통찰: 토큰 확률이 변하지 않더라도 위치 성분의 정책 경사를 무시할 경우 기대 보상(expected return)을 개선할 수 있는 방향을 놓칠 수 있음을 입증했습니다.
최첨단 성능 (SOTA): 제안된 방법이 기존의 궤적 기반 및 ELBO 기반 방법들보다 일관되게 향상된 성능을 보임을 실험적으로 검증했습니다.
실험 결과
본 방법은 LLaDA-8B-Instruct를 기본 모델로 사용하여 수학적 추론(GSM8K, MATH500) 및 코드 생성(HumanEval, MBPP) 벤치마크에서 평가되었습니다.
성능 향상: 제안된 방법은 모든 벤치마크에서 새로운 SOTA 결과를 달성했습니다.
GSM8K: 정확도 87.1% (생성 길이 128 기준), 이전 최고 성능인 StepMerge 대비 +2.5% 향상.
MATH500: 정확도 53.4%, +4.0% 향상.
HumanEval: 정확도 47.1%, +2.2% 향상.
MBPP: 정확도 53.4%, +2.5% 향상.
절제 연구 (Ablation Studies):
블록 크기 (Block Size): 제안된 방법과 베이스라인 간의 성능 격차는 추론 블록 크기가 커질수록(32에서 전체 시퀀스까지) 넓어졌습니다. 이는 위치 선택 결정 공간의 복잡성이 증가함에 따라 위치 선택을 최적화하는 것의 이점이 커짐을 시사합니다.
효율성: 제안된 방법은 강력한 ELBO 기반 베이스라인인 SPG와 비교했을 때 실제 시간(wall-clock time) 기준으로 더 빠르게 수렴했습니다. SPG의 최종 정확도에 도달하는 데 약 15시간이 걸린 반면 SPG는 약 18시간이 걸렸으며, 최종 정확도 또한 더 높았습니다 (GSM8K에서 81.0% vs 78.5%).
일반화: 이 접근 방식은 다른 베이스 모델(Dream-7B)과 계획(planning) 작업(Sudoku, Countdown)에도 일반화되었으며, 일관되게 베이스라인을 능가했습니다.
DCoLT와의 비교: 위치 선택을 위한 별도의 헤드를 사용하는 LLaDOU와 비교했을 때, 제안된 방법은 추가 모듈이 필요하지 않아 훨씬 낮은 훈련 비용(~160 GPU 시간 vs ~800 GPU 시간)으로 유사하거나 더 나은 정확도를 달 achievement 했습니다.
의의 및 주장
본 논문은 MDLM의 유연성이 자기회귀 모델보다 더 풍부한 궤적 구조, 즉 위치를 언제 언마스킹할 것인가라는 결정과 함께 온다는 점을 주장합니다. 이 결정을 정책의 일부로 취급하고 미분 가능한 확률적 프레임워크를 통해 최적화함으로써, 저자들은 **토큰 예측과 위치 선택의 공동 최적화(joint optimization)**가 RL 환경에서 확산 언어 모델의 성능을 극대화하는 데 필수적임을 입증했습니다.
저자들은 이 접근 방식이 기존의 로짓을 활용하므로 새로운 파라미터나 구조적 변경을 요구하지 않아 실용적으로 효율적이며, StepMerge와 같은 표준 MDLM 훈련 인프라와 호환된다는 점을 강조합니다. 결과는 "마스킹 경사(masking gradient)"가 기존의 토큰 가능도에만 집중했던 방법들이 포착하지 못했던 의미 있는 훈련 신호를 제공한다는 것을 시사합니다.