DiLaDiff: Distilled Latent-Augmented Diffusion for Language Modeling

기술 요약: DiLaDiff: 언어 모델링을 위한 증류된 잠재 증강 확산

문제 제기
확산 언어 모델 (DLM) 은 샘플링 품질과 처리량 (throughput) 간의 근본적인 트레이드오프에 직면해 있습니다. 이산 확산 모델 (특히 마스킹 확산 언어 모델, MDLM) 은 범주형 데이터에 효과적이지만, 해독된 토큰 간의 상관관계를 본질적으로 포착하지 못합니다. 이러한 한계는 탈노이즈기 (denoiser) 가 사후 확률을 독립적인 토큰별 주변 분포의 곱으로 근사하기 때문에 발생합니다. 그 결과, 여러 토큰을 병렬로 생성하면 불일치하는 문장이 생성되어 순차적 생성에 의존하거나 병렬성을 축소해야 하며, 이는 처리량을 심각하게 제한합니다. 반면, 텍스트 표현 (예: 임베딩) 에 적용된 연속 확산 모델은 원칙적인 샘플링 및 증류 기술을 제공하지만, 표현력 - 해독 가능성 트레이드오프에 종종 어려움을 겪습니다. 즉, 문맥적 표현을 이산 토큰으로 다시 해독하기 어렵거나, 이산 대응 모델의 성능을 따라가지 못하는 경우가 많습니다.

방법론
저자들은 토큰 상관관계 문제를 해결하면서도 높은 처리량을 유지하도록 설계된 하이브리드 연속 - 이산 프레임워크인 DiLaDiff(증류된 잠재 증강 확산 언어 모델) 를 제안합니다. 아키텍처는 세 가지 주요 구성 요소로 이루어져 있습니다:

1. 잠재 증강 오토인코더 (LaDiff):

   • 인코더: 트랜스포머 기반 인코더 ($E_\phi$) 가 문맥적 표현 (BERT 은닉 상태에서 유래) 으로부터 연속 잠재 공간 ($z$) 을 학습합니다. 이 잠재 공간은 전역 토큰 상관관계와 의미 정보를 포착하도록 설계되었습니다.
   • 디코더: 사전 훈련된 MDLM 에서 초기화된 이산 디코더 ($x_\theta$) 는 잠재 변수 $z$에 조건부적으로 텍스트 토큰을 생성합니다.
   • 훈련: 시스템은 오토인코딩과 마스킹 확산을 결합한 수정된 증거 하한 (ELBO) 을 사용하여 훈련됩니다. 핵심적으로, 사후 분포는 잠재 $z$에 조건부적으로 토큰에 대해 분해됩니다. $z$가 상관관계를 포착하므로 조건부 독립 가정이 성립하여 정확한 병렬 토큰 생성이 가능해집니다.
   • 정규화: 생성 모델링을 위해 잠재 공간이 너무 희소해지는 것을 방지하기 위해, 저자들은 좌표별 정규화, 특징 및 잠재의 무작위 마스킹, 가우시안 노이즈 주입과 같은 견고성 강제 휴리스틱을 사용합니다.

2. 잠재 사전 학습:

   • 추론 시 깨끗한 시퀀스를 사용할 수 없으므로, 인코더 출력에 대한 사전 분포 $p(z)$를 학습하기 위해 연속 확산 모델 (잠재 탈노이즈기 $z_\psi$) 이 훈련됩니다. 이를 통해 모델은 디코딩 전에 역 확률 흐름 ODE 를 통해 고품질 잠재 $z$를 샘플링할 수 있습니다.

3. 일관성 증류 (DiLaDiff):

   • 연속 확산 단계의 계산 오버헤드를 추가로 줄이기 위해, 저자들은 MeanFlow 증류를 적용합니다. 학생 모델 ($u_\eta$) 은 순간 속도가 아닌 두 시간 단계 ($t$와 $r$) 사이의 ODE 궤적의 평균 속도를 예측하도록 훈련됩니다.
   • 이 증류를 통해 잠재 $z$는 몇 단계 (예: 5 단계) 만으로도 생성될 수 있으며, 많은 단계 (예: 200 단계) 가 필요한 교사 모델 (LaDiff) 에 가까운 성능을 유지합니다.

주요 기여

