[논문 리뷰] Unified Latents (UL): 잠재 공간과 생성 품질의 통합

TL;DR

기존 잠재 확산 모델(Latent Diffusion Model)은 잠재 공간의 정보량과 최종 생성 품질 사이의 미묘한 균형을 맞추는 데 어려움을 겪었습니다. 정보가 너무 많으면 생성 품질이 저하되고, 너무 적으면 원본 복원력이 떨어집니다. **Unified Latents (UL)**는 이 문제를 해결하기 위해 제안된 새로운 프레임워크입니다. UL은 **확산 사전 분포(diffusion prior)**와 **확산 복호기(diffusion decoder)**를 함께 학습시키고, 인코더가 출력하는 잠재 변수의 노이즈 수준을 사전 분포의 최소 노이즈 수준과 일치시킵니다. 이 간단하면서도 강력한 아이디어를 통해 복잡한 학습 목표를 안정적인 MSE 손실 함수로 대체했으며, ImageNet 512x512에서 FID 1.41, Kinetics-600에서 FVD 1.3이라는 SOTA(최첨단) 성능을 달성하며 학습 효율성까지 크게 높였습니다.

연구 배경 및 동기

잠재 확산 모델은 고해상도 이미지 생성의 핵심 기술로 자리 잡았습니다. 이 모델의 성공은 픽셀 공간 대신 압축된 잠재 공간(latent space)에서 확산 과정을 수행하여 계산 효율을 높인 덕분입니다. 하지만 이는 '잠재 공간에 얼마나 많은 정보를 담아야 하는가?'라는 근본적인 질문을 낳습니다.

정보량이 너무 많으면 (Low Compression): 원본 이미지를 거의 완벽하게 복원할 수 있지만, 잠재 공간의 분포가 너무 복잡해져 새로운 이미지를 생성하는 사전 분포 모델(prior model)이 학습하기 어렵습니다. 이는 결국 생성 품질 저하로 이어집니다.
정보량이 너무 적으면 (High Compression): 잠재 공간의 분포는 단순해져 생성은 쉬워지지만, 원본 이미지의 중요한 디테일이 소실되어 복원 품질이 떨어집니다.

이러한 정보량(rate), 왜곡(distortion), 인식 품질(perception) 간의 상충 관계(trade-off)를 해결하는 것이 이 연구의 핵심 동기입니다. 본 연구는 "어떻게 잠재 공간의 정보량을 체계적으로 제어하면서 복원과 생성 품질을 모두 극대화할 수 있는가?"라는 질문에 답하고자 합니다.

연구	잠재 공간 정보량 제어	생성 품질	학습 안정성	학습 비용
Stable Diffusion	어려움	높음	보통	높음
VAE	용이	낮음	높음	낮음
GANs	보통	높음	낮음	보통
Diffusion Models	보통	높음	높음	매우 높음
UL (본 논문)	용이	매우 높음	높음	상대적 낮음

핵심 기여

Unified Latents (UL) 프레임워크 제안: 인코더, 확산 사전 분포, 확산 복호기를 하나의 목표 아래 통합적으로 학습하는 새로운 프레임워크를 제시했습니다.
체계적인 잠재 공간 정보량 제어: 인코더 출력 노이즈와 사전 분포의 최소 노이즈 수준을 연동시켜, 잠재 공간의 정보량을 명시적으로 제어하는 메커니즘을 고안했습니다.
단순하고 효율적인 학습 목표: 복잡한 KL-Divergence 항을 피하고 간단한 MSE 손실 함수만으로 모델을 안정적으로 학습시킬 수 있음을 보였습니다. 이는 잠재 공간의 비트레이트(정보량)에 대한 명확한 상한선을 제공합니다.
최고 수준의 성능 달성: ImageNet 512x512에서 FID 1.41, Kinetics-600에서 FVD 1.3을 달성하여 기존 SOTA 모델들을 능가하는 성능을 입증했습니다.

제안 방법론

