기술 요약: SGFlow – 비-그래디언트 벡터 플로우를 통한 플로우 맵 학습

1. 문제 정의

확산(Diffusion) 및 플로우 기반 생성 모델은 이미지 및 단백질 생성과 같은 영역에서 상당한 성공을 거두었습니다. 그러나 이러한 모델로부터 샘플링하는 작업은 일반적으로 상미분 방정식(ODE) 또는 확률 미분 방정식(SDE)을 수치적으로 적분해야 하며, 이는 신경망의 여러 차례 순방향 패스(forward pass)를 필요로 합니다. 이는 높은 샘플링 지연 시간과 계산 비용을 초래합니다.

Consistency Models (CMs) 및 Flow Map Matching (FMM)과 같은 최근의 접근 방식들은 확률 흐름 ODE 궤적을 따라 노이즈에서 데이터로 직접 매핑(flow maps)하는 것을 학습함으로써 이러한 적분 비용을 우회하고자 합니다. 그러나 기존 방법들은 다음과 같은 중대한 이론적 및 계산적 장애물에 직면해 있습니다:

• 가역성 제약 (Invertibility Constraints): Lagrangian Flow Map Matching (L-FMM)과 같은 방법은 모델이 명시적으로 가역적(invertible)이어야 하므로, 아키텍처 설계와 학습을 복잡하게 만듭니다.
• 중첩 미분 (Nested Differentiation): 일부 접근 방식은 중첩된 모델 호출(예: 모델의 출력과 그 역함수 또는 후속 단계를 계산)을 통한 역전파(backpropagation)를 요구하며, 이는 계산 비용이 매우 높습니다.
• 이론적 보장의 부재 (Lack of Theoretical Guarantees): 중첩 미분을 피하기 위해 "stop-gradient" (SG) 연산을 사용하는 방식들(예: 특정 변형의 FMM 또는 Meanflow)은 진정한 플로우 맵이 해당 손실 함수의 정지점(stationary point)이라는 증명이 부족합니다. 결과적으로, 최적화가 속도장(velocity field)을 올바르게 적분하지 못하는 가짜 해(spurious solutions)로 수렴할 수 있습니다.

2. 방법론: SGFlow

저자들은 명시적인 가역성 제약과 비용이 많이 드는 모델 반복을 통한 미분을 피하는 플로우 맵 학습법인 SGFlow (StopGrad Flow)를 소개합니다.

핵심 공식화

SGFlow는 시간 $t$에서 $u$까지 속도장 $v$를 적분하는 플로우 맵 $f(t, u, x)$가 다음의 수송 편미분 방정식(transport PDE)을 만족한다는 관찰에 기초합니다:
$$\partial_t f + (\partial_x f)v(t, x) = 0$$
여기서 경계 조건은 $f(u, u, x) = x$입니다.

저자들은 경계 조건을 자동으로 만족하는 파라미터화 $f_\theta(t, u, x) = x + (u - t)\tilde{f}_\theta(t, u, x)$를 제안합니다. 목표는 PDE 잔차(residual)를 최소화하는 $\tilde{f}_\theta$를 찾는 것입니다.

SGFlow 손실 함수

직접적으로 PDE의 제곱 잔차를 최소화하는 것은 실제 속도 $v(t, x) = \mathbb{E}[\dot{X}_t | X_t]$에 대한 지식을 필요로 하지만, 이는 알려져 있지 않습니다. 표준 Flow Matching은 샘플 수준의 속도 $\dot{X}_t$를 $v$ 대신 사용하지만, 이는 비선형 항(특히 자코비안 $\partial_x f$가 포함된 경우)이 포함될 때 편향(bias)을 유발합니다.

SGFlow는 알려지지 않은 $v$를 현재 타임 스텝에서의 모델 자체 속도 추정치의 stop-gradient 근사치로 대체하면서도, 올바른 그래디언트 흐름을 보장하는 손실 함수 $L_{sg}$를 구성합니다:
$$L_{sg} = \mathbb{E}_{X_t} \left[ \| \partial_t f_\theta + (\partial_x f_\theta)\dot{X}_t \|^2 - \| (\partial_x f_\theta)(\dot{X}_t - \text{SG}[\tilde{f}_\theta(t, t, \cdot)]) \|^2 \right]$$
여기서 $\text{SG}[\cdot]$는 stop-gradient 연산자를 나타냅니다. 첫 번째 항은 모델이 샘플 속도를 사용하여 수송 방정식을 만족하도록 유도합니다. 두 번째 항은 보정 역할을 합니다. 즉, 자코비안이 포함된 비선형 항에서 $\dot{X}_t$를 $v$로 대체함으로써 발생하는 편향을 보상합니다.

비-그래디언트 역학 (Non-Gradient Dynamics)

