기술 요약: 공간적 접지 플로우 매칭 (StructFlow)

문제 정의

현재 이미지 생성을 위한 플로우 매칭 모델은 일반적으로 단순하고 등방성인 가우시안 소스 분포($p_0 = \mathcal{N}(0, I)$)로부터 복잡한 자연 이미지의 데이터 분포로 샘플을 운송하는 법을 학습합니다. 효과적이긴 하지만, 이 표준 정식화는 근본적인 불일치 문제를 겪습니다. 자연 이미지는 강한 공간적 국소성(인접한 픽셀들이 서로 밀접하게 연관되어 일관된 영역을 형성함)을 보이는 반면, i.i.d. 가우시안 사전 분포는 공간적 위치 전반에 걸친 완전한 독립성을 가정하기 때문입니다.

저자들은 이러한 불일치가 모델로 하여금 속도장(velocity field)을 통해 공간적 국소성을 명시적으로 재구성하도록 강요하며, 이는 귀납적 편향(inductive bias)으로 상속받는 대신 모델이 이를 학습하게 만든다고 가설을 세웠습니다. 이는 모델에 불필요한 부담을 주며, 잠재적으로 학습 효율을 저해하고 제어 가능성을 제한할 수 있습니다. 즉, 모델은 로컬 구조를 처음부터 "학습"해야 하며, 이는 소스 분포가 이미지 도메인의 내재적 귀납적 편향과 상충되는 상황입니다.

방법론: StructFlow

이를 해결하기 위해 저자들은 표준적인 등방성 가우시안 소스를 대신하여, 공간적 국소성을 명시적으로 인코딩하는 구조화된 소스 분포를 사용하는 StructFlow 프레임워크를 제안합니다.

핵심 메커니즘: 계층적 노이즈 샘플링

모든 픽셀에 대해 독립적으로 노이즈를 샘플링하는 대신, StructFlow는 작은 공간 영역(세그먼트) 내의 픽셀들이 공통된 노이즈 성분을 공유하는 계층적 가우시안 사전 분포를 구축합니다.

1. 세그멘테이션(Segmentation): SLIC 세그멘테이션(또는 추론 시의 대리 모델)을 사용하여 이미지 공간을 $K$개의 영역(슈퍼픽셀)으로 분해합니다.
2. 계층적 샘플링:
   • 레벨 1 (앵커): 각 세그먼트 $k$에 대해 공유 앵커 벡터 $z_k$를 샘플링합니다.
   • 레벨 2 (로컬 노이즈): 세그먼트 $k$에 속하는 각 픽셀 $(i, j)$에 대해, 앵커를 중심으로 하는 가우시안 분포에서 노이즈 $\epsilon_{i,j}$를 샘플링합니다: $\epsilon_{i,j} \sim \mathcal{N}(z_k, \lambda_\varsigma^2 I)$.
   • 재스케일링(Rescaling): 단위 주변 분산을 보장하기 위해 노이즈를 전역적으로 재스케일링합니다.
3. 상관관계 제어: 파라미터 $\lambda_\varsigma$는 국소성의 강도를 제어합니다. $\lambda_\varsigma \to 0$이 될수록 동일 세그먼트 내의 픽셀들은 매우 높은 상관관계를 가지며, $\lambda_\varsigma \to \infty$가 될수록 분포는 표준 i.i.d. 가우시안 노이즈에 가까워집니다.

안정화를 위한 학습 전략

구조화된 노이즈를 사용한 학습은 증가된 그래디언트 분산으로 인해 최적화 불안정성을 초래할 수 있습니다. 저자들은 이를 안정화하기 위해 세 가지 상호 보완적인 전략을 제안합니다.

1. 코사인 유도 공간 앵커링 (CISA): 앵커 $z_k$를 무작위로 샘플링하는 대신, 공간 격자 위의 결정론적인 대각 코사인 그리드(diagonal cosine grid)로부터 유도합니다. 이는 확률적 변동성을 줄이고 일관된 공간적 단서를 제공하여 속도장의 학습을 안정화합니다.
2. 점진적 일관성 어닐링 (Progressive Coherence Annealing): 모델을 커리큘럼 방식으로 학습시키는데, $\lambda_\varsigma$가 큰 값(near i.i.d. Gaussian)에서 시작하여 목표 값까지 선형적으로 감소합니다. 이를 통해 모델이 쉬운 근사 가우시안 소스로부터 먼저 학습한 후, 점진적으로 구조적 제약에 적응할 수 있게 합니다. 결정적으로, 이는 큰 $\lambda_\varsigma$에서 시작하여 어닐링을 통해 기존의 사전 학습된 플로우 매칭 모델을 사후 학습(post-training) 할 수 있게 합니다.
3. 혼합 일관성 학습 (Mixed-Coherence Training): 생성 품질(높은 $\lambda_\varsigma$)과 로컬 편집 가능성(낮은 $\lambda_\varsigma$) 사이의 이분법적 트레이드오프를 피하기 위해, 학습 종료 시점에 $\lambda_\varsigma$를 $\lambda_{min}$ 근처의 범위 $[\lambda_{min} - \delta, \lambda_{min} + \delta]$에서 균등하게 샘플링하여 학습합니다. 이를 통해 단일 모델이 추론 시 다양한 일관성 스펙트럼에서 작동할 수 있게 합니다.

