기술 요약: ROMS-IMLE

문제 정의

생성 모델링은 변이형 오토인코더(VAE)와 생성적 적대 신경망(GAN)에서 확률적 보간(Stochastic Interpolant, SI) 모델인 확산 모델(Diffusion) 및 플로우 매칭(Flow Matching)으로 진화해 왔습니다. SI 모델은 이미지 합성 분야에서 최첨단 기술을 정의하고 있지만, 복잡한 메커니로리(complex machinery)에 의존합니다. 즉, 정교한 아키텍처(예: Transformer), 고도화된 샘플링 기법, 그리고 반복적인 수치 적분이 필요합니다. 기존의 믿음은 고품질 생성을 위해서는 노이즈에서 데이터로의 변환을 여러 개의 작고 순차적인 단계로 분해해야 한다는 것입니다. 그러나 이러한 반복적인 과정은 계산 비용이 많이 들며, 단 하나의 샘플을 생성하기 위해 여러 번의 순방향 패스(높은 NFE, 함수 평가 횟수)를 요구합니다.

저자들은 이러한 반복적인 분해가 진정으로 필요한 것인지에 대해 의문을 제기합니다. 그들은 SI 모델의 성공이 반복적인 샘사플링 과정 자체에서 오는 것이 아니라, 최소한의 단일 단계 프레임워크에 적응할 수 있는 특정 학습 특성, 즉 **단계별 지도 학습(per-stage supervision)**과 **스펙트럼 특화(spectral specialization)**에서 비롯된다고 가정합니다. 목표는 반복적인 디노이징의 계산적 오버헤드 없이 단 한 번의 순방향 패스로 고품질의 샘플을 생성하는 경쟁력 있는 생성 모델을 설계하는 것입니다.

방법론: ROMS-IMLE

제안된 방법론인 **Robust Multi-Stage IMLE (ROMS-IMLE)**는 **암묵적 최대 가능도 추정(Implicit Maximum Likelihood Estimation, IMLE)**을 기반으로 구축되었습니다. IMLE는 생성된 샘플 집합 내의 실제 이미지와 가장 가까운 이웃 사이의 거리를 최소화함으로써 사전 분포(prior distribution)를 데이터 분포로 매핑하는 단일 단계 생성기입니다. 저자들은 VAE 경계, 적대적 학습, Transformer를 배제한 IMLE의 단순성을 유지하면서도, SI 모델의 강점을 모방하기 위해 두 가지 핵심적인 수정 사항을 학습 목적 함수에 도입했습니다.

1. 스펙트럼 인지 단계별 지도 학습 (Spectrally-Aware Per-Stage Supervision)

저자들은 SI 모델에 대해 "테스트 중심적" 관점을 채택합니다. 확산을 반복적인 디노이징 과정으로 보는 대신, 샘플러를 함수의 합성 $h_\theta = g^1_\theta \circ g^2_\theta \circ \dots \circ g^T_\theta$로 간주합니다. 이 관점에서 SI 모델은 단계별 지도 학습을 제공합니다. 즉, 모든 중간 레이어(특정 노이즈 레벨 또는 타임 스텝에 대응하는 레이어)가 직접적인 학습 신호를 받습니다. 또한, 이 지도 학습은 스펙트럼 구조를 가집니다. 초기 단계(높은 노이즈)는 거친 저주파 구조를 학습하고, 후기 단계(낮은 노이즈)는 미세한 고주파 세부 사항을 학습합니다.

이를 단일 단계 IMLE 생성기 $f_\theta$에 적용하기 위해 저자들은 다음과 같이 수행합니다:

• 아키텍처 스캐폴드(Architectural Scaffold): 생성기를 업샘플링 블록의 스택($f_\theta = f_{\theta_L} \circ \dots \circ f_{\theta_1}$)으로 구성하며, 각 단계는 점진적으로 높아지는 해상도에서 작동합니다.
• 단계별 손실(Per-Stage Loss): 모든 단계 $l$에 출력 헤드를 도입합니다. 학습 목적 함수는 모든 단계의 손실을 합산하며, 단계 $l$의 출력($\tilde{x}_l$)을 해당 단계의 해상도로 다운샘플링된 실제 이미지($K_l(x_i)$)와 비교합니다.
• 스펙트럼 특화(Spectral Specialization): 다운샘플링 커널(저역 통과 필터)을 타겟으로 사용함으로써, 모델은 초기 단계에서는 거친 구조를 학습하고 후기 단계에서는 고주파 세부 사항을 학습하도록 암묵적으로 강제됩니다. 이는 확산 모델의 스펙트럼 진행을 모방합니다.

2. 불일치 쌍을 위한 강건한 손실 함수 (Robust Loss for Mismatched Pairs)

