Latent Stochastic Interpolants

1. 연구 배경 및 문제점 (Problem)

• 확률적 보간자 (SI) 의 한계: SI 는 두 개의 임의의 확률 분포를 유연하게 연결하는 강력한 생성 모델링 프레임워크입니다. 그러나 기존 SI 는 관측 공간 (Observation Space) 에서 직접 작동하며, 사전 분포 ($p_0$) 와 목표 분포 ($p_1$) 의 샘플에 대한 직접적인 접근을 전제로 합니다.
• 잠재 변수 모델과의 통합 부재: 잠재 변수 모델 (인코더, 디코더, 생성 모델) 을 공동으로 학습 (Joint Learning) 하려면, 인코더가 추론한 잠재 공간에서의 분포와 생성 모델이 학습해야 할 분포가 일치해야 합니다. 하지만 SI 는 고정된 분포를 가정하므로, 인코더와 생성 모델이 동시에 최적화되는 과정에서 잠재 공간의 분포가 변하는 상황을 처리할 수 없어 적용이 불가능했습니다.
• 계산 비용: 고차원인 관측 공간 (예: 고해상도 이미지) 에 SI 를 직접 적용하면 계산 비용이 매우 큽니다. 반면, 잠재 공간은 차원이 낮아 효율적이지만, 기존 방법론은 이를 활용하지 못했습니다.
• 단순한 사전 분포의 제약: 기존 확산 모델 (Diffusion Models) 은 주로 단순한 가우시안 사전 분포를 사용하는데, 이는 표현력의 한계를 가집니다.

2. 제안 방법론 (Methodology)

저자들은 잠재 확률적 보간자 (LSI) 를 제안하여 위의 문제들을 해결했습니다.

• 연속 시간 기반의 증거 하한 (ELBO) 유도:
  • LSI 는 잠재 공간에서 인코더, 디코더, 잠재 SI 모델을 엔드 투 엔드 (End-to-End) 로 공동 학습합니다.
  • 이를 위해 연속 시간 (Continuous Time) 동역학을 기반으로 한 새로운 증거 하한 (ELBO, Evidence Lower Bound) 목적 함수를 유도했습니다. 이 목적 함수는 관측 데이터의 로그 가능도 (Log-likelihood) 를 제어하면서도, 잠재 공간에서의 변분 근사 (Variational Approximation) 를 가능하게 합니다.