UL 프레임워크는 인코더(Encoder), 확산 사전 분포(Diffusion Prior), 확산 복호기(Diffusion Decoder) 세 가지 구성 요소를 함께 학습시킵니다.

핵심 아이디어: 정보량 제어와 학습 목표 단순화

UL의 가장 중요한 아이디어는 인코더가 출력하는 잠재 변수 $z_0$ 의 노이즈 수준을 확산 사전 분포가 가정하는 최소 노이즈 수준과 일치시키는 것입니다.

인코더의 역할: 인코더 $E$ 는 입력 이미지 $x$ 를 '깨끗한' 잠재 변수 $z_{clean} = E(x)$ 로 압축합니다. 그리고 여기에 미리 정해진 분산 $\sigma_E^2$ 을 갖는 가우시안 노이즈 $\epsilon$ 을 더해 최종 잠재 변수 $z_0 = z_{clean} + \sigma_E \cdot \epsilon$ 를 만듭니다.
사전 분포와의 연결: 확산 사전 분포 모델 $p(z_0)$ 는 순수한 노이즈에서 시작하여 $z_0$ 를 생성하는 모델입니다. UL은 이 사전 분포의 확산 과정(diffusion process)이 시작되는 최소 노이즈 수준을 인코더가 추가하는 노이즈 수준 $\sigma_E^2$ 과 정확히 일치시킵니다.
학습 목표 단순화: 이 연결 덕분에 VAE의 학습 목표에 등장하는 복잡한 $KL(q(z_0|x) || p(z_0))$ 항을 최적화하는 문제가, 사전 분포 모델이 $z_0$ 에서 노이즈를 예측하는 간단한 MSE 손실( $L_z$ ) 문제로 바뀝니다. 즉, 잠재 공간의 정규화와 생성을 위한 학습이 하나의 간단한 목표로 통합됩니다.

모델 아키텍처

UL 프레임워크의 학습 및 추론 과정은 다음과 같습니다.

                  [학습 과정]
                  -----------
[이미지 x] -> [인코더 E] -> [z_clean] --(고정 노이즈 추가)--> [z_0]
                                                               |
      +--------------------------------------------------------+
      |                                                        |
      v                                                        v
[확산 복호기 D_dec] <--(x_t, t)-- [원본 이미지 x]         [확산 사전 분포 D_prior] <--(z_t', t')-- [z_clean]
      |                                                        |
      v                                                        v
   [L_x: 복원 손실]                                        [L_z: 사전 분포 손실]

                  [추론(생성) 과정]
                  -----------------
