기술 요약: 잠재 공간과 함께 잠재 확산 언어 모델 학습 방법

문제 정의

현재 자동회귀 모델이 텍스트 생성을 지배하고 있지만, 왼쪽에서 오른쪽으로 순차적으로 생성되는 특성으로 인해 병렬 생성과 이전 오류 수정이 제한됩니다. 확산 언어 모델은 병렬로 반복적 정제를 통해 텍스트를 생성함으로써 대안을 제시합니다. 그러나 기존 접근법은 상당한 장애물에 직면해 있습니다:

• 이산 확산 (Discrete Diffusion): 토큰 공간에서 직접 작동하지만, 분산된 결합 확률 분포로 인해 각 단계에서 토큰을 독립적으로 예측하는 문제가 발생합니다. 이는 여러 토큰을 동시에 생성하는 능력을 제한하여 고품질 단단 단계 생성을 어렵게 만듭니다.
• 연속 확산 (Continuous Diffusion): 잠재 공간에서 작동하여 모든 위치의 점진적 정제를 가능하게 합니다. 이산 방법보다 우수하지만, 현재 연속 접근법은 잠재 공간을 구성하기 위해 동결된 사전 훈련 인코더(예: Transformer 기반 인코더)에 의존합니다. 이러한 인코더는 다음 토큰 예측과 같은 대리 작업으로 훈련되며, 확산 탈노이즈 과정의 특정 요구 사항에 최적화되지 않았습니다. 반면, 얕은 토큰 임베딩 (원-핫 또는 고정 벡터) 은 확산에 맞게 조정되더라도 최적이지 않은 잠재 공간을 생성합니다.

핵심적인 과제는 인코더 훈련과 확산 훈련을 분리함으로써 발생하는 불일치 없이, 동시에 복호화 가능(토큰으로 매핑 가능) 하고 확산 친화적(탈노이즈가 용이) 인 잠재 공간을 구축하는 것입니다.

방법론: LDLM

저자들은 잠재 인코더, 확산 모델, 디코더가 함께 훈련되는 잠재 확산 언어 모델 (LDLM) 프레임워크를 제안합니다. 이를 통해 확산 목적 함수가 잠재 공간을 직접 형성할 수 있습니다.

아키텍처

1. 2 단계 인코딩:
   • 동결된 토큰 인코더: 입력 토큰 $w$는 먼저 동결된 사전 훈련 토큰 인코더 (GPT-2) 를 통해 처리되어 컨텍스트 숨겨진 상태 $h$를 추출합니다.
   • 학습 가능한 잠재 인코더: 학습 가능한 인코더 $E_\theta^z$는 이러한 숨겨진 상태 $h$를 확산 잠재 공간 $z_0$으로 매핑합니다. 이는 확산에 적합하도록 표현을 재구성하면서도 복호화 가능성을 유지합니다.
2. 2 단계 디코딩:
   • 잠재 디코더: 학습 가능한 디코더 $D_\theta^h$는 훈련 중 노이즈로 교란된 잠재 $z_0$를 숨겨진 상태 $\hat{h}$로 다시 매핑합니다.
   • 토큰 디코더: 토큰 디코더 $D_\theta^w$는 $\hat{h}$를 토큰 $w$로 변환합니다.
3. 훈련 목적 함수:
   모델은 결합된 손실 함수 $L(\theta) = L_{diff} + L_h + L_w$를 최소화합니다.
   • $L_{diff}$: 노이즈가 있는 관측치에서 깨끗한 잠재 상태를 복원하기 위한 표준 연속 확산 손실 (MSE).
   • $L_h$(숨겨진 상태 MSE): 잠재 공간에서 동결된 인코더의 숨겨진 상태 $h$를 재구성합니다. 핵심적으로, 이는 교차 엔트로피 (CE) 대신 **평균 제곱 오차 (MSE)**를 사용합니다.
   • $L_w$(토큰 CE): 숨겨진 상태로부터 토큰을 재구성합니다. 기울기는 재구성된 숨겨진 상태에서 차단되어 잠재 인코더에 영향을 미치지 않도록 합니다.

핵심 훈련 레시피 ("4 가지 요소" 솔루션)

저자들은 단순한 결합 훈련이 실패한다 (낮은 다양성의 생성물 생성) 고 지적하며, 이를 작동시키기 위한 구체적인 레시피를 제안합니다:

1. MSE 디코더 손실: 이전 연구에서 교차 엔트로피를 사용한 것과 달리, LDLM 은 숨겨진 상태를 재구성하기 위해 MSE 를 사용합니다. CE 는 토큰을 구별하기 위해 잠재 상태를 잘 분리되도록 강제하여 확산에 필요한 매끄러운 기하학과 충돌합니다. MSE 는 덜 엄격하여 기본 숨겨진 상태가 유사하다면 잠재 상태가 서로 가까워지도록 허용하여 잠재 붕괴를 방지하면서도 매끄러운 다양체를 유지합니다.
2. 확산 - 인코더 워밍업: 훈련 초기에는 재구성 손실 (큰 잠재 크기를 장려) 과 확산 손실 (탈노이즈를 단순화하기 위해 작은 크기를 장려) 이 인코더를 반대 방향으로 당깁니다. 저자들은 확산 기울기를 시작 시 0 에 가깝다가 1 로 증가하는 계수 $\gamma(s)$로 스케일링하는 워밍업을 구현합니다. 이를 통해 인코더가 확산 목적 함수에 의해 형성되기 전에 먼저 복호화 가능한 잠재 공간을 학습할 수 있습니다.
3. 적응적 타임스텝 샘플링: 잠재 공간이 훈련 중에 진화하므로 고정된 노이즈 스케줄은 최적이지 않습니다. 저자들은 탈노이즈 손실이 타임스텝에 따라 선형적으로 증가하도록 타임스텝 샘플링 분포를 동적으로 조정하여 모든 타임스텝이 불확실성 감소에 동등하게 기여하도록 합니다.
4. 디코더 입력 노이즈: 훈련 중 잠재 디코더 입력에 가우시안 노이즈 ($\sigma_{dec}\epsilon$) 를 주입합니다. 이는 세 가지 목적을 수행합니다:
   • 인코더가 정보를 효율적으로 저장하도록 장려하여 (사용되지 않는 차원이 신호를 손상시키는 것 방지).
   • 데이터 증강 역할을 하여 디코더가 확산 오류에 강건하게 만듭니다.
   • 잠재 좌표의 분산을 안정화하여 확산 파라미터화를 개선합니다.

결과

저자들은 LDLM 을 OpenWebText (OWT) 및 LM1B에서 평가하여 최첨단 이산 (MDLM, Duo) 및 연속 (CANDI, FLM) 확산 베이스라인과 비교했습니다.

• 품질 - 다양성 트레이드오프: LDLM 은 생성 퍼플렉시티 (Gen. PPL) 와 엔트로피 사이의 가장 좋은 파레토 프론티어를 달성합니다. 동등한 엔트로피 수준에서 베이스라인보다 Gen. PPL 에서 현저히 우수하며, 실제 텍스트의 품질에 근접합니다.
• 잠재 공간 비교: 제거 실험 (Ablation studies) 은 LDLM 의 함께 학습된 잠재 상태가 다음 것들보다 우수함을 보여줍니다:
  • 처음부터 훈련되거나 동결된 얕은 임베딩.
  • 동결된 인코더의 컨텍스트 임베딩 (TEncDM 스타일).
  • 명시적 평활화 목적 함수로 훈련된 오토인코더의 잠재 상태 (COSMOS 스타일).
    이는 잠재 공간이 확산 목적 함수에 맞게 특별히 조정되어야 함을 확인시켜 줍니다.
• 계산 효율성: LDLM 은 샘플링 시 경쟁 방법보다 2~13 배 빠릅니다. 매 단계에서 어휘로 투영하는 토큰 수준 확산 모델과 달리, LDLM 은 생성 과정 내내 저차원 잠재 공간에서 작동하며 마지막에 한 번만 토큰으로 투영합니다.

중요성 및 주장

이 논문은 잠재 공간을 함께 학습하는 것이 텍스트 생성에 있어 잠재 확산을 경쟁력 있게 만들기 위한 중요하고 이전에 다루지 않았던 단계라고 주장합니다.

• 단순한 결합 훈련 극복: 이 연구는 단순히 인코더와 확산 모델을 함께 훈련하는 것만으로는 부족하며, 훈련을 안정화하고 고품질을 달성하기 위해 MSE 손실, 워밍업, 적응적 샘플링, 디코더 노이즈와 같은 특정 기법이 필요함을 보여줍니다.
• 컨텍스트 잠재의 우월성: 결과는 고수준 시퀀스 정보 (동결된 GPT-2 인코더가 포착) 가 필수적이지만, 확산을 위해 학습 가능한 인코더에 의해 재구성되어야 함을 나타냅니다.
• 효율성: 매 단계별 어휘 투영을 피함으로써 LDLM 은 상당한 속도 향상을 제공하여 확산 기반 텍스트 생성의 주요 병목 현상을 해결합니다.

저자들은 탈노이즈 목적 함수로 잠재 공간을 직접 형성하는 것이 효과적인 잠재 텍스트 확산 모델을 구축하는 데 필수적이며, 최첨단 생성 품질을 달성하면서도 우수한 추론 속도를 유지한다고 결론지었습니다.