• 시뮬레이션 없는 훈련 (Simulation-free Training):
  • 일반적인 잠재 변수 모델에서는 SDE(확률 미분 방정식) 의 해를 구하기 위해 매 단계마다 시뮬레이션이 필요해 계산 비용이 높습니다.
  • LSI 는 확산 브리지 (Diffusion Bridge) 와 도브의 h-변환 (Doob's h-transform) 을 활용하여, SDE 시뮬레이션 없이도 잠재 변수 $z_t$ 를 직접 샘플링할 수 있는 폐쇄형 해 (Closed-form solution) 를 유도했습니다. 이는 관측 공간 확산 모델의 효율성을 잠재 공간으로 확장한 것입니다.
• 잠재 공간 보간자 구성:
  • 잠재 변수 $z_t$ 는 사전 분포 $z_0$ 와 인코더가 생성한 잠재 표현 $z_1$ 을 연결하는 보간자로 정의됩니다.
  • $z_t = \eta_t \epsilon + \kappa_t z_1 + \nu_t z_0$ 형태로 파라미터화되며, 여기서 $\epsilon$ 은 노이즈입니다. 이를 통해 인코더와 디코더가 학습되는 동안에도 생성 과정이 유연하게 조정됩니다.
• InterpFlow 파라미터화:
  • 훈련 안정성과 성능을 높이기 위해 제안된 InterpFlow라는 새로운 파라미터화 방식을 사용합니다. 이는 기존 방법들 (OrigFlow, NoisePred 등) 보다 낮은 분산의 그래디언트를 제공하여 더 빠른 수렴과 더 좋은 결과를 보입니다.

3. 주요 기여 (Key Contributions)

1. LSI 프레임워크: 연속 시간 동역학을 가진 잠재 변수 생성 모델을 인코더, 디코더, 생성 모델과 함께 공동 학습할 수 있는 새로운 프레임워크를 제시했습니다.
2. 통합적 관점: SI 의 유연한 연속 시간 공식을 잠재 변수 모델에 통합하는 새로운 관점을 제시했습니다. 이는 유동 기반 (Flow-based) 방법과 확산 기반 (Diffusion-based) 방법의 장점을 결합합니다.
3. 원칙적인 ELBO 목적 함수: SI 의 장점 (단순한 시뮬레이션 없는 훈련, 유연한 사전 분포 선택) 을 유지하면서, 잠재 공간에서의 공동 학습 이점을 제공하는 새로운 ELBO 목적 함수를 제안했습니다.

4. 실험 결과 (Results)

연구진은 ImageNet 데이터셋을 사용하여 LSI 의 성능을 광범위하게 평가했습니다.

• 성능 (FID):
  • LSI 는 관측 공간에서 직접 훈련된 동급 크기의 SI 모델과 비슷하거나 더 좋은 FID(Fréchet Inception Distance) 점수를 기록했습니다.
  • 예를 들어, 128x128 해상도에서 LSI 는 3.12 의 FID 를 달성했으며, 이는 관측 공간 모델 (3.46) 보다 우수했습니다.
• 계산 효율성:
  • LSI 는 인코더와 디코더를 포함하지만, 샘플링 단계에서는 인코더를 사용하지 않고 디코더를 한 번만 사용합니다. 생성 모델 (잠재 SI) 만 반복적으로 실행되므로, 샘플링 시 FLOPs(연산량) 가 크게 감소합니다.
  • 128x128 이미지 샘플링 시 약 29.7%, 256x256 이미지에서는 약 64.9% 의 FLOPs 절감 효과를 보였습니다.
• 공동 학습의 효과:
  • 인코더와 생성 모델을 분리하여 학습하는 경우 ($\beta \to 0$) 보다, 공동 학습 ($\beta > 0$) 을 할 때 FID 가 약 17% 개선되었습니다. 이는 인코더가 생성 과정에 최적화된 잠재 표현을 학습했음을 의미합니다.
• 유연한 사전 분포:
  • 가우시안뿐만 아니라 균일 (Uniform), 라플라스 (Laplacian), 가우시안 혼합 (Gaussian Mixture) 등 다양한 사전 분포 ($p_0$) 를 사용할 수 있음을 실험적으로 증명했습니다. 이는 SI 의 핵심 강점인 유연성을 잠재 공간에서도 유지함을 보여줍니다.
• 샘플링 유연성:
  • 분류기 없는 안내 (Classifier-Free Guidance, CFG) 와 DDIM 역변환 (Inversion) 을 지원하여, 특정 클래스에 대한 샘플 생성이나 이미지 편집이 가능함을 보였습니다.

5. 의의 및 결론 (Significance)

• 효율성과 성능의 균형: LSI 는 고차원 관측 공간에서의 SI 적용에 따른 계산 부담을 잠재 공간으로 이동시켜 해결하면서도, SI 가 가진 유연한 분포 변환 능력을 유지합니다.
• 새로운 학습 패러다임: 단순한 사전 분포에 의존하는 기존 확산 모델의 한계를 넘어, 학습된 잠재 공간에서 복잡한 생성 과정을 학습할 수 있는 새로운 길을 열었습니다.
• 확장성: ImageNet 과 같은 대규모 데이터셋에서 경쟁력 있는 성능을 보여주며, 차세대 고효율 생성 모델의 기초가 될 수 있음을 입증했습니다.

결론적으로, 이 논문은 잠재 공간에서의 확률적 보간을 가능하게 함으로써, 생성 모델의 학습 효율성, 계산 비용, 그리고 표현력을 동시에 향상시키는 중요한 이정표를 제시합니다.