[랜덤 노이즈] -> [확산 사전 분포 D_prior] -> [생성된 z_0] -> [확산 복호기 D_dec] -> [생성된 이미지 x']

인코더 ( $E$ ): 이미지를 저차원의 깨끗한 잠재 변수 $z_{clean}$ 으로 매핑합니다.
확산 사전 분포 ( $D_{prior}$ ): 잠재 공간의 분포를 학습합니다. 학습 시에는 $z_0$ 에서 노이즈를 제거하여 $z_{clean}$ 을 예측하고, 추론 시에는 순수 노이즈로부터 새로운 $z_0$ 를 생성합니다.
확산 복호기 ( $D_{dec}$ ): 잠재 변수 $z_0$ 를 조건으로 받아 원본 이미지를 복원합니다. $z_0$ 가 주어졌을 때, 노이즈 낀 이미지 $x_t$ 에서 원본 이미지 $x$ 를 예측하도록 학습됩니다.

핵심 수식

UL의 학습은 VAE의 증거 하한(ELBO)을 최적화하는 관점에서 이해할 수 있습니다. $\log p(x) \ge \mathbb{E}_{q(z_0|x)}[\log p(x|z_0)] - KL(q(z_0|x) || p(z_0))$ UL은 이 두 항을 각각 확산 모델의 손실 함수로 치환하여 최적화합니다.

복호기 손실 (Decoder Loss, $L_x$ ): 위 식의 첫 번째 항인 복원 확률( $\log p(x|z_0)$ )을 최대화하는 과정입니다. 이는 $z_0$ 를 조건으로 이미지 $x$ 의 노이즈를 예측하는 표준적인 조건부 확산 모델의 손실 함수와 같습니다. $L_x(\theta) = \mathbb{E}_{x, z_0, \epsilon, t} \left[ w(t) \| \epsilon - \epsilon_\theta(x_t, t, z_0) \|^2 \right]$ 여기서 $x_t$ 는 원본 이미지 $x$ 에 시간 $t$ 만큼 노이즈 $\epsilon$ 을 추가한 것이고, $\epsilon_\theta$ 는 이를 예측하는 신경망(복호기)입니다. $w(t)$ 는 시간에 따른 가중치입니다.
사전 분포 손실 (Prior Loss, $L_z$ ): 두 번째 항인 KL-Divergence를 최소화하는 과정입니다. 인코더 출력 $q(z_0|x)$ 와 사전 분포 $p(z_0)$ 를 일치시키는 것을 목표로 합니다. UL의 핵심 설계 덕분에, 이는 $z_0$ 에서 $z_{clean}$ 을 예측하는 확산 사전 분포 모델의 손실 함수로 간단히 표현됩니다. $L_z(\theta) = \mathbb{E}_{z_{clean}, \epsilon', t'} \left[ w(t') \| z_{clean} - D_{prior}(z_{t'}, t') \|^2 \right]$ 여기서 $z_{t'}$ 는 $z_{clean}$ 에 노이즈를 추가한 것이고, $D_{prior}$ 는 $z_{clean}$ 을 예측하는 신경망(사전 분포 모델)입니다.
전체 손실: 최종적으로 두 손실을 더하여 전체 모델을 end-to-end로 학습합니다. $L(\theta) = L_x(\theta) + L_z(\theta)$

실험 설정

데이터셋: ImageNet 512x512 (이미지), Kinetics-600 (비디오)
평가 지표:
- 생성 품질: FID (이미지), FVD (비디오)
- 복원 품질: PSNR, rFID (Reconstruction FID)
베이스라인: LDM (Latent Diffusion Model), VQ-GAN, StyleGAN-XL 등 기존 SOTA 모델
아키텍처: 인코더/복호기는 2D/3D U-Net 구조를, 사전 분포 모델은 DiT(Diffusion Transformer)를 사용했습니다.

실험 결과 분석

UL은 이미지 및 비디오 생성 태스크에서 기존 모델들을 압도하는 성능을 보였습니다.

데이터셋	평가지표	UL 성능	이전 SOTA
ImageNet 512x512	FID (lower is better)	1.41	1.77 (LDM-B)
Kinetics-600	FVD (lower is better)	1.3	1.8 (StyleGAN-V)

생성 품질: ImageNet에서 FID 1.41을 달성하여, 픽셀 공간에서 직접 학습한 ADM-G 모델(FID 2.97)이나 기존 잠재 확산 모델 LDM(FID 1.77)보다 월등히 높은 품질의 이미지를 생성했습니다.
복원 품질: 높은 압축률에도 불구하고 뛰어난 복원 품질(rFID)을 유지하여, 정보량과 복원력 사이의 균형을 성공적으로 맞추었음을 보여주었습니다.
학습 효율성: UL은 기존 SOTA 모델보다 훨씬 적은 계산량으로 동등하거나 더 나은 성능을 달성했습니다.

Ablation Study

주요 설계 요소의 중요성을 검증하기 위한 분석 결과, 인코더의 고정된 출력 노이즈와 확산 사전 분포를 통합적으로 학습하는 것이 성능 향상의 핵심 요인임이 밝혀졌습니다. 만약 두 요소를 분리하거나 기존 VAE처럼 KL-Divergence를 직접 최적화할 경우, 성능이 크게 하락했습니다.

비판적 평가

강점:

우아한 통합: VAE의 잠재 공간 제어와 확산 모델의 생성 능력을 이론적으로나 실용적으로 매우 깔끔하게 통합했습니다.
뛰어난 성능과 효율성: 적은 계산 자원으로 SOTA를 달성하여 잠재 확산 모델의 새로운 표준을 제시했습니다.
체계적 제어 가능성: 잠재 공간의 정보량을 $\sigma_E$ 라는 단일 하이퍼파라미터로 제어할 수 있어, 응용 목적에 맞게 모델을 튜닝하기 용이합니다.

한계점과 개선 방향:

공간적 다운샘플링의 영향: 잠재 공간의 채널 수에는 덜 민감하지만, 공간적 다운샘플링 비율(spatial downsampling factor)은 여전히 성능에 큰 영향을 미칩니다. 최적의 비율을 찾는 추가적인 연구가 필요합니다.
2단계 학습: 논문에서는 복호기를 먼저 학습시킨 후 사전 분포 모델을 학습시키는 2단계 접근법을 사용했습니다. 완전한 end-to-end 동시 학습의 안정성에 대한 추가 검증이 필요할 수 있습니다.

재현성 평가: 제안된 방법론은 명확한 수식과 아키텍처를 기반으로 하며, 공식 코드가 공개되어 재현성이 높습니다. 다만, 대규모 데이터셋과 컴퓨팅 자원이 필요하여 개인 연구자가 동일한 규모로 재현하기에는 어려움이 따를 수 있습니다.

향후 연구 방향

UL 프레임워크는 이미지, 비디오를 넘어 다양한 데이터 모달리티로 확장될 잠재력이 큽니다.

다양한 응용 분야: 텍스트-이미지 생성, 3D 에셋 생성, 오디오 합성 등 다양한 생성 태스크에 UL의 아이디어를 적용할 수 있습니다.
모델 경량화: 더 효율적인 인코더 및 사전 분포 모델 아키텍처를 탐구하여 모바일 등 경량 환경에서의 사용 가능성을 모색할 수 있습니다.
조건부 생성 강화: 복잡한 조건(예: 상세한 텍스트 설명, 레이아웃 맵)을 처리하기 위한 메커니즘을 통합하여 제어 가능성을 더욱 높이는 연구가 기대됩니다.

실무 적용 가이드

실무 환경에서 UL을 도입할 때 고려할 점은 다음과 같습니다.

하이퍼파라미터 튜닝: 가장 중요한 하이퍼파라미터는 인코더의 출력 노이즈 분산 $\sigma_E$ 입니다. 이 값을 조절하여 생성 품질과 복원 품질 간의 균형을 맞출 수 있습니다.
데이터 전처리: 모든 생성 모델과 마찬가지로, 고품질의 정제된 학습 데이터는 최종 결과물의 질을 결정하는 핵심 요소입니다.
사전 학습 모델 활용: 처음부터 모든 것을 학습시키기보다, 공개된 사전 학습 모델을 기반으로 특정 도메인에 맞게 파인튜닝하는 것이 효율적입니다.

결론

**Unified Latents (UL)**는 잠재 확산 모델의 오랜 숙제였던 정보량과 생성 품질 간의 상충 관계를 효과적으로 해결한 혁신적인 프레임워크입니다. 인코더, 확산 사전 분포, 확산 복호기를 하나의 통합된 목표 아래 학습시키는 독창적인 접근법을 통해, 기존 모델들을 뛰어넘는 SOTA 성능과 학습 효율성을 동시에 달성했습니다. UL은 향후 다양한 생성 모델 연구에 중요한 이론적, 실용적 기반을 제공할 것으로 기대됩니다.

참고 자료

논문 원문: Unified Latents: Taming the Latent Space of Diffusion Models
프로젝트 페이지: https://unified-latents.github.io/
공식 코드 저장소: https://github.com/google-research/google-research/tree/master/ul

[논문 리뷰] Unified Latents (UL): How to train your latents