이 논문의 핵심적인 이론적 통찰은 $L_{sg}$에 의해 유도되는 최적화 역학이 **비-보존적(non-conservative)**이라는 점입니다. stop-gradient 연산자가 업데이트를 단일 스칼라 목적 함수의 그래디언트로 표현하는 데 필요한 대칭성을 깨뜨리기 때문에, 시스템은 "비-그래디언트 벡터 플로우(non-gradient vector flow)"를 따릅니다. 이 체제에서 목표는 스칼라 손실을 최소화하는 것이 아니라, 벡터장의 **정지점(stationary point)**을 찾는 것입니다.

3. 주요 기여

1. 정지점에 대한 이론적 보장: 본 논문은 진정한 플로우 맵이 SGFlow 목적 함수의 유일한 정지점임을 증명합니다 (Theorem 1). 이는 특정 보정이 없는 Meanflow나 증명되지 않은 FMM 변형과 같은 이전의 stop-gradient 방식들에 비해 중요한 개선입니다. 이들은 진정한 플로우 맵이 정지점이 될 것이라는 보장이 없었기에, 잠재적으로 하위 최적의 해로 수렴할 수 있었습니다.
2. 역함수 불필요 학습 (Inverse-Free Training): Flow Map Matching과 달리, SGFlow는 모델이 가역적일 필요가 없거나 학습 중에 모델의 역함수를 계산할 필요가 없습니다. 이 방식은 오직 Jacobian-Vector Product (JVP)에만 의존하며, 이는 전체 자코비안 행렬을 구현하지 않고도 순방향 자동 미분을 통해 효율적으로 계산할 수 있습니다.
3. 중첩 미분 배제: 이 방법은 일관성 증류(consistency distillation) 및 자기 증류(self-distillation) 방법에서 흔히 발생하는 병목 현상인 중첩된 모델 호출(예: $f(f(x))$)을 통한 역전파를 피합니다.
4. 유연한 샘플링: SGFlow는 ODE 궤적을 따라 임의의 단계로 생성을 가능하게 하여, 1-step 방식과 다단계(many-step) 방식 사이의 간극을 메웁니다.

4. 실험 결과

저자들은 CIFAR-10 무조건부 이미지 생성 벤치마크에서 SGFlow를 Flow Matching, Meanflow, 그리고 (stop-gradient 배치가 적용된) Lagrangian map matching과 비교 평가했습니다.

• 성능: 단일 방법이 모든 샘플링 단계 수에서 압도적으로 우세하지는 않았습니다.
  • Meanflow는 1-step에서 가장 좋은 성능을 보였습니다.
  • Lagrangian map matching은 50 및 100단계에서 가장 좋은 성능을 보였습니다.
  • SGFlow는 **10단계 샘플링 시 최상의 Fréchet Inception Distance (FID) (FID 4.00)**를 달성했으며, 다른 단계 수에서도 경쟁력 있는 성능을 유지했습니다.
• 이론 vs 실증: SGFlow는 비교 대상 중 stop-gradient 기반 역학에 대해 증명된 정지점 보장을 제공하는 유일한 방법입니다. 결과는 이러한 이론적 엄밀함이 실증적 성능의 저하를 초래하지 않음을 보여줍니다. 실제로 SGFlow는 중간 단계 수(10단계)에서 최첨단(SOTA) 결과를 달성했습니다.

5. 의의 및 주장

본 논문은 SGFlow가 다음과 같은 특성을 제공함으로써 1-step 및 다단계 생성 모델링 사이의 "설계 공간(design space)" 문제를 해결한다고 주장합니다:

• 계산 효율성: 비용이 많이 드는 역함수 계산과 중첩 미분을 피합니다.
• 이론적 타당성: 진정한 플로우 맵이 유일한 정지점이라는 엄격한 증명을 제공하며, 이는 여러 현대적 stop-gradient 접근 방식에서 결여된 보장입니다.
• 실용적 효과: 표준 벤치마크, 특히 중간 정도의 샘플링 단계(예: 10단계)에서 경쟁력 있거나 더 우수한 생성 품질(FID)을 달성합니다.

저자들은 SGFlow에서 stop-gradient의 사용이 단순한 계산적 트릭이 아니라 정확성을 위한 필수 요소임을 강조합니다. 특정 stop-gradient 배치와 보정 손실 항이 없다면, 가우시안 반례(counterexample)에서 보여준 것처럼 최적화 역학이 진정한 플로우 맵에서 벗어나게 됩니다(Meanflow는 stop-gradient 없이 가짜 끌개(spurious attractor)로 수렴함).

이 연구는 SGFlow를 가역성, 계산 비용, 이론적 최적성 사이의 트레이드오프를 균형 있게 조절하며 플로우 맵을 학습하는 강력한 대안으로 포지셔닝하며, 전통적인 ODE 적분의 오버헤드 없이 고품질의 효율적인 생성 모델링을 위한 길을 열어줍니다.