CrossFlow: One-Step Generation Across Latent and Pixel Spaces

기술 요약: CrossFlow: 잠재 공간과 픽셀 공간을 가로지르는 단일 단계 생성 (One-Step Generation)

1. 문제 정의

현재의 최첨단 이미지 생성 시스템, 특히 잠재 확산 모델(Latent Diffusion Models, LDMs)과 플로우 매칭(flow-matching) 방법론은 사전 확률(prior), 확률 경로(probability path), 그리고 예측 대상(prediction target)을 단일 표현 공간 내에 정의하여 작동합니다. LDMs는 압축된 오토인코더 잠재 공간(latent space)에서 생성 궤적을 수행함으로써 계산 효율성을 개선하지만, 잠재 공간에서의 반복적인 디노이징(denoising) 과정과 이를 다시 픽셀로 매핑하는 별도의 디코딩 단계라는 2단계 파이프라인에 의존합니다.

이러한 구조는 근본적인 불일치를 초래합니다:

• 최적화 불일치(Optimization Mismatch): 생성기는 잠재 변위(latent displacement)를 예측하도록 최적화되는 반면, 최종 이미지 품질은 별도로 학습된 디코더에 의해 결정됩니다.
• 분포 변화(Distribution Shift): 추론 시, 디코더는 훈련 과정 중 생성된 깨끗한 잠재 변수와 분포적으로 다를 수 있는 생성된 잠재 변수를 입력받게 됩니다. 이러한 변화는 오류를 증폭시킬 수 있습니다.
• 감독의 한계(Supervision Limitation): 생성기의 출력이 별도의 디코더를 통과해야만 하므로, 생성기가 이미지 수준의 지각적(perceptual) 또는 적대적(adversarial) 손실 함수를 직접 사용하여 학습될 수 없으며, 이는 직접적인 픽셀 공간 감독을 약화시킵니다.

본 논문은 노이즈 공간과 출력 공간이 반드시 동일해야 하는가라는 질문을 던지며, 잠재 표현의 효율성을 유지하면서도 이미지 자체를 직접 예측할 수 있는 프레임워크를 제안합니다.

2. 방법론: CrossFlow

CrossFlow는 노이즈가 섞인 잠재 입력을 단일 함수 평가를 통해 픽셀 공간의 이미지로 직접 매핑하는 교차 공간 플로우(cross-space flow) 정식화입니다. 이 방식은 감독 대상(supervised target)에서 잠재 공간의 속도(latent-space velocity)를 제거하는 목적 함수를 도출함으로써 공간 간의 불일치를 해결합니다.

핵심 기술 정식화:

• 통합 프레임워크: 저자들은 단일 단계 생성 패러다임을 위한 통합된 타겟 함수 $Z$를 정의합니다. 표준 방식들은 잠재 공간의 속도장(velocity fields)에 의존하며, 이는 교차 공간 생성을 방해합니다.
• 설계 기준: 교차 공간 감독을 가능하게 하기 위해, 저자들은 두 가지 기준을 설정합니다:
  1. 잠재 속도 의존성 제거: 목적 함수는 구조적으로 잠재 공간의 속도와 분리되어야 합니다. 즉, 네트워크의 출력은 잠재 변위 벡터가 아니어야 합니다.
  2. 종단 경계 조건(Endpoint Boundary Conditions): 프레임워크는 목적지 경계($r=0$)에서 유효한 이미지를 산출해야 합니다.
• CrossFlow 목적 함수: 이러한 기준을 충족함으로써, 저자들은 네트워크 $F_\theta$가 잠재 속도가 아닌 픽셀 공간의 이미지 $x_0$를 예측하는 일반적인 손실 함수(식 10)를 도출합니다.
  • 잠재 궤적(정의: $\alpha(t)$ 및 $\beta(t)$)은 여t전히 깨끗한 잠재 변수가 어떻게 노이즈화되는지를 결정합니다.
  • 감독 대상은 깨끗한 픽셀 이미지입니다.
  • 실험에 사용된 특정 목적 함수(식 11)는 잠재 궤적을 따른 야코비안-벡터 곱(Jacobian-Vector Product, JVP)을 통해 계산된 시간 미분 항을 포함하며, 이를 통해 모델이 잠재 속도를 직접 예측하지 않고도 올바른 플로우 역학을 학습하도록 보장합니다.

구현 세부 사항:

• 아키텍처: 모델은 이미지를 잠재 변수로 매핑하기 위한 고정된 사전 학습 이미지 인코더(VA-VAE)와, 노이즈가 섞인 잠재 변수로부터 픽셀로 직접 매핑하는 비전 트랜스포머(ViT) 생성기($F_\theta$)를 사용합니다.
• 훈련 파이프라인: 안정성을 보장하기 위해 2단계 접근 방식을 사용합니다:
  1. 표준 오토인코더 목적 함수(재구성 + 적대적)를 사용하여 인코더 $E$를 훈련합니다.
  2. $E$를 고정하고 CrossFlow 목적 함수를 사용하여 생성기 $F_\theta$를 처음부터 훈련합니다.
