[논문 리뷰] Latent Forcing: 확산 모델의 생성 순서를 재정의하여 픽셀을 지배하다

TL;DR

기존 이미지 생성 모델들은 고품질과 높은 계산 효율 사이에서 어려운 줄다리기를 해왔습니다. Latent Forcing은 이 문제를 해결하기 위해 잠재 공간(Latent Space)과 픽셀 공간(Pixel Space)을 동시에 활용하는 새로운 패러다임을 제시합니다. 이 방법론은 잠재 공간을 전체적인 구조를 잡는 '스케치북'으로, 픽셀 공간을 세밀한 디테일을 채우는 '캔버스'로 사용하는 Coarse-to-Fine (거시에서 미시로) 전략을 채택합니다. ImageNet 데이터셋 실험 결과, Latent Forcing은 기존 SOTA 모델들을 능가하는 성능을 보였으며, 특히 조건 없는(unconditional) 생성에서도 압도적인 결과를 달성했습니다.

연구 배경 및 동기

이미지 생성 분야의 양대 산맥은 픽셀 공간 확산 모델과 **잠재 공간 확산 모델(LDM)**이었습니다. 픽셀 공간 모델(e.g., DDPM, DiT)은 이미지 픽셀에 직접 노이즈를 가하고 제거하며 최고 수준의 품질을 달성하지만, 고해상도 이미지에서는 계산량이 막대하게 증가합니다. 반면, LDM은 이미지를 저차원 잠재 공간으로 압축하여 계산 효율을 극적으로 높였지만, 이 과정에서 발생하는 정보 손실로 인해 미세한 텍스처 복원에 한계를 보입니다.

Latent Forcing은 "왜 둘 중 하나를 선택해야 하는가?"라는 질문에서 출발합니다. 이 연구는 두 공간의 장점만을 결합하여, 잠재 공간에서 이미지의 전반적인 의미와 구조를 빠르게 잡아낸 뒤, 이 정보를 바탕으로 픽셀 공간에서 손실된 고주파 디테일을 완벽하게 복원하는 새로운 길을 제시합니다.

연구	접근 방식	장점	단점
DDPM	픽셀 공간 확산	최고 수준 품질	높은 계산 비용
VAE	잠재 공간 인코딩	높은 효율성	정보 손실, 흐릿함
GAN	적대적 학습	사실적인 이미지	훈련 불안정
LDM	잠재 공간 확산	높은 효율성	정보 병목, 디테일 손실
DiT	트랜스포머 기반 픽셀 확산	최고 수준 품질	매우 높은 계산 비용
Latent Forcing	잠재+픽셀 공간 동시 활용	최고 수준 품질 + 높은 효율	증가된 모델 복잡도

핵심 기여

잠재 공간과 픽셀 공간의 시너지: 두 공간을 동시에 처리하는 단일 확산 모델을 제안하여, 효율성과 품질을 모두 잡았습니다.
Latent-First 스케줄링: 잠재 공간의 노이즈를 먼저 제거하고, 이를 기반으로 픽셀 공간의 디테일을 복원하는 순차적 노이즈 제거 전략의 효과를 입증했습니다.
다중 시간 변수와 스케줄링: 잠재 공간과 픽셀 공간에 각각 독립적인 시간 변수( $t_{latent}$ , $t_{pixel}$ )와 노이즈 스케줄을 적용하여 유연하고 효과적인 제어를 가능하게 했습니다.
통합된 손실 함수: 다중 모달리티(잠재, 픽셀)를 동시에 학습하기 위한 일반화된 손실 함수를 설계했습니다.
실험적 SOTA 달성: ImageNet 데이터셋에서 기존 모든 모델을 능가하는 FID 점수를 기록하며 그 성능을 입증했습니다.

제안 방법론

Latent Forcing의 핵심은 단일 확산 모델(본 논문에서는 DiT 사용)이 [압축된 잠재 표현 z, 원본 픽셀 x] 쌍을 동시에 입력받아, 정보 생성의 '순서'를 제어하는 것입니다.

모델 아키텍처 및 Latent-First 스케줄링

모델은 노이즈 제거 과정을 두 단계로 나누어 수행하도록 훈련됩니다.

1단계 (초기 타임스텝, 예: $t=1000 \rightarrow 500$ ): 모델은 잠재 공간( $z$ )의 노이즈 제거에 집중합니다. 이 단계에서 이미지의 전반적인 구조, 의미, 색상 분포 등 거시적인 정보가 결정됩니다. 픽셀 공간( $x$ )의 정보는 보조적으로만 사용됩니다.
2단계 (후기 타임스텝, 예: $t=500 \rightarrow 0$ ): 1단계에서 거의 깨끗해진 잠재 공간의 정보를 강력한 조건(condition)으로 활용하여, 픽셀 공간( $x$ )의 노이즈를 제거합니다. 이 단계에서 잠재 공간 인코딩 시 손실되었던 고주파의 세밀한 디테일과 텍스처가 복원됩니다.

이러한 'Latent-First' 스케줄링은 마치 화가가 처음엔 연필로 전체 구도를 스케치하고, 나중에 물감으로 세부를 칠하는 과정과 유사합니다.

핵심 수식

