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

문제 정의

연속 확산은 이미지, 비디오, 오디오의 고차원 생성을 위한 표준 프레임워크가 되었으나, 언어 생성 분야에서는 마스크된 이산 확산 모델 (MDM) 에 비해 일관되게 뒤처져 왔습니다. 기존 연속 확산 언어 모델은 일반적으로 분산 보존 (VP) 또는 분산 폭발 (VE) 가산 노이즈를 사용하여 노이즈가 섞인 토큰 임베딩 위에서 작동합니다. 저자들은 병목 현상이 연속성 그 자체가 아니라 표현의 선택에 있다고 주장합니다.

표준 연속 확산에서, 노이즈 제거 과정은 시간 단계 인덱스가 지정된 노이즈 체제에 의존합니다. 모델은 노이즈가 섞인 임베딩이 포함하고 있는 증거의 양을 보정하기 위해 시간 단계 또는 노이즈 수준 레이블을 받아야 합니다. 결과적으로, 이 표현 방식의 노이즈 제거기는 마스크된 (정보 없음) 상태와 불확실한 (유한 신호대잡음비) 상태를 구별하기 위해 명시적인 시간 조건화에 매우 민감합니다. 이러한 명시적 시간 조건화에 대한 의존성은 별도의 샘플러 전용 모델이나 재학습 없이 단일 네트워크가 무작위 순서 자기회귀 샘플링이나 하이브리드 연속 - 이산 디코딩과 같은 다양한 추론 경로를 자연스럽게 지원하지 못하게 합니다.

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

본 논문은 명목 신호대잡음비 (SNR) 에 대해 베이지안 최적 노이즈 제거기를 불변으로 만드는 연속 상태 프레임워크인 **이산 확률적 국소화 (DSL)**를 소개합니다.

1. 단위 구 토큰 기하학

핵심 혁신은 모든 깨끗한 토큰 임베딩을 단위 구 ($\|x_i\|_2 = 1$) 위에 위치하도록 제약하는 것입니다. 노이즈가 섞인 관측치가 $z_t = t x + \sqrt{t}\epsilon$인 확률적 국소화 채널 하에서, 베이지안 사후 분해는 $t \|x\|_2^2$에 비례하는 항을 포함합니다. 임베딩이 단위 노름을 가지므로, $\|x\|_2^2$는 주어진 길이의 문장에 대해 일정합니다. 이 항은 정규화 과정에서 상쇄되어, 사후 확률 및 따라서 최적 노이즈 제거기가 명목 SNR $t$에 독립이 됩니다.

최적 노이즈 제거기는 노이즈가 섞인 상태 $z$만의 함수가 됩니다:
$$\hat{x}(z) = \mathbb{E}[x | z] = \frac{\mathbb{E}_P[x e^{x \cdot z}]}{\mathbb{E}_P[e^{x \cdot z}]}$$
이를 통해 단일 시간 무관 네트워크가 토큰별 SNR 경로 전체를 지원할 수 있게 됩니다.

2. 토큰별 SNR 경로

노이즈 제거기가 전역 시간 단계를 필요로 하지 않기 때문에, DSL 은 각 토큰 위치가 자체 임의의 SNR 경로 $\gamma_i(t)$ 하에서 진화할 수 있게 합니다. 이는 하나의 모델 하에서 여러 추론 전략을 통합합니다:

• 무작위 순서 자기회귀 (ROAR): 토큰을 하나씩 드러내는 것은 다른 토큰들을 0 에 고정시킨 채 특정 $\gamma_i$를 0 에서 $\infty$로 구동하는 것에 해당합니다.
• 마스크 정제: 토큰을 다시 마스크하는 것은 해당 $\gamma_i$를 0 으로 재설정하는 것에 해당합니다.
• 연속 확산: 모든 $\gamma_i$를 동시에 증가시키는 것은 표준 연속 노이즈 제거에 해당합니다.

3. 학습 목표 및 아키텍처

이러한 다양한 경로를 위한 단일 모델을 학습시키기 위해, DSL 은 혼합 지원 SNR 샘플링 전략을 사용합니다:

• 연속 분기: 유한 SNR 상태를 커버하기 위해 로그 정규 분포에서 SNR 을 샘플링합니다.
• 엔드포인트 분기: 마스크된 ($\gamma=0$) 과 깨끗한 ($\gamma=\infty$) 극단 근처에 집중된 ROAR 스타일 분포에서 샘플링합니다.

두 분기 모두 단일 토큰 수준 교차 엔트로피 (CE) 손실 함수를 사용하여 학습됩니다. 저자들은 범주형 토큰 사후 확률에 대한 CE 를 최소화하는 것이 정확한 경로 적분 평균 제곱 오차 (MSE) NLL 의 대리 변수이며, 직접 임베딩 공간 회귀 (임베딩 붕괴 또는 암기 등) 와 관련된 병리 현상을 피한다고 보여줌으로써 이를 정당화합니다.

