의존성기반샘플러를활용한마스크확산언어모델향상
마스크 확산 언어 모델(MDLM)의 기존 디코딩은 토큰 수준의 불확실성에만 의존해 전체 시퀀스의 의존성을 충분히 반영하지 못한다. 본 논문은 트랜스포머의 어텐션 행렬을 활용해 토큰 간 의존성을 추정하고, 이를 기반으로 마스크된 토큰을 선택·업데이트하는 훈련‑프리 디코딩 전략인 Dependency‑Oriented Sampler(DOS)를 제안한다. DOS는 블록 구조에 구애받지 않으며, 코드 생성 및 수학 추론 등에서 기존 방법을 능가하는 품질을…
저자: Xueyu Zhou, Yangrong Hu, Jian Huang
본 논문은 마스크 확산 언어 모델(Masked Diffusion Language Models, MDLMs)의 디코딩 효율성과 품질을 동시에 향상시키는 새로운 전략인 Dependency‑Oriented Sampler(DOS)를 제안한다. 서론에서는 기존 대형 언어 모델이 주로 다음 토큰 예측을 기반으로 한 자동 회귀(AR) 방식에 의존해 왔으며, 이는 순차적인 생성으로 인해 병렬 디코딩이 어려운 한계를 가지고 있음을 지적한다. 이러한 문제를 해결하고자 마스크 확산 모델이 제안되었지만, 현재까지 제안된 디코딩 전략들은 토큰‑레벨 불확실성(Confidence, Entropy, Margin Confidence 등)에만 초점을 맞추어 시퀀스‑레벨의 토큰 간 의존성을 충분히 활용하지 못한다는 점을 강조한다.
관련 연구 파트에서는 두 갈래의 기존 접근을 정리한다. 첫 번째는 이산 확산 모델 자체에 대한 연구로, D3PM, CTMC 기반 연속 시간 마스크, 그리고 최근의 대규모 MDLM(LLaD‑A, Dream 등)이 소개된다. 두 번째는 디코딩 전략에 관한 연구로, 무작위 마스크 해제, 토큰‑레벨 불확실성 기반 순서 결정, 블록‑와이즈 병렬 디코딩, 그리고 속도‑정확도 균형을 위한 동적 샘플링 기법(Fast‑dLLM, KLASS, EB‑sampler) 등이 논의된다. 그러나 이들 모두가 토큰 간 상호 의존성을 명시적으로 모델링하지 못한다는 공통된 한계가 있다.
논문의 핵심 이론적 분석은 ‘분포 불일치’를 중심으로 전개된다. MDLM은 마스크된 토큰을 조건부로 예측하도록 훈련되지만, 실제 생성 목표는 전체 시퀀스의 조인트 분포를 복원하는 것이다. 저자는 toy example을 통해 병렬 디코딩이 어떻게 서로 다른 팩터라이제이션을 만들고, 잘못된 블록 구조가 조인트 분포 복원을 방해하는지를 시각화한다. 특히, 마스크된 토큰을 현재 관측된 컨텍스트와 가장 강하게 의존하는 순서대로 해제하면(예시 d) 조인트 분포를 정확히 복원할 수 있음을 보인다.
이론적 통찰을 바탕으로 제안된 DOS는 트랜스포머의 어텐션 행렬을 이용해 토큰 간 의존성을 정량화한다. 구체적인 절차는 다음과 같다. (1) 현재 시점 t에서 모델의 출력과 어텐션 가중치를 추출한다. (2) 다중 헤드 어텐션을 평균해 L×L 크기의 어텐션 매트릭스를 만든다. (3) 마스크된 위치 m에 대해, 비마스크 토큰 집합 U에 대한 어텐션 가중치 합을 계산해 의존성 점수 dep(m)=∑_{u∈U}A_{mu}를 얻는다. (4) 의존성 점수가 높은 토큰을 우선적으로 디코딩한다. 이 과정은 순전파 단계에서만 수행되며 추가 학습이나 파라미터 수정이 필요 없으므로 ‘training‑free’이다.
실험 설정은 LLaD‑A Instruct‑8B 모델을 기반으로 두 주요 벤치마크, HumanEval(코드 생성)과 GSM8K(수학 추론)를 사용한다. 비교 대상은 기존의 Confidence, Entropy, Margin, Fast‑dLLM, KLASS 등이다. 결과는 다음과 같다. (1) 단일 블록(512 토큰) 설정에서도 DOS는 기존 방법보다 일관적으로 높은 정확도를 기록한다. 특히, 블록 크기가 커질수록 성능이 급락하는 Fast‑dLLM과 달리, DOS는 32와 512 모두에서 거의 동일한 성능을 유지한다. (2) 코드 생성에서는 HumanEval 점수가 46.2%→48.5%로 상승했고, 수학 추론에서는 71.3%→74.9%로 개선되었다. (3) DOS를 Fast‑dLLM과 결합하면 샘플링 단계 수를 약 30% 줄이면서도 동일하거나 더 높은 정확도를 달성한다. 이는 DOS가 의존성‑우선 순서를 제공하고, 이후 단계에서 기존의 토큰‑레벨 불확실성 기준을 그대로 활용할 수 있음을 의미한다.
논의에서는 DOS의 장점과 한계를 짚는다. 장점으로는 (가) 블록 구조에 독립적이며, (나) 기존 병렬 샘플링 기법과 호환 가능하고, (다) 추가 연산 비용이 거의 없다는 점을 들었다. 한계로는 어텐션이 실제 의미적 의존성을 완전히 반영하지 않을 가능성, 깊은 레이어와 얕은 레이어 간 어텐션 차이, 대규모 모델에서 어텐션 매트릭스 저장·연산 비용이 증가할 수 있다는 점을 제시한다. 향후 연구 방향으로는 다중 스케일 어텐션 통합, 그래프 기반 의존성 모델링, 어텐션 근사 기법 등을 통해 효율성을 더욱 높이는 방안을 제시한다.
결론적으로, DOS는 MDLM 디코딩에서 ‘어디를 먼저 풀어야 하는가’를 토큰‑레벨 불확실성이 아닌 실제 컨텍스트 의존성에 기반해 결정함으로써, 기존 방법이 놓친 시퀀스‑레벨 구조를 복원한다. 이는 마스크 확산 기반 텍스트 생성이 보다 안정적이고 효율적으로 확장될 수 있는 중요한 발판을 제공한다.
원본 논문
고화질 논문을 불러오는 중입니다...
댓글 및 학술 토론
Loading comments...
의견 남기기