기술 요약: 적응형 블록 확산 (Adaptive Block Diffusion, ABD)

1. 문제 정의

확산 언어 모델(Diffusion Language Models, DLMs)은 병렬 디코딩을 가능하게 하여 추론 지연 시간을 줄일 수 있는 잠재력을 가진, 자기회귀(Autoregressive, AR) 모델의 유망한 대안을 제공한다. 그러나 MDLM(Masked Discrete Diffusion Language Models) 및 블록 확산(예: BD3LM)과 같은 최근의 구조적 확산 방식은 고정된 컨텍스트 구조에 의존한다. 이러한 모델들에서 노이즈를 제거(denoising)해야 하는 토큰의 집합은 훈련 과정 동안 미리 결정된 패턴(예: 고정된 블록 크기 및 정렬된 접두사)을 따른다.

이러한 설계는 근본적인 **훈련-추론 불일치(training–inference mismatch)**를 야기한다:

• 훈련: 모델은 특정 구성(the "training grid")에 대해 최적화된 노이즈 제거 위험을 최적화한다.
• 추론: 실제 디코딩 전략은 임의의 접두사 길이와 윈도우 크기(예: 슬라이딩 윈도우, 적응형 블록 크기 또는 스트리밍)를 필요로 하는 경우가 많다.

고정 블록 모델을 훈련된 범위를 벗어난 구성("off-grid")에서 평가할 때, 이들은 퍼플렉시티(perplexity)의 급격한 저하와 불안정한 생성 동작을 보인다. 이들은 일반적인 노이즈 제거 절차를 학습하는 대신, 훈련 중에 본 특정 구성에만 특화된다. 본 논문은 단일 확산 모델이 이러한 불일치를 제거하여 디코딩 구성에 대해 불변성을 유지할 수 있어야 한다고 주장한다.

2. 방법론: 적응형 블록 확산 (ABD)

저자들은 디노이징 구성을 고정된 하이퍼파라미터가 아닌 확률 변수로 취급함으로써 훈련-추론 불일치를 해결하는 적응형 블록 확산 (ABD) 프레임워크를 제안한다.

핵심 메커니즘

고정된 블록 정렬 서브셋으로 훈련을 제한하는 대신, ABD는 구성 공간 $\Omega := \{(k, \ell) \mid 0 \le k \le L-1, 1 \le \ell \le L-k\}$를 정의한다. 여기서 $k$는 접두사 길이이고 $\ell$은 디노이징 윈도우 길이이다.

• 확률적 훈련: 모델은 $\Omega$ 위에 정의된 분포 $\pi$에 따라 훈련된다. 각 훈련 단계에서 구성 $(k, \ell)$은 $\pi$로부터 샘플링된다.
• 디노이징 과정: 인덱스 $S(k, \ell) = \{k+1, \dots, k+\ell\}$에 해당하는 토큰들은 순방향 확산 과정에 의해 오염(corrupted)되는 반면, 접두사 $x_{1:k}$는 깨끗한 상태를 유지하며 조건부 컨텍스트 역할을 한다.
• 목적 함수: 훈련 목적 함수는 분포 $\pi$에 대한 기대 노이즈 제거 위험을 최소화한다:
  $$\mathcal{L}_{ABD}(\theta) = \mathbb{E}_{(k,\ell)\sim\pi, x_0\sim D, t} \left[ \sum_{i \in S(k,\ell)} \text{KL}\left(q(x^i_0 | x^i_t, t) \parallel p_\theta(x^i_0 | x^i_t, x^i_{1:k}, t)\right) \right]$$

구현 세부 사항

• 아키텍처: ABD는 기본 확산 모델(예: 트랜스포머)에 대한 아키텍처 변경을 요구하지 않는다. 오직 훈련 중에 마주치는 구성의 분포만을 수정한다.
• 어텐션 마스킹: 각 단계마다 샘플링된 $(k, \ell)$에 기반하여 동적 어텐션 마스크가 구축된다. 이 마스크는 노이즈가 섞인 토큰들이 동일한 블록 내의 다른 노이즈 토큰들과 이전 블록의 깨끗한 토큰들에만 어텐션을 수행하도록 보장하며, 깨끗한 토큰들은 이전의 깨끗한 토큰들에 대해 인과적(causal)으로 어텐션을 수행한다.
• 분포 선택: 논문은 $\pi$에 대한 다양한 분포(범주형 지수 분포, 균등 분포, 로그-정규 분포)를 탐색한다. 주요 실험에서는 2의 거듭제곱 $\{1, 2, 4, \dots, 128\}$에 대한 범주형 분포를 사용하며, 이는 작은 블록 크기를 우선시하도록 지수적 감소를 적용한다.

3. 이론적 분석

본 논문은 훈련 분포의 서포트(support)와 일반화 성능 사이의 관계를 연결하는 공식적인 보증을 제공한다.

