기술 요약: XYZFlow

문제 정의

확산 모델(Diffusion Models)과 플로우 매칭(Flow Matching)을 이용한 고충실도 이미지 생성은 샘플링 속도와 출력 품질 사이의 근본적인 트레이드오프에 직면해 있습니다. DDPM 및 Score-SDE와 같은 모델은 놀라운 시각적 충실도를 달ach하지만, 샘 sample당 수백 번의 반복적인 신경 함수 평가(Neural Function Evaluations)가 필요하여 실시간 애플리케이션에는 부적합합니다.

효율성에 초점을 맞춘 현재의 접근 방식들은 주로 **증류(Distillation)**에 의존하며, 사전 학습된 모델을 적은 단계의 샘플러(예: Rectified Flows, Consistency Models, Shortcut Models)로 압축하려고 시도합니다. 그러나 이러한 방법들은 다음과 같은 중대한 과제에 직면해 있습니다:

1. 모델 의존성: 이들의 성공은 교사 모델(Teacher Model)의 품질에 크게 좌우됩니다.
2. 최적화 복잡성: 복잡한 분포를 결정론적 함수로 증류하는 것은 매우 어려운 최적화 문제입니다.
3. 궤적 모호성(Trajectory Ambiguity): 전통적인 디노이징은 특히 전체 이미지를 단 한 단계로 디노이징하거나 충분한 제약 조건 없이 시도할 때, 겹치거나 모호한 확률 경로(Probability Paths) 문제를 겪습니다.

본 논문은 핵심 과제가 단순히 증류 알고리즘을 개선하는 것이 아니라, 모델 아키텍처를 통해 확률 경로를 본질적으로 더 유니크하고 학습 가능하게 만듦으로써 엄격한 단계 제한 하에서 생성 모델의 **표현력(Expressive Power)**을 확장하는 것이라고 주장합니다.

방법론: XYZFlow

저자들은 XYZFlow를 제안하며, 이는 확률 흐름(Probability Flow)을 시간적(Temporal) 및 공간적(Spatial)이라는 두 개의 직교하는 차원을 따라 스케일링함으로써 효율적인 생성을 재고하는 프레임워크입니다. 핵심 철학은 "집약적 스케일링(Intensive Scaling)"으로, 단순히 파라미터를 추가하거나 증류 단계를 늘리는 것이 아니라 구조화된 다차원 조건화를 통해 확률 흐름의 표현력을 강화하는 것입니다.

1. 이론적 토대: 흐름 강화를 위한 자기회귀 모델링(Autoregressive Modeling)

본 논문은 자기회귀 모델링을 단순한 순차적 예측 전략이 아니라, **흐름 직선화(Flow Straightening)**를 위한 암묵적 메커즘으로 재해석합니다.

• 개념: 자기회귀 모델이 시퀀스를 생성함에 따라, 컨텍스트(이전 상태들)는 점진적으로 더 강력한 제약을 가합니다.
• 효과: 이러한 제약은 조건부 분포 $p(x_p | \cdot)$의 분산을 감소시켜, 노이즈에서 데이터로 가는 확률 흐름 경로를 효과적으로 직선화합니다. 이는 역과정(Reverse Process)을 모호한 일대다(One-to-Many) 매핑에서 거의 결정론적인 일대일(One-to-One) 함수로 변환합니다.

2. 다차원 조건화 (Multidimensional Conditioning)

XYZFlow는 이 이론을 이중 경로 조건화 아키텍처를 통해 구현합니다.

• 시간적 스케일링 (패치 내부 조건화 - Intra-patch Conditioning):

  • 각 디노이징 단계를 현재 상태에만 의존하는 마르코프적(Markovian) 과정으로 취급하는 대신, XYZFlow는 각 단계를 패치의 생성 궤적 전체($H^p_t$)에 조건화합니다.
  • 이를 통해 생성 경로를 안정화하는 시간적 좌표계를 생성하며, 이는 강력한 생성적 사전 지식(Generative Prior) 역할을 합니다.

• 공간적 스케일링 (패치 간 조건화 - Inter-patch Conditioning):

  • 이미지는 패치(예: $2 \times 2$ 그리드)로 나뉩니다.
  • 패치 $x_p$의 생성은 이전 패치들의 최종 출력뿐만 아니라, 그들의 전체 디노이징 궤적($T_{<p}$)에 의해 조건화됩니다.
  • 이전 패치들의 전체 궤적을 사전 지식으로 전달함으로써, 모델은 더 풍부한 컨텍스트 신호를 얻게 되며, 이는 후속 패치들의 솔루션 공간을 더욱 줄여줍니다.

3. 차기 숏컷 예측 (Next Shortcut Prediction)

중심적인 운영 구성 요소는 차기 숏컷 예측으로, 이는 스케일링 패러다임을 자원 집약적 방식에서 제약 집약적 방식으로 전환합니다.