• 하이브리드 아키텍처: 이 논문은 디코더가 사전 훈련된 이산 확산 베이스라인에서 초기화된 텍스트 오토인코더 훈련 방법을 제시합니다. 이는 의미 정보와 토큰 상관관계를 포착하는 제어 가능하고 정규화된 잠재 공간을 결과물로 낳습니다.
• 잠재 유도 병렬성: 연속 확산을 통해 잠재 사전을 학습함으로써, 모델 (LaDiff) 은 텍스트 품질을 희생하지 않고 여러 토큰을 병렬로 샘플링할 수 있게 합니다. 이는 이산 모델이 놓치는 토큰 종속성을 잠재 공간이 해결하기 때문입니다.
• 텍스트 잠재에 대한 최초의 성공적인 ODE 증류: 저자들은 문맥적 텍스트 잠재에 대한 ODE 궤적을 증류한 최초의 성공적인 시도인 DiLaDiff를 제시합니다. 이는 이산 디코딩에 비해 잠재 생성 비용을 무시할 수 있는 수준으로 줄여줍니다.

결과
실험은 시퀀스 길이 1024 토큰의 OpenWebText 데이터셋에서 수행되었습니다.

• 베이스라인 대비 성능: 증류 없이도 LaDiff는 마스킹 확산 베이스라인 (MDLM) 을 능가합니다. 무조건부 생성에서 LaDiff 는 1024 이산 단계를 가진 MDLM 대비 생성 퍼플렉시티 (GenPPL) 에서 28.5 의 절대적 개선과 MAUVE 에서 0.08 의 절대적 개선을 달성합니다.
• 처리량: LaDiff 는 최적의 MDLM 설정 대비 배치 크기 32 에서 추론 속도를 7 배 가속화합니다.
• 증류 효율성: DiLaDiff는 교사 모델 (LaDiff) 에 비해 5 개의 잠재 확산 단계만 사용하여 (교사 모델은 200 단계) 교사 모델에 가까운 성능을 달성합니다. 잠재 확산 구성 요소의 계산 오버헤드는 전체 월타임의 약 **5%**로 감소하여, 잠재 생성 비용이 이산 디코딩 비용에 비해 무시할 수 있게 됩니다.
• 소수 단계 생성: DiLaDiff 는 이산 디코더를 명시적으로 증류하지 않고도 최첨단 소수 단계 방법 (MDLM+SDTT 및 DUO+DCD 등) 과 경쟁하며, 잠재 증강의 효과를 입증합니다.
• 의미적 다양성: 분석 결과, 잠재 공간은 노이즈가 있는 잠재에 대한 BERTScore 를 통해 검증된 대로 의미 정보를 보존하며, 디코더가 낮은 온도에서도 엔트로피를 유지할 수 있다는 사실로 입증된 대로 다양한 생성을 가능하게 합니다.

의의 및 주장
이 논문은 DiLaDiff 가 연속 표현의 표현력과 이산 디코딩의 견고성 사이의 간극을 성공적으로 메웠다고 주장합니다. 잠재 공간에서 전역 상관관계를 포착함으로써, 모델은 병렬 토큰 생성과 관련된 이산 확산 모델의 본질적 한계를 극복합니다. 저자들은 이 접근 방식이 추론 지연을 크게 줄이면서도 고품질 텍스트 생성을 가능하게 한다고 강조합니다.

이 논문은 한계에 대해 겸손하게 서술합니다:

• DiLaDiff 는 아직 소수 단계 영역에서 교사 모델 (LaDiff) 의 성능을 완전히 따라잡지 못하여, 증류에 대한 추가 작업이 필요함을 시사합니다.
• 이산 디코더 자체는 증류되지 않았습니다. 저자들은 잠재 증류와 이산 디코더 증류를 결합하는 것이 진정한 최첨단 소수 단계 성능을 달성하기 위한 자연스러운 미래 확장이라고 지적합니다.
• 정규화 전략과 확산 스케줄 설계는 현재 휴리스틱이지만, 저자들은 부록에서 더 원칙적인 시도를 제시합니다.

요약하자면, DiLaDiff 는 잠재 증류로 강화된 하이브리드 연속 - 이산 접근 방식이 생성 품질을 유지하거나 향상시키면서 언어 모델링 추론을 크게 가속화할 수 있음을 보여주며, 현재 확산 언어 모델에 내재된 처리량 - 품질 트레이드오프를 효과적으로 해결합니다.