모델은 다중 모달리티에 대한 일반화된 v-prediction 손실 함수를 사용하여 훈련됩니다. v-prediction은 노이즈( $\epsilon$ ) 대신 데이터의 '속도(velocity)'를 예측하는 방식으로, 훈련 안정성과 생성 품질을 높이는 최신 기법입니다.

L = \sum_{i \in \{z, x\}} \lambda_i \mathbb{E} \|v_{\theta,i}(z_{t_z}, x_{t_x}, t_z, t_x) - v_i(i_{t_i})\|^2

$i$ : 데이터 양식 (잠재 공간 $z$ 또는 픽셀 공간 $x$ )
$z_{t_z}, x_{t_x}$ : 각기 다른 노이즈 스케줄에 따라 노이즈가 추가된 데이터
$v_{\theta,i}(...)$ : 모든 데이터 양식을 입력받아 $i$ 번째 데이터의 속도를 예측하는 단일 모델
$v_i(i_{t_i})$ : 모델이 예측해야 할 실제 목표 속도 (깨끗한 데이터를 향하는 방향 벡터)
$\lambda_i$ : 각 데이터 양식의 손실에 대한 가중치

이 손실 함수를 통해 단일 모델 $v_\theta$ 는 잠재-픽셀 공간 간의 복잡한 관계를 학습하고, 스케줄링된 시간표에 따라 특정 공간의 노이즈를 우선적으로 제거하는 방법을 터득하게 됩니다.

실험 설정

데이터셋: ImageNet 256x256, 512x512
평가 지표: FID (Fréchet Inception Distance), 낮을수록 좋음
베이스라인: DiT (픽셀 공간), LDM (잠재 공간) 등 SOTA 모델
백본 아키텍처: Diffusion Transformer (DiT-L/2)

하이퍼파라미터	값
배치 사이즈	1024
학습률	1e-4
옵티마이저	AdamW
VAE 인코더	LDM에서 사용된 사전 훈련된 VAE

실험 결과 분석

Latent Forcing은 ImageNet 256x256 클래스 조건부 생성 태스크에서 FID 1.98을 달성하며, 기존 최강자였던 DiT-XL/2 (FID 2.27)를 큰 차이로 앞질렀습니다. 이는 픽셀 공간에서 작동하는 확산 모델의 새로운 SOTA(State-of-the-Art) 기록입니다.

모델 (ImageNet 256x256)	FID 점수 (낮을수록 좋음)
LDM-4 (ViT-G/2)	3.60
DiT-XL/2	2.27
Latent Forcing (DiT-L/2)	1.98

특히 주목할 점은, 훨씬 작은 모델(DiT-L)을 사용했음에도 더 큰 모델(DiT-XL)을 능가했다는 것입니다. 이는 Latent Forcing의 방법론 자체가 매우 효율적이고 강력함을 시사합니다.

Ablation Study

Ablation Study(제거 연구)를 통해 Latent Forcing의 핵심 요소들의 중요성을 검증했습니다.

스케줄링 전략: 'Latent-First' 스케줄이 'Pixel-First'나 두 공간을 동시에 처리하는 'Joint' 스케줄에 비해 월등히 우수한 성능을 보였습니다. 이는 정보 생성의 '순서'가 최종 품질에 결정적인 영향을 미친다는 가설을 강력하게 뒷받침합니다.
다중 시간 변수: 잠재 공간과 픽셀 공간에 독립적인 시간 변수를 사용하는 것이 단일 시간 변수를 공유하는 것보다 더 좋은 결과를 낳았습니다.

비판적 평가

강점

효율성과 품질의 완벽한 조화: LDM처럼 빠르게 전체 구조를 잡으면서도, DiT와 같은 픽셀 기반 모델 이상의 최종 품질을 달성합니다.
'Latent-First' 전략의 유효성 입증: '무엇을' 생성하는가 만큼 '어떤 순서로' 정보를 구성하는가가 중요함을 실험적으로 증명했습니다.
압도적인 무조건부 생성 성능: 클래스 레이블 같은 조건 없이 이미지를 생성하는 무조건부 생성에서도 SOTA를 달성하여, 모델이 데이터 분포 자체를 깊이 이해하고 있음을 보여주었습니다.

한계점

증가된 복잡도 및 메모리 요구량: 모델이 두 가지 데이터 양식을 동시에 처리해야 하므로, 단일 공간 모델에 비해 학습 시 더 많은 GPU 메모리를 요구할 수 있습니다.
VAE 의존성: 잠재 공간 표현을 위해 사전 훈련된 VAE 인코더에 의존하므로, VAE의 성능이 전체 모델의 성능에 영향을 미칠 수 있습니다.

향후 연구 방향

Latent Forcing의 'Coarse-to-Fine' 접근 방식은 이미지 생성을 넘어 다양한 분야로 확장될 잠재력이 큽니다.

비디오 생성: 초기 프레임에서 전체적인 장면과 움직임을 생성하고, 후반 프레임에서 디테일을 추가하는 방식
3D 모델링: 저해상도 메쉬로 전체 형태를 잡고, 고해상도 텍스처와 지오메트리를 추가하는 방식
음악 생성: 멜로디와 코드 진행을 먼저 구성하고, 그 위에 세부적인 악기 편곡과 음향 효과를 입히는 방식

논문 링크: arXiv:2402.11401
관련 자료: ImageNet Dataset

[논문 리뷰] Latent Forcing: Reordering the Diffusion Trajectory for Pixel-Space Image Generation