• 점진적 디노이징 스케줄: 모델은 초기 패치는 더 많은 단계를 거치고, 후기 패치는 더 적은 단계를 거치는 비균등 디노이징 스케줄(예: $5 \to 4 \to 3 \to 2$ 단계)을 채택합니다.
• 메커니즘: 후기 패치는 초기 패치들이 구축한 컨텍스트와 "직선화된" 궤적을 활용하여 더 빠르게 수렴합니다.
• 학습 목표: 모델은 교사 궤적 증류(Teacher Trajectory Distillation)를 통해 학습됩니다. 모델은 역사적 궤적 $x_{T(p):t}$와 공간적 컨텍스트 $T_{<p}$가 주어졌을 때 다음 상태 $x_{t-1}$을 예측하는 법을 배웁니다. 목표는 학생 모델의 예측된 흐름과 교사 모델의 실제 조건부 흐름 사이의 편차를 최소화하는 것입니다.
• 추론: 첫 번째 패치는 견고한 앵커를 구축하기 위해 전체 단계 예산을 사용하여 생성됩니다. 이후 패치들은 KV 캐시(Key-Value Cache)를 효율적으로 활용하여 역사적 궤적 정보를 관리하며 자기회귀적으로 생성됩니다.

주요 기여

1. 새로운 스케일링 패러다임: 본 논문은 생성 모델을 확장할 때 단순한 모델 규모 확대("폭 스케일링")나 증류 단계 추가에 의존하는 대신, 구조화된 다차원 조건화를 통해 확률 흐름의 표현력을 높이는 것을 옹호하며, 제약 차원의 스케일링이 매핑의 유일성을 확보하는 원칙적인 경로임을 주장합니다.
2. XYZFlow 프레임워크: 효율적인 추론 시점의 교차 패치 암묵적 궤적 가이딩을 특징으로 하는 시간적 및 공간적 흐름 스케일링의 실질적인 구현체인 Next Shortcut Prediction을 제시합니다.
3. 이론적 정당화: 저자들은 자기회귀 모델링을 명시적인 흐름 강화로 공식화하여, 구조화된 조건화가 역과정의 분산을 제로(Zero)로 유도하여 매핑의 유일성을 보장할 수 있음을 입증했습니다.

실험 결과

실험은 다양한 크기의 교사 모델(Base: 172M, Large: 608M, Huge: 1.1B)을 사용하여 ImageNet 256×256 클래스 조건부 생성 벤치마크에서 수행되었습니다.

• 성능 대비 속도: XYZFlow는 최첨단(SOTA) 효율성-품질 트레이드오프를 달성합니다.
  • **XYZFlow-B (172M 파라미터)**는 0.018초의 추론 시간과 함께 1.63의 FID를 달성합니다.
  • 이는 676M 파라미터 규모의 1-스텝 모델인 *MeanFlow-XL/2+*의 속도(0.018s)와 일치하면서도, 훨씬 더 우수한 FID(1.63 vs 2.20)를 기록했습니다.
  • 또한 820M 파라미터 규모의 DART-FM보다 FID(1.63 vs 3.82)와 속도(0.018s vs 0.32s) 모두에서 앞섭니다.
• 가속도: 표 1에 따르면, XYZFlow는 각자의 교사 모델 대비 상당한 속도 향상을 달립니다. 구체적으로 Base 모델(XYZFlow-B)의 경우 36.1배, Large 모델(XYZFlow-L)은 9.9배, Huge 모델(XYZFlow-H)은 4.0배의 속도 향상을 달성했습니다.
• 강건성: 본 프레임워크는 표준 증류 방식이 크게 저하되는 약한 교사 모델(예: DiT/XL-2)로부터 증류될 때도 강건함을 유지합니다.
• 어블레이션 연구 (Ablation Studies):
  • 전체 이력 가이드 (Full History Guidance): 패치 간 궤적 조건화를 제거하면 FID가 약 0.5 정도 크게 악화되어, 패치의 최종 콘텐츠만 사용하는 것보다 전체 궤적 정보가 우수함을 확인했습니다.
  • 숏컷 예측 (Shortcut Prediction): 점진적 단계 감소($5 \to 4 \to 3 \to 2$) 방식은 일정한 단계 전략과 거의 동일한 품질을 달성하면서도 총 단계 수는 약 30% 줄였습니다.
  • 셀 크기 (Cell Size): $8 \times 8$ 토큰 셀 크기를 사용했을 때 가장 좋은 FID(2.02)를 얻었으며, 이는 이미지를 단일 단위로 처리하는 것보다 패치 수준에서 컨텍스트를 조건화하는 것이 최적임을 시사합니다.

의의 및 주장

본 논문은 XYZFlow가 단순히 모델 규모를 키우는("폭 스케일링") 또는 증류 단계를 추가하는 기존 방식에 대한 유망한 대안을 제공한다고 주장합니다. "딥 스케일링(Deep Scaling)"을 통해 다차원 제약 조건화로 접근함으로써, 이 방법은 다음을 달성합니다:

1. 시간적 조건화를 통해 생성 궤적을 직선화합니다.
2. 공간적 분해를 통해 패치 단위 생성을 단순화합니다.
3. 이 두 차원 사이의 시너지 효과를 활용하여 훨씬 적은 단계와 더 작은 모델 크기로 고충실도 합성을 구현합니다.

저자들은 이러한 이점이 특정 강력한 교사 모델에 대한 의존성이 아니라, 제안된 시공간 아키텍처 자체에서 비롯된다는 점을 강조합니다. 이는 제약의 차원을 확장하는 것이 효율적이고 고충실도인 생성 모델링을 향한 유효한 경로임을 시사합니다.