• 손실 함수: 전체 목적 함수는 픽셀 공간의 보조 손실들과 결합됩니다:
  • L1 재구성 손실(L1 Reconstruction Loss): 픽셀 단위 정렬.
  • 지각적 손실(Perceptual Loss): 구조적 유사성.
  • 적대적(GAN) 손실: 고주파 성분의 충실도.
  • 참고: 훈련의 불안정성을 방지하기 위해 지각적 손실 계산 시 플로우 목적 함수의 분석적 가중치 계수는 제외됩니다.

3. 주요 기여

1. 교차 공간 프레임워크: 입력 사전 확률(잠재 변수)과 출력 이미지(픽셀)가 서로 다른 표현 공간에 존재하는 교차 공간 생성 프레임워크인 CrossFlow를 도입했습니다.
2. 속도 없는 단일 단계 목적 함수: 모델이 단일 단계 잠재-픽셀 생성기이자 디코더 대체제 역할을 할 수 있게 하여, 추론 시 별도의 디코더가 필요 없는 새로운 유도 방식을 제시했습니다.
3. 직접적인 픽셀 공간 감독: 생성된 이미지에 대해 지각적 및 적대적 손실을 직접 사용하여 생성기를 훈련할 수 있으며, 중간 디코더 단계를 우회합니다.
4. 실증적 검증: 단일 함수 평가(1-NFE)만으로 ImageNet-1k (256×256)에서 경쟁력 있는 FID 점수를 달성하며, 기존의 단일 단계 잠재 방식들을 능가하고 다단계 잠재 확산 베이스라인과 대등한 성능을 보임을 입증했습니다.

4. 실험 결과

실험은 ImageNet-1k 256×256 해상도에서 클래스 조건부 생성을 대상으로 수행되었습니다.

• 주요 성능:

  • CrossFlow-XL (461M 파라미터)은 1-NFE에서 FID 1.62를 달성했습니다.
  • 이는 DiT-XL/2 (250 스텝에서 2.27 FID) 및 SiT-XL/2 (250 스텝에서 2.06 FID)와 같은 다단계 잠재 확산 베이스라인과 경쟁할 만한 수준입니다.
  • 또한 다른 단일 단계 잠재 방식(예: iMF-XL/2, 1.72 FID)을 크게 앞서며 별도의 디코더가 필요하지 않습니다.
  • CrossFlow-B (95M)와 CrossFlow-L (315M)은 각각 2.54와 1.69의 FID를 기록하며 확장성을 보여주었습니다.

• 절제 연구(Ablation Studies):

  • 인코더 선택: VA-VAE 인코더(32 채널)가 SD-VAE(4.08 FID)나 FLUX.2(3.81 FID)보다 뛰어난 2.54 FID를 기록하며 가장 좋은 결과를 냈습니다. 이는 잠재 표현 구조가 매우 중요하다는 것을 시사합니다.
  • 손실 구성 요소: 지각적 손실을 제거하면 FID가 30.21로 급격히 악화되었고, GAN 손실을 제거하면 8.86으로 떨어졌습니다. CrossFlow 플로우 목적 함수가 구조적 신호를 제공한다면, 지각적 및 적대적 손실은 고주파 충실도에 필수적입니다.
  • 디코더 대체: LightningDiT-B/1에서 생성된 잠재 변수의 디코더로 사용했을 때, CrossFlow는 훈련 전반에 걸쳐 표준 VA-VAE 디코더보다 일관되게 향상된 FID를 보여주었습니다.

5. 의의 및 주장

본 논문은 CrossFlow가 잠재 표현의 효율성과 직접적인 픽셀 공간 감독의 충실도 사이의 간극을 성공적으로 메웠다고 주장합니다. 생성기와 디코더를 하나의 모델로 통합하여 노이즈가 섞인 잠재 입상과 픽셀 공간 타겟을 사용하여 학습함으로써, CrossFlow는 다음과 같은 성과를 거두었습니다:

• 추론 시 명시적인 디코더 단계를 제거했습니다.
• 훈련 과정 중에 모델이 단순히 깨끗한 인코더 출력만이 아니라, 생성 과정에서 발생하는 노이즈 섞인/불완전한 잠재 변수에 노출되도록 했습니다.
• 교차 공간 플ow 목적 함수가 잠재 확산의 계산적 이점과 직접적인 픽셀 수준 적대적/지각적 훈련의 품질 향상을 결격할 수 있음을 입증했습니다.

저자들은 잠재 표현과 보조 손실이 중요하지만, 핵심 기여는 별도의 디코더 모듈 없이도 단일 단계 교차 공간 생성을 가능하게 하는 아키텍처 및 목적 함수의 전환에 있다고 언급하며 신중한 어조를 유지하고 있습니다.