Consistent Diffusion Language Models

"일관된 확산 언어 모델 (Consistent Diffusion Language Models, CDLM)"논문에 대한 상세한 기술적 요약입니다.

1. 문제 제기

확산 언어 모델 (DLM) 은 병렬 토큰 생성과 하선형 시간 추론을 가능하게 함으로써 자기회귀 (AR) 모델에 대한 유망한 대안을 제시합니다. 그러나 중요한 병목 현상이 남아 있습니다: 고품질 텍스트 생성은 일반적으로 수백 단계의 정제 과정을 필요로 하여, 병렬화가 약속한 효율성 향상을 무효화합니다.

연속 영역 (예: 이미지) 에서 **일관성 모델 (Consistency Models)**은 **확률 흐름 상미분 방정식 (PF-ODE)**으로 정의된 고유한 결정론적 궤적을 따라 모델이 불변하도록 학습시킴으로써 확산을 성공적으로 가속화했습니다.

• 격차: 이산 영역 (텍스트 등) 에서는 유사한 샘플 공간 PF-ODE 가 존재하지 않습니다. 이산적 오염 과정 (예: 마스킹 또는 균일 노이즈) 은 단일 결정론적 경로를 따르지 않으며, 대신 확률적 전이를 수반합니다.
• 결과: 연속 일관성 방법을 이산 데이터에 직접 적용하는 것은 잘 정의되지 않습니다. 기존 해결책들은 종종 (교사 모델을 사용하는) 다단계 증류 파이프라인이나 연속 완화 대리 모델을 의존하는데, 이는 복잡하고 계산 비용이 많이 들며 종종 모드 붕괴나 다양성 손실을 겪습니다.

2. 방법론: 다중 경로 이산 일관성 (MPDC)

저자들은 근본적인 전환을 제안합니다: 존재하지 않는 결정론적 경로를 찾는 대신, 이산 확산 과정에 내재된 **정확한 사후 브리지 (exact posterior bridges)**를 활용합니다.

핵심 개념: 사후 브리지

이산 확산 과정에서 임의의 두 노이즈 수준 $s < t$에 대해, 닫힌 형태의 정확한 사후 브리지 $q(x_s | x_t, x_0)$가 존재합니다. 이 브리지는 실제 데이터 $x_0$에 조건부인 노이즈가 있는 상태 $x_t$를 더 깨끗한 상태 $x_s$로 연결하는 유효한 확률적 경로를 정의합니다.

• 연속 PF-ODE 와 달리, 이 브리지는 확률적입니다.
• 결정적으로, 이러한 브리지는 **반군 (semigroup)**을 형성 (채프먼 - 콜모고로프 성질) 하여 경로가 전이적으로 결합될 수 있음을 의미합니다.

MPDC 원칙

저자들은 **다중 경로 이산 일관성 (Multi-Path Discrete Consistency, MPDC)**을 도입하여, 노이즈 제거기를 기대값 내에서 경로 불변하도록 학습시킵니다.

• 정의: 모델 $f_\theta(x_t, t)$가 일관성이 있다는 것은, 노이즈가 있는 상태 $x_t$에서의 예측이 정확한 사후 브리지를 통해 더 깨끗한 상태 $x_s$로 "이동 (hop)"한 후 이루어진 예측과 (기대값에서) 일치함을 의미합니다.
• 목적: 모델은 $x_t$로부터의 직접 예측과 $x_s$ (브리지에서 샘플링됨) 로부터의 기대 예측 간의 발산을 최소화합니다.
• 수학적 공식화:
  $$\mathcal{L}_{CDLM} = \mathbb{E}_{t, \delta, x_t} \left[ \sum_{i \in \mathcal{M}(x_t)} w(t, \delta) \cdot D\left( f_\theta(x_t, t)_i \parallel \tilde{f}_\theta(x_s, s)_i \right) \right]$$
  여기서 $x_s \sim q(x_s | x_t, x_0)$이며, $\tilde{f}_\theta$는 타겟 네트워크 (EMA) 이고, $D$는 발산 지표 (일반적으로 제센 - 샤논 발산) 입니다.

안정성을 위한 주요 설계 선택

모드 붕괴 (모델이 일관성을 무난히 만족시키기 위해 반복적인 텍스트를 출력하는 현상) 와 균일한 드리프트를 방지하기 위해, 저자들은 세 가지 안정화 메커니즘을 도입합니다:

1. 커리큘럼으로서의 단계 크기: 특정 범위 (예: $1/8$에서 $3/8$) 에서 단계 크기 $\delta$를 샘플링하고 $1/\delta$로 가중치를 부여하여 짧고 긴 경로 전반에 걸쳐 학습 신호를 균형 있게 조정합니다.
2. 최대-단계 확산 앵커: 표준 확산 목적을 사용하여 훈련 손실의 작은 비율을 실제 데이터 ($s=0$) 에 앵커링합니다. 이는 모델이 퇴화된 상수 해로 수렴하는 것을 방지합니다.
3. 최적화 비대칭 및 발산: 타겟 네트워크에 스톱-그래디언트를 적용하고 Forward/Backward KL 대신 **제센 - 샤논 발산 (JSD)**을 사용합니다. JSD 는 모드 탐색과 모드 커버링 사이의 안정적인 균형을 제공하여 처음부터 훈련할 때 KL 발산의 불안정성을 피합니다.

