기술 요약: 구형 잠재 인코더를 활용한 효율적인 이미지 합성

문제 제기

소수 단계 이미지 생성은 일관성 및 평균 흐름 기반 방법을 통해 크게 발전했으나, 이러한 접근법은 종종 훈련 불안정성, 하이퍼파라미터에 대한 민감도, 그리고 최적화 목표 간의 내재적 충돌로 고통받습니다. 대안으로 제안된 구형 인코더 (Sphere Encoder) [40] 는 잠재 표현을 초구면 (hypersphere) 에 투영함으로써 소수 단계에서 고품질 생성을 달성합니다. 그러나 구형 인코더 프레임워크는 두 가지 중요한 한계에 직면해 있습니다:

1. 계산 비효율성: 생성 과정은 모든 샘플링 단계마다 픽셀 공간과 잠재 공간 간의 반복적인 전환 (이미지로 디코딩한 후 다시 인코딩) 을 요구합니다. 이는 훈련 및 추론 모두에서 상당한 계산 오버헤드를 초래합니다.
2. 목표 충돌: 이 프레임워크는 단일 인코더 - 디코더 아키텍처 내에서 재구성과 생성을 공동으로 최적화합니다. 선행 연구는 이러한 작업들이 서로 다른 표현을 선호함을 시사합니다. 예를 들어, 낮은 노이즈 수준은 재구성을 개선하지만 초구면 커버리지를 감소시켜 생성 성능을 저하시킵니다. 이러한 긴장은 샘플 품질을 유지하기 위해 더 큰 아키텍처를 사용하도록 강요합니다.

방법론

저자들은 반복적인 픽셀 공간 연산을 제거하기 위해 재구성과 생성을 분리하는 구형 잠재 인코더 (Sphere Latent Encoder) 를 제안합니다. 핵심 혁신은 고정된 구형 잠재 공간 내에서 전체 생성 과정을 수행하는 데 있습니다.

아키텍처

이 방법은 세 가지 주요 구성 요소로 이루어집니다:

1. 고정 이미지 토크나이저: 사전 훈련된 표현 오토인코더 (RAE) [43] 가 고정된 인코더 ($E$) 와 디코더 ($D$) 로 사용됩니다. DINOv2 기반의 인코더는 이미지를 잠재 표현으로 매핑하고, 디코더는 이미지를 재구성합니다. 이 구성 요소는 훈련 중 고정됩니다.
2. 잠재 탈노이즈 모델: 별도의 트랜스포머 기반 탈노이즈 네트워크 ($G$, SiT 아키텍처 [25] 활용) 가 잠재 공간에서만 작동합니다. 이는 노이즈가 있는 구형 잠재 입력을 받아 깨끗한 잠재 값을 예측합니다.
3. 구형 투영: 잠재 표현은 RMSNorm 기반의 구형화 함수 ($F$) 를 사용하여 초구면으로 투영되어 균일한 잠재 구조를 보장합니다.

훈련 목적

픽셀 공간에서 손실을 계산하는 구형 인코더와 달리, 구형 잠재 인코더는 효율성을 높이기 위해 모든 감독 신호를 잠재 공간으로 이동시킵니다. 훈련은 서로 다른 스케일에서 가우시안 노이즈를 추가하여 생성된 두 개의 노이즈가 있는 구형 잠재, $v_{NOISY}$ 및 $v_{noisy}$를 포함합니다.

• 재구성 손실 ($L_{recon}$): 예측된 깨끗한 잠재 $G(v_{noisy})$ 와 실제 깨끗한 잠재 $z$ 사이의 $\ell_1$ 손실과 코사인 유사도 손실을 결합합니다. 이는 탈노이즈된 잠재가 깨끗한 표현과 정렬되도록 합니다.
• 일관성 손실 ($L_{cons}$): 서로 다른 노이즈 수준 ($v_{NOISY}$ 및 $v_{noisy}$) 에서의 예측 간 일관성을 강제합니다. 낮은 노이즈 잠재에서의 예측은 (stop-gradient 를 통해) 고정된 목표로 간주되어 초구면에서의 국부적 부드러움을 장려합니다.
• 잠재 일관성 손실 제거: 구형 인코더에서 사용되던 잠재 일관성 손실 (디코딩된 이미지를 다시 인코딩하는 방식) 을 명시적으로 제거했습니다. 저자들은 분리된 설정에서는 성능을 향상시키지 않으면서 계산 오버헤드만 추가한다고 판단했습니다.

샘플링 절차

샘플링은 완전히 잠재 공간에서 발생합니다 (알고리즘 1):