아키텍처에는 노이즈가 섞인 임베딩 $z_i$를 바운드된 토큰 혼합 표현 $m_i$로 매핑하는 사후 뷰 변환기가 포함되어 있습니다. 이는 SNR 수준에 관계없이 Transformer 백본의 입력이 어휘에 대한 확률 분포로 안정적이고 해석 가능하도록 보장합니다.

주요 기여

1. 이산 생성을 위한 연속 상태 프레임워크: DSL 은 단위 구 임베딩 위의 확률적 국소화 채널을 도입하여, 베이지안 최적 노이즈 제거기가 노이즈가 섞인 상태만의 함수가 되도록 함으로써 명시적 시간 조건화의 필요성을 제거합니다.
2. 실용적 레시피: 논문은 엔드포인트 (마스크/드러내기) 와 유한 SNR 상태 모두에서 단일 노이즈 제거기를 학습하는 혼합 지원 SNR 샘플링 전략을 제안하며, 이는 연속 사후 기하학을 표준 Transformer 백본에 연결하는 변환기와 결합됩니다.
3. 경로에 걸친 증거: 저자들은 단일 DSL 체크포인트가 재학습 없이 마스크 정제, 무작위 순서 자기회귀 샘플링, 그리고 하이브리드 연속 - 이산 샘플링을 모두 지원할 수 있음을 입증합니다.

실험 결과

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

• OpenWebText 생성: DSL 은 기준선 MDLM 에 비해 모든 단계 예산 ($T=128$에서 $T=1024$) 에서 분포적 충실도 (MAUVE 로 측정) 를 크게 향상시킵니다.
  • 마스크 정제 디코더 (ReMDM) 를 사용하면, DSL 은 $T=1024$에서 0.722의 MAUVE 를 달성하여 PRISM-loop (0.527) 및 ADLM (0.699) 과 같은 이전 최첨단 방법들을 능가합니다.
  • 반복 정제 없이도 (간단한 ROAR 또는 일반 MDLM 디코더 사용), DSL 미세 조정 체크포인트는 기준선을 크게 능가하며, 이는 개선이 샘플링 스케줄 자체보다는 학습된 사후 확률에서 비롯됨을 시사합니다.
• 하이브리드 샘플링: 짧은 연속 노이즈 제거 단계를 수행한 후 이산 마스크 정제를 수행하는 하이브리드 샘플러는 총 48 단계만에 일관된 생성을 달성하여 (MAUVE 0.662), 훨씬 더 많은 단계 수의 비정제 기준선들을 능가합니다.
• Text8 가능성: Text8 벤치마크에서 DSL 은 문자당 비트 (BPC) $\le 1.45$를 달성합니다. 이는 상한선 (ELBO) 이며 자기회귀 모델의 정확한 NLL 과 직접 비교할 수는 없으나, 이산 확산 모델 (MDLM $\le 1.40$) 의 범위와 근접하고 다른 연속 확산 기준선 (Plaid $\le 1.48$) 보다 개선된 결과를 보입니다.

중요성 및 주장

본 논문은 언어 분야에서 연속 확산과 이산 확산 간의 격차가 연속성의 본질 때문이 아니라 노이즈 체제의 표현 때문이라고 주장합니다. 노이즈가 섞인 상태만이 깨끗한 토큰 사후 확률을 결정하도록 오염을 매개변수화함으로써, DSL 은 마스크 확산의 "엔드포인트 드러내기" 행동과 확산 모델의 "연속 정제" 능력을 통합합니다.

핵심 주장은 엔드포인트 드러내기, 마스크 정제, 그리고 연속/하이브리드 노이즈 제거가 별도의 샘플러 전용 모델이 아닌 하나의 사후 추정기에 의해 지원될 수 있다는 것입니다. 이는 모델이 시간 무관 사후 매핑을 학습하고, 추론 경로 (AR, 마스크, 또는 하이브리드) 가 샘플링 시점에 동적으로 결정되는 비자기회귀 생성을 위한 새로운 패러다임을 시사합니다.

저자들은 한계에 대해 겸손하게 언급하며, 평가가 OWT 와 Text8 로 제한되어 있으며, 보고된 샘플러 (ReMDM 계열 및 하이브리드) 는 포괄적으로 최적화된 디코딩 정책이 아니라 DSL 사후의 유연성을 입증하는 개념 증명임을 지적합니다. 향후 연구 방향으로는 최적화된 리마스크 스케줄, 더 큰 모델, 그리고 조건부 생성 등이 포함됩니다.