단일 단계 IMLE에서 모델은 실제 이미지를 가장 가까운 생성된 샘플과 매칭합니다. 샘플링의 무작위성으로 인해, 실제 이미지가 실제로는 올바른 모드(mode)를 나타내지 않는 생성된 샘플과 매칭될 수 있습니다(즉, "불일치 쌍"). 표준 제곱 $\ell_2$ 손실은 이러한 이상치(outlier)에 민감합니다. 큰 잔차(residual)가 불균형하게 큰 그래디언트를 생성하여 학습을 불안정하게 만들 수 있기 때문입니다.

저자들은 표준 거리 메트릭을 강건한 손실 함수(구체적으로 Geman-McClure)로 대체합니다. 이 함수는 잔차에 대해 이차 함수보다 느리게 증가(sub-quadratically)하여, 불일치 쌍의 영향을 줄이고 잘 매칭된 쌍이 그래디언트 업데이트를 주도할 수 있게 합니다. 이를 통해 근접 이웃 매칭이 불완전하더라도 안정적인 학습을 보장합니다.

3. 라운드 트립 리젝션 (Round-Trip Rejection, 잠재 공간)

고해상도 생성(ImageNet 256)을 위해 모델은 동결된 오토인코더의 잠재 공간(latent space)에서 작동합니다. 생성기가 디코딩 성능이 떨어지는 잠재 변수를 생성할 수 있는 분포 불일치 문제를 해결하기 위해, 저자들은 **라운드 트립 리젝션(round-trip rejection)**을 도입합니다. 추론 과정에서 생성된 이미지를 다시 인코딩 및 디코딩합니다. 만약 원래 생성된 이미지와 재디코딩된 이미지 사이의 거리(라운드 트립 비용)가 높으면 해당 샘플을 폐기합니다. 이는 고충실도(high-fidelity) 샘플만을 출력하도록 하는 필터 역할을 합니다.

주요 기여

1. 최소한의 단일 단계 설계: 본 논문은 단순한 ConvNeXt 기반 아키텍처와 단일 단계 순방향 패스를 사용하여 경쟁력 있는 생성 모델을 구축할 수 있음을 입증하며, Transformer와 반복적 샘플링의 복잡성을 피했습니다.
2. IMLE에서의 단계별 지도 학습: 저자들은 반복적인 SI 모델의 중간 단계 지도 학습 개념을 단일 단계 IMLE 프레임워크로 성공적으로 전이시켰으며, 여러 해상도에서의 직접적인 지도가 품질에 결정적임을 보여주었습니다.
3. 강건한 학습 목적 함수: 강건한 손실 함수(Geman-McClure)의 통합은 IMLE의 근접 이웃 불일치로 인한 불안정성을 효과적으로 완화합니다.
4. 효율성: 결과 모델은 기존의 확산 기반 베이스라인 모델들과 비교하여 훨씬 적은 파라미터와 함수 평가 횟수(NFE)로 최첨단 성능을 달s성합니다.

실험 결과

모델은 CIFAR-10, CelebA-HQ, ImageNet 256 데이터셋에 대해 평가되었습니다.

• CIFAR-10 (Unconditional): ROMS-IMLE는 단일 NFE로 FID 4.01을 달성했습니다. 특히, 비교된 모든 방법 중 가장 높은 **정밀도(Precision, 0.91)**와 **재현율(Recall, 0.80)**을 기록하며, 100단계 DDIM 베이스라인보다 높은 커버리지와 충실도를 보였습니다.
• CelebA-HQ (Unconditional): 모델은 FID 6.70을 달성했으며, 단일 단계 방법들 중 최고의 정밀도(0.96)와 재현율(0.61)을 기록하여 Shortcut(128 단계)과 같은 다단계 베이스라인을 두 지표 모두에서 능가했습니다.
• ImageNet 256 (Class-Conditional): 310M 파라미터를 사용하여 잠재 공간에서 작동하며(DiT-XL/2의 675M보다 적음), 분류기 없는 가이던스(classifier-free guidance) 없이 단일 순방향 패스로 FID 2.56에 도달했습니다. 이는 다단계 확산 모델(예: 250 단계의 SiT-XL/2)과 경쟁할 만한 수준이며, 훨씬 더 효율적입니다(250배 적은 함수 평가).

의의 및 주장

본 논문은 현대 생성 모델의 높은 품질이 본질적으로 반복적인 샘플링에 묶여 있는 것이 아니라, 스펙트럼을 인지하는 단계별 지도 학습에 달려 있다고 주장합니다. 이러한 특정 학습 신호를 최소한의 IMLE 프레임워크에 주입함으로써, 저자들은 단일 단계 생성기가 복잡한 반복 모델의 샘플 품질과 다양성을 따라잡을 수 있음을 입증했습니다.

이 연구의 의의는 이미지 생성을 더 저렴하고 단순하게 만드는 데 있습니다. ROMS-IMLE는 반복적인 추론의 계산 비용 없이 경쟁력 있는 FID, 정밀도 및 재현율을 달성함으로써, 생성을 점진적인 변환으로 분해해야 할 필요성에 의문을 제기합니다. 저자들은 적절한 목적 함수를 갖춘 단순한 단일 단계 모델이 인기 있는 SI 모델의 강점을 유지할 수 있다고 결론짓습니다.