주요 기여

1. 소스 불일치 식별: 본 논문은 표준 i.i.d. 가우시안 소스 분포의 공간적 구조 결여가 플로우 매칭의 주요 한계임을 식별하고, 국소성이 데이터로부터 완전히 학습되는 것이 아니라 소스에 구축되어야 한다고 주장합니다.
2. StructFlow 프레임워크: 로컬 영역 내의 노이즈를 상관시키는 공간 구조적 소스 분포를 도입하여 이미지 국소성의 귀납적 편향을 직접 인코딩합니다.
3. 안정화 기술: 구조화된 노이즈와 관련된 최적화 불안정성을 해결하기 위해 CISA, 점진적 어닐링, 혼합 일관성 학습을 개발했습니다.
4. 창발적 능력: 이 간단한 수정이 생성 품질을 희생하지 않으면서도 미세한 로컬 편집, 구조 보존, 매끄러운 시맨틱 보간(semantic interpolation)과 같은 새로운 속성을 어떻게 해제하는지 입증했습니다.

결과 및 평가

저자들은 ImageNet 및 FFHQ 데이터셋에 대해 DiT 아키텍처를 사용하여 무조건적(unconditional), 클래스 조건부(class-conditional), 텍스트 조건부(text-conditioned) 환경에서 StructFlow를 평가합니다.

• 생성 품질: StructFlow는 FID, sFID, Precision, Recall, ImageReward와 같은 지표에서 표준 플로우 매칭과 대등하거나 이를 능가하는 경쟁력 있는 이미지 생성 품질을 달성합니다.
• 미세 편집 (Fine-Grained Editing): 선택된 세그먼트 내의 노이즈만 재샘플링하고 나머지는 고정함으로써, StructFlow는 정밀한 로컬 재합성(re-synthesis)을 가능하게 합니다. 정량적 분석 결과, 낮은 $\lambda_\varsigma$로 학습된 모델은 편집되지 않은 영역을 현저히 더 잘 보존(낮은 LPIPS/RMSE "블리드")하면서도 편집된 영역에서의 다양성을 유지함을 보여줍니다.
• 구조 보존: 생성된 이미지는 제공된 세그멘테이션 마스크를 자연스럽게 따르므로, 사용자가 질감과 디테일을 변화시키면서 글로벌 레이아웃을 가이드할 수 있게 합니다.
• 초기 특징 출현: 중간 표현(intermediate representation) 분석을 통해, StructFlow 모델이 표준 모델에 비해 디노이징 과정의 훨씬 초기 단계에서 공간적으로 일관된 특징(예: 객체 경계)을 발달시킨다는 것을 밝혀냈습니다.
• 사후 학습 적용 가능성: 이 방법은 코herence 어닐링을 통한 경량 사후 학습을 통해 기존 사전 학습된 모델(예: SANA 텍스트-이미지 모델)에 성공적으로 통합될 수 있으며, 이는 그 이점이 파운데이션 모델까지 확장됨을 입증합니다.

의의 및 주장

본 논문은 소스 분포 설계가 생성 모델을 개선하기 위한 강력하고 아직 충분히 탐구되지 않은 축임을 주장합니다. 소스 분포를 타겟 도메인의 귀납적 편향(공간적 국소성)과 정렬함으로써, StructFlow는 다음과 같은 효과를 거둡니다:

• 로컬 구조를 재구성해야 하는 학습된 속도장의 부담을 줄입니다.
• 표준 i.i.d. 사전 분포로는 달성하기 어려운 제어 가능성과 해석 가능성(예: 국소적 편집, 구조 보존)을 확보합니다.
• 이러한 이점들이 경량 사후 학습을 통해 기존 파운데이션 모델에 통합될 수 있음을 보여줌으로써, 접근 방식의 실용성을 증명합니다.

저자들은 이러한 기능들이 특수한 편집 아키텍처나 복잡한 인버전(inversion) 기술을 요구하는 것이 아니라, 구조화된 소스로부터 자연스럽게 창발된다는 점을 강조합니다.