기술 요약: 비자율 생성을 위한 이산 확률적 국소화

1. 문제 제기

연속 확산 모델은 이미지, 비디오, 오디오의 고차원 생성 분야에서 표준으로 자리 잡았습니다. 그러나 토큰 임베딩에 대한 연속 확산을 통한 언어 생성 적용은 일관되게 마스크된 이산 확산 모델 (MDMs) 보다 낮은 성능을 보였습니다. 기존 연속 확산 언어 모델들은 일반적으로 분포 충실도 (MAUVE 와 같은 지표로 측정) 측면에서 MDM 들보다 뒤처집니다.

저자들은 병목 현상이 연속성 자체에 있는 것이 아니라 사용된 표현 (representation) 에 있다고 주장합니다. 표준 접근법 (분산 보존 또는 분산 폭발) 은 무정보 한계를 마스킹과 유사한 토큰 상태가 아닌 등방성 가우시안 노이즈로 표현합니다. 또한, 이러한 프레임워크에서 노이즈 제거기는 노이즈가 포함된 임베딩 내 증거의 양을 보정하기 위해 시간 단계 또는 노이즈 수준 레이블을 받아야 합니다. 이로 인해 노이즈 제거기는 시간 조건에 매우 민감해지며, 마스킹된 상태와 불확실한 상태에 대해 구별된 처리가 필요합니다. 결과적으로 연속 확산 모델은 유한한 신호대노이즈비 (SNR) 상태의 연속체가 제공하는 풍부한 감독 정보를 활용하지 못했습니다.

2. 방법론: 이산 확률적 국소화 (DSL)

본 논문은 노이즈 제거 과정을 명목상 SNR 에 무관하게 만드는 연속 상태 프레임워크인 이산 확률적 국소화 (Discrete Stochastic Localization, DSL) 를 소개합니다.

핵심 설계 선택

1. 단위 구 토큰 임베딩: DSL 은 모든 깨끗한 토큰 임베딩이 단위 구 위에 있도록 ($\|x_i\|_2 = 1$) 제한합니다. 이는 손실 정규화기가 아닌 기하학적 구조적 제약입니다.
2. 확률적 국소화 채널: 오염 과정은 $z_t = t x + \sqrt{t} \epsilon$로 정의되며, 여기서 $t$는 SNR 을 나타냅니다. 단위 구 제약 하에서 베이즈 사후분포에 포함된 $x$의 노름과 관련된 항은 상수가 되어 정규화 과정에서 상쇄됩니다.
3. SNR 불변 노이즈 제거기: 이 기하학 하에서 베이즈 최적 노이즈 제거기 (사후 평균) 는 명목상 SNR $t$가 아닌 노이즈가 포함된 토큰 상태 $z$에만 의존한다는 것이 핵심 이론적 결과입니다.
   $$\hat{x}(z, t) = \hat{x}(z) = \frac{\mathbb{E}_P[x e^{x \cdot z}]}{\mathbb{E}_P[e^{x \cdot z}]}$$
   이를 통해 단일 시간 무관 네트워크가 모든 SNR 영역 (완전 마스킹에서 거의 깨끗한 상태까지) 의 노이즈 제거기로 기능할 수 있습니다.

학습 및 아키텍처

• 혼합 지원 학습: 연속 경로 적분 가능도와 마스킹 확산의 엔드포인트 가능도를 모두 활용하기 위해 DSL 은 혼합 지원 학습 목표를 사용합니다. 두 가지 브랜치에서 SNR 을 샘플링합니다:
  • 연속 경로: 로그 정규 분포에서 추출된 SNR (유한 SNR 상태).
  • ROAR (랜덤 순서 자율 생성) 엔드포인트: 극단적인 값에서 추출된 SNR (마스킹된 경우 0 에 가깝고, 노출된 경우 $\infty$에 가깝습니다).
    두 브랜치 모두 토큰 사후분포에 대한 단일 교차 엔트로피 손실로 학습되며, 이는 임베딩 공간 최적화 병리 (예: 임베딩 붕괴) 를 피하면서 경로 적분 평균 제곱 오차 (MSE) 의 대용량 역할을 합니다.
• 사후 뷰 변환기: 노이즈 상태 $z$는 연속 임베딩 공간에 존재하지만 백본 (Transformer/DiT) 은 토큰과 유사한 입력을 기대하므로, 변환기가 $z$를 "토큰 혼합" 표현으로 매핑합니다. 이는 노이즈 상태를 단위 노름 임베딩의 유계 조합으로 매핑하여 사후 기하학 (방향은 토큰 선호도를 제어하고 노름은 신뢰도를 제어함) 을 보존하면서 안정적인 자기 주의를 보장합니다.

추론 유연성

노이즈 제거기가 시간 무관하므로 DSL 은 단일 생성 실행 내에서 임의의 토큰별 SNR 경로를 지원합니다:

• 마스킹 정제: 표준 반복적 재마스킹 (ReMDM).
• 랜덤 순서 자율 생성 (ROAR): 재마스킹 없이 무작위 순서로 토큰을 하나씩 노출합니다.
• 하이브리드 샘플링: 연속 노이즈 제거 단계 후 이산 마스킹 정제 단계를 거칩니다.

