기술 요약: 속도 계획: 마스킹 확산 언어 모델을 위한 확장 스케줄링

1. 문제 제기

마스킹 확산 언어 모델 (MDLM) 은 병렬적이고 비자기회귀적인 텍스트 생성을 가능하게 함으로써 기존 자기회귀 (AR) 모델에 대한 유망한 대안을 제시합니다. 그러나 추론 효율성 측면에서 치명적인 병목 현상이 존재합니다. MDLM 은 이론적으로 임의의 순서로 생성할 수 있지만, 기존 샘플링 전략들은 이러한 잠재력을 실현하지 못하는 경우가 많습니다.

현재 접근 방식은 일반적으로 모델의 신뢰도나 엔트로피에 기반하여 마스크를 해제할 토큰을 선택하는 **자기 계획기 (self-planners)**에 의존합니다. 해당 논문은 이러한 신뢰도 기반 계획기가 우연히 자기회귀적 행동을 모방한다고 주장합니다. 가장 신뢰도가 높은 토큰들 (종종 공간적으로 인접한 토큰들) 을 탐욕적으로 선택함으로써 병렬 마스크 해제 단계 간의 상호작용을 무시하게 됩니다. 이로 인해 모델이 실제로 느린 토큰 단위 생성으로 회귀하는 '단시안적 (short-sighted)' 과정이 발생하며, 길이 $G$인 시퀀스에 대해 $O(G)$개의 디노이저 호출이 필요하게 됩니다.

반-자기회귀 블록 확산과 같은 대안적 전략들은 시퀀스를 연속된 구간으로 나누지만, 여전히 크기 $B$인 블록당 $O(B)$개의 디노이저 호출을 요구하여 병렬성을 완전히 활용하지 못합니다. furthermore, 최근 연구들은 마스크 해제를 추론 시간 계획 문제로 규정하며 종종 외부 가이드 모델이나 복잡한 학습 목적 함수를 요구하여, 즉시 적용 가능한 솔루션으로서의 실용성을 제한합니다.

2. 방법론: 확장 마스크 해제 스케줄러 (DUS)

저자들은 기존 디노이저를 수정하거나 재학습을 요구하지 않고 각 디노이징 단계에서 결합 엔트로피 증가를 최소화하도록 설계된, 순수한 추론 시간 및 모델 중립적 계획기인 **확장 마스크 해제 스케줄러 (Dilated Unmasking Scheduler, DUS)**를 제안합니다.

핵심 개념

DUS 는 크기 $B$인 시퀀스 블록을 인접하지 않은 **확장된 그룹 (dilated groups)**으로 분할하고 이를 병렬로 마스크 해제합니다. 스케줄러는 계산 가능한 합-주변분 (sum-of-marginals) 목적 함수와 계산 불가능한 결합 엔트로피 목적 함수 사이의 **엔트로피 간격 ($\Delta$)**을 최소화하는 것이 고품질 병렬 샘플링의 핵심이라는 원리에 기반합니다.

$$\Delta(I_t; S_t) = \sum_{i \in I_t} H(X_i | S_t) - H(X_{I_t} | S_t)$$

여기서 $I_t$는 $t$단계에서 마스크가 해제된 인덱스 집합이고, $S_t$는 현재 상태입니다. 선택된 토큰들이 조건부 독립일 때 이 간격은 최소화됩니다.

알고리즘 설계

DUS 는 결정론적이고 coarse-to-fine(대략적에서 정밀하게) 한 스케줄을 사용합니다:

1. 확장 패턴: 모델 점수에 기반하여 토큰을 선택하는 대신, DUS 는 고정된 확장 패턴을 사용하여 마스크를 해제할 인덱스를 미리 정의합니다.
2. 반복적 정제: 크기 $B$인 블록의 경우, 스케줄러는 $R = \lceil \log_a B \rceil$회 반복 (여기서 $a > 1$은 일반적으로 2 인 밑수) 을 통해 블록을 완료합니다.
   • 초기 반복: 넓은 간격으로 떨어진 희소 토큰 집합 (큰 스트라이드 $s_t$) 을 마스크 해제합니다. 이는 빠른 혼합을 보이는 가변 길이 마르코프 체인 (VLMC) 에 대한 가정 하에 동시에 샘플링되는 토큰 간의 조건부 의존성을 감소시킵니다.
   • 후기 반복: 노출된 이웃으로 인해 컨텍스트 $S_t$가 더 풍부해짐에 따라 더 높은 해상도로 나머지 간격을 채웁니다.
3. 복잡도: 이 접근 방식은 블록당 디노이저 호출 횟수를 $O(B)$에서 $O(\log B)$로 줄여, 총 복잡도를 $O(\frac{G}{B} \log B)$로 달성합니다.

이론적 보장

