확률적 공동 샘플링을 가능하게 하는 경량 샘플러 ADJUST

초록

이 논문은 마스크된 확산 언어 모델에서 한 번의 전체 모델 순전파로 여러 토큰을 동시에 샘플링하면서도 진정한 공동 분포에 가깝게 추정하는 방법을 제안한다. 기존의 병렬 마스크 해제는 토큰별 주변 분포의 곱으로 샘플링해 품질이 급락했지만, 저자들은 기존 대형 모델 위에 단일 트랜스포머 레이어인 ADJUST를 두어 순차적으로 토큰을 “초안”하고, 각 초안 토큰을 다음 초안에 반영함으로써 근사 공동 샘플링을 구현한다. 실험 결과, 48 토큰을 한 번에 해제해도 MAUVE 점수가 0.50.8 수준으로 크게 회복되었으며, 수학·코딩 벤치마크에서도 정확도가 크게 향상되었다.

상세 분석

마스크된 확산 언어 모델(Diffusion LM)은 입력 시퀀스의 마스크 토큰을 점진적으로 복원하는 방식으로 텍스트를 생성한다. 한 번의 순전파에서 모델은 각 위치마다 주변(마진) 확률 pᵢ(·|x)를 출력한다. 따라서 K개의 토큰을 동시에 해제하면, 실제로는 p₁·p₂·…·p_K 와 같은 독립적인 곱분포에서 샘플링하게 되며, 이는 원래 모델이 내재하고 있는 “진정한” 공동 분포 p*와는 근본적으로 다르다. 이 차이는 특히 토큰 간 상호 의존성이 강한 문맥(예: 수학식, 코드)에서 품질 저하로 직결된다.

저자들은 이 문제를 “근사 공동 샘플링”이라는 새로운 패러다임으로 접근한다. 핵심 아이디어는 대형 확산 모델 f 를 고정하고, 그 위에 아주 가벼운 단일 레이어 트랜스포머 g (ADJUST)를 추가하는 것이다. 전체 흐름은 다음과 같다. ① 기본 모델 f 에 전체 마스크 입력을 넣어 h₀ = f(x) 를 얻는다. ② 첫 번째 토큰은 f 그 자체의 마진 p₁ 에서 직접 샘플링한다. ③ 이후 토큰들은 g 에 현재 임베딩 hₖ 와 이미 샘플된 토큰들을 입력해 새로운 임베딩 hₖ₊₁ = g(hₖ, x ⊕ 샘플) 을 계산하고, 이를 통해 조건부 분포 qₖ 를 얻어 다음 토큰을 샘플링한다. 이렇게 하면 각 토큰이 이전에 샘플된 토큰들의 실제 정체성을 반영하므로, 독립적인 마진 곱이 아니라 점진적으로 업데이트되는 조건부 분포를 따르게 된다.

학습 단계에서는 고정된 f 로부터 다양한 마스크 패턴과 노이즈 레벨의 데이터를 오프라인으로 생성하고, g 가 p* 와 동일한 연쇄 확률을 모방하도록 KL‑다이버전스 기반 손실을 최소화한다. 특히 g 는 “가장 최신 임베딩 h”와 “현재 문자열”을 동시에 입력받아, 토큰 간 상호작용을 학습하도록 설계되었다. 이는 기존의 스펙큘러 디코딩(초안‑검증)과는 달리 별도의 검증 모델을 필요로 하지 않으며, 최종 출력 자체가 원본 확산 모델이 표현한 분포와 다른 새로운 근사 분포가 된다.

실험에서는 Dream‑7B‑Base, Llada‑7B‑Base 등 사전학습 모델과 Dream‑7B‑Instruct, Dream‑7B‑Coder 등 인스트럭션 튜닝 모델을 대상으로, 토큰당 NLL, MAUVE, GSM8K, MBPP, HEval 등 다양한 지표를 평가했다. 48 토큰을 한 번에 해제했을 때도 ADJUST는 기존 병렬 샘플링 대비 NLL을 0.20.4 낮추고, MAUVE를 0.50.8 포인트 상승시켰으며, 특히 수학·코딩 벤치마크에서 정확도가 1520% 포인트 향상되는 효과를 보였다. 속도 면에서는 전체 모델 순전파는 동일하게 유지하면서, 가벼운 g 레이어를 K‑1 번 추가 실행하므로, 순수 병렬 샘플링 대비 20~30% 정도만 느려졌다. 이는 품질‑속도 트레이드오프에서 매우 유리한 위치를 차지한다.

이 논문의 주요 공헌은 (1) 확산 LM에서 “한 번에 여러 토큰을 샘플링하면 주변 분포의 곱을 따르게 된다”는 이론적 분석, (2) 단일 레이어 경량 샘플러를 통해 근사 공동 샘플링을 구현한 설계, (3) 오프라인 데이터 생성·특수 손실 함수를 통한 효과적인 학습 방법, (4) 다양한 모델·태스크에 걸친 실증적 검증이다. 한계점으로는 g 가 여전히 f 에 비해 파라미터가 적어 복잡한 의존성을 완전히 포착하지 못할 수 있으며, 매우 큰 K (예: 16~32)에서는 여전히 품질 저하가 관찰된다. 향후 연구에서는 다중 레이어 g 구조, 동적 K 조정, 그리고 다른 종류의 마스크 스케줄과 결합하는 방안을 탐색할 여지가 있다.