3. 주요 기여

1. 이산 생성을 위한 연속 상태 프레임워크: DSL 은 노이즈가 있는 토큰에 대한 베이즈 최적 노이즈 제거기가 노이즈 상태만의 함수가 되는 단위 구 임베딩 위의 확률적 국소화 채널을 도입합니다. 이를 통해 단일 시간 무관 네트워크가 토큰별 SNR 경로 전체를 지원할 수 있습니다.
2. 실용적인 학습 레시피: 저자들은 엔드포인트 (마스킹/노출) 와 유한 SNR 상태 모두에서 동일한 노이즈 제거기를 학습하는 혼합 지원 SNR 샘플링 전략을 제안합니다. 표준 Transformer 아키텍처와 연속 사후 뷰를 정렬하기 위해 토큰 혼합 변환기를 활용합니다.
3. 경로 전반의 증거: 단일 DSL 체크포인트는 마스킹 정제 하에서 OpenWebText 의 몇 단계 생성 성능을 향상시키고, 재학습 없이 랜덤 순서 AR 샘플링을 지원하며, Text8 에서 경쟁력 있는 가능도를 달성함을 보여줍니다.

4. 실험 결과

저자들은 OpenWebText (OWT) 와 Text8 에서 사전 훈련된 MDLM 체크포인트를 미세 조정하여 DSL 을 평가했습니다.

• OpenWebText 생성 (MAUVE):
  • 비정제: 반복적 재마스킹 없이도 DSL 미세 조정 체크포인트는 기준 MDLM 보다 훨씬 우수한 성능을 보였습니다. $T=1024$에서 MAUVE 는 0.042(MDLM) 에서 0.495(DSL-FT + MDLM) 및 0.551(DSL-FT + ROAR) 로 향상되었습니다.
  • 반복 정제 (ReMDM): ReMDM 계열 디코더와 결합되었을 때, DSL 은 비참조 기준 중 최첨단 결과를 달성했습니다. 가장 좋은 구성 (원칙적인 $\eta_{cap}$을 사용한 ReMDM-루프) 은 $T=1024$에서 MAUVE 0.722에 도달하여 PRISM-loop(0.527) 및 ADLM(0.699) 과 같은 이전 방법들을 능가했습니다.
  • 소수 단계 효율성: DSL 은 낮은 단계 예산에서도 강력한 성능을 보였습니다. $T=128$에서 DSL-FT + ReMDM-루프는 0.639 MAUVE 를 달성한 반면, 기준 MDLM 은 0.015 였습니다.
• 하이브리드 샘플링: 하이브리드 샘플러 (연속 노이즈 제거 후 이산 정제) 는 총 48 단계 만에 MAUVE 0.662를 달성하여 $T=1024$의 비정제 MDLM 결과를 초과했습니다.
• Text8 가능도: DSL 은 $\le$ 1.45의 문자당 비트 (BPC) 를 달성하여 마스킹 이산 확산 모델 (MDLM $\le$ 1.40, MD4 $\le$ 1.37) 의 범위에 근접하고 이전 연속 확산 기준 (Plaid $\le$ 1.48) 보다 향상된 성능을 보였습니다.

5. 중요성 및 주장

본 논문은 DSL 의 주요 이점이 추론 경로의 통합에 있다고 주장합니다. 시간 단계 인덱싱된 노이즈 제거가 아닌 노이즈 토큰에 대한 사후 상태 노이즈 제거로 연속 확산을 매개변수화함으로써 DSL 은 다음을 가능하게 합니다:

• 엔드포인트 노출, 마스킹 정제, 연속/하이브리드 노이즈 제거를 단일 사후 추정기가 처리할 수 있습니다.
• 체크포인트 호환성: 단일 모델을 재학습이나 증류 없이 여러 SNR 경로 (예: 랜덤 순서 AR, 마스킹 정제, 하이브리드) 를 따라 쿼리할 수 있습니다.
• 향상된 분포 충실도: 연속 유한 SNR 상태는 더 풍부한 감독을 제공하여 표준 연속 확산 또는 마스킹 확산 기준에 비해 모든 단계 예산에서 MAUVE 점수가 크게 향상됩니다.

저자들은 언어를 위한 연속 확산에서 결여된 요소가 연속성이 아니라 표현 선택이었다고 결론지었습니다. 단위 구 임베딩과 확률적 국소화를 채택함으로써 그들은 연속 및 이산 확산 간의 격차를 해소하여 단일 모델이 높은 효율성과 품질로 다양한 생성 전략을 지원할 수 있도록 했습니다.

한계: 현재 평가는 OpenWebText 와 Text8 으로 제한되어 있습니다. 보고된 샘플러 (ReMDM 계열, 하이브리드) 는 포괄적으로 최적화된 디코딩 정책은 아니지만, 근본적인 DSL 사후의 유연성을 입증하는 역할을 합니다. 향후 연구 방향으로는 재마스킹 일정 최적화, 하이브리드 핸드오프 규칙, 조건부 생성 탐구 등이 포함됩니다.