빠른 혼합 VLMC 에 대한 가정 하에, 논문은 인덱스들을 거리 $d$만큼 간격을 두면 그들 간의 상호 정보가 기하급수적으로 감소함을 증명합니다. 결과적으로, 넓게 간격을 둔 토큰들의 결합 조건부 엔트로피는 개별 조건부 엔트로피들의 합에 접근하게 되며, 이는 병렬 샘플링을 위한 주변 엔트로피 대리 변수의 사용을 유효하게 입증합니다.

3. 주요 기여

• 추론 전용, 모델 중립적 디코딩: DUS 는 모델 아키텍처, 가중치, 또는 학습 절차에 제로 수정을 요구하는 '즉시 적용 (drop-in)' 전략입니다.
• 로그arithmic 마스크 해제 스케줄: 이는 지역적 컨텍스트를 존중하고 $O(\log B)$회 반복 내에서 결합 엔트로피를 최소화하는 결정론적이고 coarse-to-fine 한 확장 전략을 도입합니다.
• 이론적 분석: 저자들은 빠른 혼합 가정 하에서 DUS 가 각 반복에서 결합 엔트로피 한계에 접근함을 보여주는 이론적 경계를 제시하며, 이는 군집화된 토큰 선택으로 인해 큰 엔트로피 간격을 초래할 수 있는 기준선과 대조됩니다.
• 보완적 하이브리드 확장: 논문은 확장 간격 원리가 적응형 샘플러 (Entropy-Bounded 및 Confidence-Bounded) 에 대한 후처리 필터로 적용될 수 있음을 보여주며, 동일하거나 약간 더 높은 계산 예산 하에서 정확도를 향상시킵니다.

4. 실험 결과

저자들은 네 가지 벤치마크 카테고리에서 다섯 가지 MDLM 변형 (LLaDA-8B, Dream-7B, DiffuCoder-7B) 에 대해 DUS 를 평가했습니다:

• 수학: GSM8K, MATH500
• 코드: HumanEval, MBPP
• 일반 지식: BBH, MMLU-Pro
• 지시 따르기: IFEval

주요 발견 사항

• 속도 - 정확도 트레이드오프: DUS 는 기존 신뢰도 기반 계획기보다 일관되게 우수합니다. 동등한 속도 향상 시, DUS 는 수학 벤치마크에서 최대 27%, 코드 벤치마크에서 **20%**까지 하류 작업 정확도를 향상시킵니다.
• 월시계 속도 향상: DUS 는 디노이저를 수정하지 않고도 토큰 단위 MDLM 디코딩 대비 최대 5.8 배의 월시계 속도 향상을 달성합니다. 예를 들어, LLaDA-8B 를 사용한 GSM8K 에서 DUS(블록 크기 $B=32$) 는 5.8 배 속도 향상과 함께 49.36% 정확도를 달성하는 반면, 토큰 단위 기준선은 72.63% 정확도를 보이지만 훨씬 느립니다.
• AR 과의 비교: 특정 구성 (예: GSM8K 에서 DUS 를 사용한 LLaDA-8B) 에서 이 방법은 경쟁력 있는 정확도를 유지하면서 자기회귀 Llama-3-8B 기준선보다 5.6 배 빠릅니다.
• 강건성: 이 방법은 코드 생성 및 일반 지식 추론 등 다양한 도메인에서 잘 수행되어, 엄격한 빠른 혼합 가정을 위반할 수 있는 작업에서도 확장 스케줄의 유효성을 입증합니다.
• 하이브리드 개선: 적응형 샘플러 (EB/CB) 에 대한 후처리 필터로 적용될 때, 확장 간격은 밀집된 토큰의 선택을 방지함으로써 정확도를 향상시킵니다 (예: EB 의 경우 HumanEval 에서 +13.4%).

5. 중요성 및 주장

이 논문은 DUS 가 마스킹 확산 모델에 내재된 '품질 - 속도 트레이드오프'를 성공적으로 해결한다고 주장합니다. 성능을 저하시키는 일반적인 가속화 방법과 달리, DUS 는 기존 병렬 샘플링 전략에 비해 추론을 가속화할 뿐만 아니라 출력 품질도 향상시킵니다.

저자들은 DUS 가 다음과 같은 것들 없이 확산 디코딩의 고유한 병렬성을 해제한다고 강조합니다:

• 모델 재학습.
• 보조 계획기 모델 추가 (BERT 기반 계획기 등).
• 디노이저의 학습 목적 함수 수정.

품질 - 속도 트레이드오프를 블록 크기 $B$의 결정론적이고 예측 가능한 함수로 변환함으로써, DUS 는 효율적이고 고품질의 마스킹 확산 언어 모델을 배포하기 위한 실용적인 경로를 제공합니다. 이 방법은 단계당 비용을 줄이는 방식으로 직교적으로 작동하는 KV 캐싱과 같은 다른 효율성 향상 기법들을 보완하는 기술로 제시됩니다.