3. 주요 기여

1. 이산 생성 모델링을 위한 새로운 원칙: 이 논문은 정확한 사후 브리지가 결여된 PF-ODE 를 위한 엄밀하고 닫힌 형태의 대안으로 작용하여, "이산 가속화"장애를 분석적 자산으로 전환한다고 규명합니다.
2. 통합된 훈련 프레임워크: CDLM 은 단일 단계, 교사 없는 목적을 제공합니다. 저자들은 표준 마스킹 확산, 연속 일관성 모델, 점진적 증류, 그리고 최근의 2 단계 증류 방법 (SDTT 및 DUO-DCD 등) 이 모두 CDLM 프레임워크의 특수한 경우, 극한, 또는 근사 결합임을 증명합니다.
3. 최고 수준의 성능: CDLM 은 무조건부 및 조건부 텍스트 생성 모두에서 우수한 결과를 달성하여 강력한 기본 모델을 능가하고, 교사를 필요로 하지 않는 복잡한 다단계 증류 모델과 맞먹거나 능가합니다.

4. 실험 결과

저자들은 OpenWebText 데이터셋에서 MDLM, DUO, SDTT, DUO-DCD 와 같은 베이스라인에 대해 CDLM (마스킹 확산용 MCDLM과 균일 확산용 UCDLM) 을 평가했습니다.

• 무조건부 생성:

  • 소수 단계 효율성: CDLM 은 단일 단계 모델에 대해 새로운 최고 수준 (SOTA) 결과를 확립했습니다. 모든 샘플링 예산에 걸쳐 기본 모델 (MDLM, DUO) 을 크게 능가합니다.
  • 증류 모델과의 비교: 주목할 만하게도, 단일 단계 CDLM 은 훨씬 더 높은 **엔트로피 (다양성)**를 유지하면서 퍼플렉시티 (PPL) 측면에서 **다단계 증류 모델 (SDTT, DUO-DCD)**을 종종 능가합니다.
  • 모드 붕괴: 탐욕적 샘플러를 사용하는 증류 모델 (예: DUO-DCD) 은 낮은 엔트로피 (3.9) 를 보여 심각한 모드 붕괴 (반복적인 문자) 를 나타냈습니다. 반면 CDLM 은 건강한 엔트로피 (5.0 이상) 를 유지하여 다양성을 보존하는 능력을 입증했습니다.
  • 속도: CDLM 은 자기회귀 베이스라인 대비 32 배~64 배, 표준 MDLM 대비 64 배~128 배의 속도 향상을 달성하면서도 비교 가능한 품질을 유지합니다.

• 조건부 생성:

  • CDLM 은 여러 데이터셋 (OpenWebText, Lambada, Wikitext-103, PTB) 에서 50% 오염된 입력으로부터 원래 문장을 복원하는 데 있어 견고한 성능을 보여주었습니다.
  • 다른 모델들이 입력 조건을 유지하는 데 실패한 경우에도, CDLM 은 SDTT 및 DUO 보다 더 나은 BLEU 점수와 더 낮은 퍼플렉시티를 달성했습니다.

5. 중요성 및 영향

• 이론적 통합: 이 논문은 이산 생성 모델링을 위한 "통합 렌즈"를 제공하여, 다양한 가속화 기법이 경로 일관성 개념을 통해 근본적으로 연결되어 있음을 보여줍니다.
• 확장성: 다단계 증류와 교사 모델의 필요성을 제거함으로써, CDLM 은 차세대 빠른 이산 생성 모델을 위한 원칙적이고 확장 가능하며 교사 없는 기반을 제공합니다.
• 일반화 가능성: 텍스트에서 입증되었지만, 이 프레임워크는 처리 가능한 사후 브리지를 가진 모든 이산 영역 (예: 생물학적 시퀀스, 그래프, 프로그램 합성) 에 적용 가능합니다.
• 패러다임 전환: CDLM 은 빠른 이산 확산을 연속 ODE 의 근사가 아닌, 확률적 브리지를 통한 경로 독립적 노이즈 제거기의 학습으로 재정의하여, 이산 공간에서의 소수 단계 생성 문제를 성공적으로 해결합니다.

결론적으로, CDLM은 연속 일관성 모델의 효율성과 이산 데이터의 현실 사이의 간극을 메우며, 증류 파이프라인의 복잡성 없이 고품질의 빠른 텍스트 생성을 달성하는 견고한 단일 단계 솔루션을 제공합니다.