1. 가우시안 분포에서 잠재 값을 샘플링합니다.
2. 이를 초구면으로 투영합니다.
3. 탈노이즈 모델 $G$ 가 분류기 없는 가이드 (classifier-free guidance) 를 사용하여 잠재 값을 정제합니다.
4. 정제된 잠재 값을 다시 투영하고 노이즈로 교란시킨 후 다음 단계의 입력으로 사용합니다.
5. 이 반복적 정제는 $T$ 단계 동안 계속됩니다.
6. 디코더는 최종 잠재 값을 픽셀로 매핑하기 위해 마지막에 한 번만 호출됩니다.

저자들은 또한 소수 단계 샘플링에 더 잘 적합하도록 균일 샘플링과 구형 인코더의 샘플링과 구별되는 로그 - 정규 분포를 사용하는 수정된 노이즈 샘플링 전략을 제안합니다.

주요 기여

1. 구형 잠재 인코더 프레임워크: 반복적인 픽셀 - 잠재 전환을 제거하는 구형 잠재 공간에서 완전히 작동하는 생성 프레임워크입니다. 이는 동일한 샘플링 단계에서 구형 인코더에 비해 추론 비용을 약 85% (약 6.5 배 적은 FLOPs) 감소시킵니다.
2. 분리된 설계: 고정된 사전 훈련된 오토인코더와 별도의 탈노이즈 모델을 활용함으로써, 재구성과 생성이 독립적으로 전문화되도록 하여 공동 최적화의 내재적 목표 충돌을 해결합니다.
3. 효율적인 최적화: 이 방법은 일부 소수 단계 흐름 방법이 요구하는 1 차 근사 목적 (예: 야코비안 - 벡터 곱) 을 피하고, 대신 구조화된 구형 다양체에서 직접적인 잠재 탈노이즈 과정을 학습합니다.

실험 결과

이 방법은 Animal-Faces, Oxford-Flowers, 그리고 ImageNet-1K(256×256) 에서 평가되었습니다.

• 구형 인코더 대비 성능: Animal-Faces 와 Oxford-Flowers 에서 제안된 방법은 2, 4, 6 단계 샘플링 전반에 걸쳐 생성 품질 (FID) 과 계산 비용 (GFLOPs) 모두에서 구형 인코더를 크게 능가했습니다. 예를 들어, Animal-Faces 에서 4 단계 시 FID 는 구형 인코더의 18.23 에서 6.89(본 방법) 로 개선되었고, GFLOPs 는 4554 에서 390으로 감소했습니다.
• ImageNet-1K 결과: 4 단계 샘플링 (가이드 포함 시 8 NFE) 에서 이 방법은 FID 2.25를 달성하여 동일한 예산 하의 구형 인코더 (4.02) 보다 크게 개선되었습니다. 또한 최근의 여러 소수 단계 베이스라인 (예: MeanFlow, $\alpha$-Flow) 을 능가했으며, 250 단계 이상의 단계를 요구하는 다단계 확산 모델 (예: DiT-XL/2) 의 성능에 근접하면서도 훨씬 적은 함수 평가를 필요로 했습니다.
• 애블레이션 연구:
  • 노이즈 분포: 로그 - 정규 샘플링이 균일 샘플링보다 우수한 성능을 보였습니다.
  • 손실 함수: 재구성과 일관성 손실을 결합한 것이 최상의 결과를 낳았으며, 잠재 일관성 손실을 추가하면 성능이 저하되었습니다.
  • 오토인코더: 토크나이저 선택이 결정적이었습니다. RAE [43](DINOv2 특징 및 $d=768$ 사용) 은 FLUX-VAE 와 GAE 보다 크게 우수한 성능을 보였으며, 강력한 의미적 표현과 고차원 잠재 공간의 필요성을 강조했습니다.

의의 및 주장

이 논문은 구형 잠재 인코더가 소수 단계 이미지 생성을 위한 더 간단하고 효율적인 파이프라인을 제공한다고 주장합니다. 반복적인 픽셀 공간 정제에 대한 의존성을 제거함으로써, 이 방법은 다음과 같은 성과를 달성합니다:

• 높은 효율성: 구형 인코더에 비해 FLOPs 와 추론 시간이 극적으로 감소합니다.
• 더 나은 품질: 재구성과 생성의 분리로 인해 탈노이즈 모델이 재구성 충실도를 희생하지 않고 생성에만 집중할 수 있어 샘플 품질 (낮은 FID) 이 향상됩니다.
• 안정성: 야코비안 - 벡터 곱과 관련된 불안정성과 공동 최적화의 충돌하는 목표를 피하는 훈련 과정입니다.

저자들은 이 작업을 구형 잠재 공간의 이점을 이전 구현의 계산적 페널티 없이 활용하는 확장 가능하고 효율적인 생성 모델링을 향한 한 걸음으로 위치시킵니다. 그들은 현재 클래스 조건부 생성에 초점을 맞추고 있으며 사전 훈련된 오토인코더의 품질에 의존한다는 한계를 인정하며, 이는 1 단계 샘플링 능력을 제한할 수 있음을 지적합니다.