정답을 이용한 마스크 해제 순서 학습으로 향상된 확산 언어 모델

초록

본 논문은 마스크 확산 언어 모델(MDLM)에서 “어디를 먼저 풀어줄지”(where‑to‑unmask) 결정이 생성 품질에 미치는 영향을 규명한다. 정답 토큰을 활용해 각 위치의 확신 정도를 나타내는 Gt‑Margin 점수를 정의하고, 이를 기반으로 만든 최적 해제 순서가 논리 추론 벤치마크에서 크게 성능을 끌어올림을 보였다. 이후 학습‑to‑rank 방식을 이용해 Gt‑Margin을 모방하는 감독 플래너를 학습시켜, 실제 추론 시 히스토리 기반 히어리스틱 대신 사용함으로써 토큰 예측 모델을 변경하지 않고도 추론 정확도를 개선한다.

상세 분석

Masked Diffusion Language Models(MDLM)은 텍스트를 마스크‑복원 과정으로 생성한다는 점에서 기존의 왼쪽‑우측 자동회귀 모델과 근본적으로 다르다. MDLM은 매 단계 두 가지 선택을 해야 하는데, 하나는 “어디를 먼저 마스크 해제할지”(where‑to‑unmask)이고, 다른 하나는 “그 위치에 어떤 토큰을 넣을지”(what‑to‑unmask)이다. 기존 학습은 what‑to‑unmask, 즉 토큰 예측에만 초점을 맞추고, where‑to‑unmask은 모델의 불확실성(예: 최고 확률, 마진) 기반 히어리스틱이나 비용이 큰 강화학습(RL)으로 처리한다. 그러나 이러한 접근은 최적의 해제 순서를 보장하지 못한다.

논문은 먼저 Gt‑Margin이라는 새로운 점수를 제안한다. Gt‑Margin은 현재 마스크된 상태에서 모델이 정답 토큰에 할당한 확률과 두 번째로 높은 토큰 확률의 차이로 정의된다. 이 값이 클수록 해당 위치를 지금 바로 해제해도 올바른 토큰이 선택될 확률이 높으며, 작은 값은 아직 충분한 컨텍스트가 제공되지 않아 해제를 미뤄야 함을 의미한다. 따라서 Gt‑Margin을 내림차순으로 정렬하면 “쉬운 → 어려운” 순서의 최적 해제 순서가 얻어진다.

실험에서는 제어된 디코딩 프로토콜(한 번에 한 위치만 해제하고, greedy argmax로 토큰을 선택) 하에 Gt‑Margin 기반 순서가 기존 Top‑Prob, Margin 히어리스틱보다 크게 정확도를 향상시켰다. 특히 GSM8K, Math, Sudoku, StrategyQA와 같은 논리·수리 추론 벤치마크에서 초기 몇 단계의 해제 순서가 전체 성능에 미치는 영향이 크게 드러났다.

그 다음 단계로, Gt‑Margin을 직접 사용할 수 없는 실제 추론 상황을 위해 학습‑to‑rank 프레임워크를 도입한다. 각 마스크된 위치의 특징(예: 토큰 임베딩, 주변 컨텍스트)을 입력으로 하는 플래너 f_ϕ를 학습시키고, 목표 레이블은 Gt‑Margin에 의해 정의된 순위이다. 논문은 리스트와이즈 NDCG의 미분 가능한 근사인 PiRank 손실을 사용해 플래너를 최적화한다. 학습된 플래너는 추론 시 현재 마스크된 상태만을 보고 순위를 예측하므로, 별도의 온‑폴리시 롤아웃 없이도 Gt‑Margin에 근접한 해제 순서를 제공한다.

결과적으로, 플래너를 적용한 MDLM 샘플링은 토큰 예측 모델을 그대로 유지하면서도 기존 히어리스틱 대비 평균 2~5%p(percentage point)의 정확도 향상을 달성했다. 특히 초기 단계에서 “쉬운” 위치를 먼저 풀어줌으로써 모델이 더 풍부한 컨텍스트를 확보하고, 이후 어려운 위치를 정확히 복원할 수 있게 된다. 이 접근법은 RL 기반 플래너가 요구하는 높은 계산 비용을 크게 낮추면서도 비슷하거나 더 나은 성능을 보이는 점에서 실용적이다.

또한, 논문은 Gt‑Margin이 모델의 불확실성을 정량화하는 새로운 지표로 활용될 수 있음을 시사한다. 향후 연구에서는 Gt‑Margin을 이용해 동적 마스크 스케줄링, 멀티턴 대화 생성, 혹은 편집 작업 등 다양한 텍스트 조작 시나리오에 적용할 가능성을 제시한다.