• 위험 최소화: 저자들은 확산 훈련을 구성 공간 $\Omega$에 대한 위험 최소화로 정식화한다.
• 정리 1 (최적성): ABD 목적 함수의 모든 전역 최소값은 $\pi$-거의 모든(almost every) 구성 $(k, \ell)$에 대해 최적의 노이즈 제거 성능을 달성한다. $\pi$의 서포트 밖에 있는 구성들은 제약되지 않은 상태로 남는다.
• 정리 2 (훈련-추론 정렬): 만약 추론 정책이 $\Pi_{inf} \ll \pi$ (즉, 추론 서포트가 훈련 서포트에 포함됨)를 유도하는 분포를 생성한다면, $\pi$ 하에서 위험을 최소화하도록 훈련된 모델은 추론 위험에 대해서도 전역 최소값이 된다.
• 시사점: 일반화는 모델 아키텍처의 속성이 아니라 훈련 분포의 서포트 커버리지의 속성이다. 고정 블록 모델은 임의의 디코딩 전략에 대해 $\Pi_{inf} \not\ll \pi_{fixed}$이기 때문에 실패한다.

4. 실험 결과

저자들은 ABD를 LM1B (One Billion Word) 및 OpenWebText (OWT) 벤치마크에서 평가하며, 이를 자기회귀(AR), MDLM, SEDD 및 고정 블록 전문가(BD3LM)와 비교한다.

주요 발견

1. 구조적 불변성:
   • 고정 블록 전문가(예: 블록 크기 4로 훈련된 BD3LM)는 훈련된 구성을 초과하는 블록 크기에서 평가될 때 급격한 퍼플렉시티 스파이크를 보인다.
   • ABD는 전체 범위에 걸쳐 블록 크기와 퍼플렉시티 사이의 **단조 관계(monotonic relationship)**를 회복한다. 블록 크기가 작아짐에 따라(자기회귀 정밀도에 접근함에 따라) 오프그리드 붕괴 없이 퍼플렉시티가 단조적으로 개선된다.
2. 성능 동등성:
   • 단일 ABD 모델은 특정 "홈" 구성에서 고정 블록 전문가와 대등하거나 더 나은 성능을 보인다. 예를 들어, LM1B에서 광범위한 분포로 훈련된 ABD는 특정 크기에 맞춰 훈련된 BD3LM 전문가와 비교했을 때 블록 크기 4 및 16에서 경쟁력 있는 퍼플렉시티를 달행한다.
3. 제로샷 일반화:
   • 홀드아웃 벤치마크(PTB, Wikitext, Lambada 등)에서 ABD는 강력한 일반화 능력을 보여준다. 특히, 블록 크기 4에서의 ABD는 과학 분야(Pubmed, Arxiv)에서 BD3LM 전문가보다 우수한 성능을 보이는데, 이는 멀티 스케일 훈련이 도메인 변화에 대한 강건성을 향상시킴을 시사한다.
4. 생성 품질:
   • 생성 퍼플렉시티(GPT2-Large를 통해 측정) 측면에서, ABD는 동등한 계산 비용(NFE)에서 가장 강력한 고정 블록 전문가의 품질과 일치하며, 이는 광범위한 훈련이 샘플 품질을 저하시지 않음을 확인시켜 준다.

절제 연구 (Ablation Studies)

• 분포 민감도: $\pi$의 선택을 통해 실무자는 스케일 간의 트레이드오프를 조절할 수 있다. 지수적 가중치는 작은 블록을 선호하고, 균등 가중치는 큰 블록을 선호하며, 로그-정규 분포는 균형을 제공한다.
• 훈련 역학: 멀티 스케일 훈련은 안정적인 정규화 도구로 작용하여, 단일 스케일에 특화되기보다는 모든 블록 크기에서 동시에 성능을 향상시킨다.
• 이산 vs 연속: 2의 거듭제곱에 대한 이산 분포가 연속 분포보다 더 다루기 쉽고 성능이 좋은 것으로 나타났는데, 이는 이들이 서로 다른 작동 지점을 나타내기 때문인 것으로 보인다.

5. 의의 및 주장

본 논문은 **구성 서포트(configuration support)**가 확산 언어 모델의 일반화를 결정하는 핵심 요소라고 주장한다. 고정된 아키텍처 제약에서 훈련 구성의 분포으로 초점을 전환함으로써, ABD는 구조적 제약이 있는 목적 함수에 대한 원칙적인 대안을 확립한다.

• 불일치 해결: ABD는 훈련-추론 불일치를 성공적으로 제거하여, 아키텍처 변경이나 추론 시 휴리스틱 없이도 단일 모델이 임의의 디코딩 전략(슬라이딩 윈도우, 적응형 블록)을 처리할 수 있게 한다.
• 통합 모델: 단일 모델이 구조적 불변성을 달성하여, 모든 디코딩 스케일에서 높은 품질을 유지하면서 블록 크기와 퍼플렉시티 사이의 기대되는 단조 관계를 회복할 수 있음을 입증한다.
• 한계점: 저자들은 성능이 여전히 $\pi$의 선택(스케일 간의 트레이드오프를 결정함)에 의존한다는 점과, 추론 효율성(테스트 시 최적의 정책을 선택하는 문제)이 여전히 미해결 과제로 남아 있다는 점을 겸허히 언급한다. 이론적 보증은 서포트 커버리지에 기반하며 유한 샘플 경계(finite-sample bounds)를 제공하지는 않는다.

결론적으로, 적응형 블록 확산은 디코딩 구성에 불변하는 확산 언어 모델을 훈련하기 위한 강력한 프레임워크를 제공하며, 현재의 구조적 확산 방식이 가진 결정적인 한계를 해결한다.