AI바라기의 인공지능

Diffusion : 논문 리뷰 : Continuous Adversarial Flow Models 본문

논문리뷰

Diffusion : 논문 리뷰 : Continuous Adversarial Flow Models

AI바라기 2026. 4. 19. 19:39

용어 설명 (Glossary)

  • Flow Matching (FM): 노이즈 분포에서 실제 데이터 분포로 이동하는 확률 흐름(probability flow)의 속도장(velocity field)을 예측하도록 학습하는 continuous-time 생성 모델링 기법.
  • Continuous Adversarial Flow Models (CAFMs): 본 논문에서 제안하는 핵심 모델. 기존 FM에 학습 가능한 Discriminator를 도입하여 adversarial objective로 학습하는 continuous-time 모델.
  • Jacobian-Vector Product (JVP): Discriminator가 모델이 예측한 속도(velocity)를 단순한 숫자가 아닌 '위치에 대한 미분값'이라는 물리적 속성으로 올바르게 인식하도록 만들기 위해 사용하는 수학적 연산.
  • Post-training: 이미 대규모 데이터로 학습이 완료된 모델(여기서는 기존 FM 모델)에 새로운 objective(CAFM)를 적용하여 짧은 시간 동안 추가 학습(fine-tuning)시켜 성능을 극대화하는 기법.
  • Out-of-distribution (OOD) samples: 모델이 학습 데이터의 실제 manifold를 제대로 파악하지 못해 생성해내는, 현실에 존재하지 않을 법한 어색하거나 퀄리티가 낮은 이미지들.
  • vmap (Vectorizing map): 여러 입력 데이터에 대해 동일한 함수(여기서는 JVP)를 병렬로 빠르게 연산하게 해주는 최적화 기능.

Purpose of the Paper

  • 기존 문제점: Flow Matching (FM)은 Generator가 예측한 속도와 실제 타겟 속도 간의 고정된 Euclidean distance (L2 loss) 를 최소화하도록 학습함. 하지만 신경망의 용량(capacity)은 유한하므로, 실제 데이터의 복잡한 manifold를 무시하고 등방성(isotropic)으로 잘못된 일반화(generalization)를 수행하여 Guidance 없이는 OOD 샘플을 생성하는 치명적인 한계가 있음.
  • 기존 연구의 한계: 이를 극복하기 위해 제안된 discrete-time 기반의 Adversarial Flow Models (AFMs)는 time step이 0에 가까워질수록 학습이 극도로 불안정해지고 gradient vanishing 문제가 발생함.
  • 논문의 목표: GAN의 장점(Discriminator가 데이터의 manifold와 perceptual detail을 잘 포착함)을 continuous-time flow modeling에 안정적으로 통합하는 것. 이를 통해 고정된 L2 loss의 편향을 없애고, 타겟 데이터 분포에 완벽하게 정렬되는 새로운 Continuous Adversarial Flow Models (CAFMs) 를 제안함.

Key Contributions

  • Continuous-time Adversarial Training 설계 (Novelty)
    • 기존 discrete-time에 머물렀던 AFMs의 한계를 깨고, continuous-time flow modeling에 adversarial training을 통합한 최초의 연구. Time step 분할에 따른 불안정성을 근본적으로 해결함.
  • JVP 기반의 혁신적인 Discriminator 구조 (Novelty)
    • 단순히 예측된 속도 vt를 Discriminator의 입력으로 바로 넣는 Naive 방식을 폐기함. 대신, 속도를 위치 xt의 미분(derivative)으로 취급하여 Jacobian-Vector Product (JVP) 공간에서 판별을 수행함.
    • 이 참신한 접근법은 Generator가 엉뚱한 타겟으로 수렴하며 진동하는 현상을 막고, 물리적으로 타당한 판별 공간을 구축하여 유의미한 gradient를 제공함.
  • 안정적인 학습 최적화 및 Gradient Penalty 제거
    • Continuous-time으로 확장하면서 gradient vanishing 문제가 수학적으로 완화됨을 증명하여, 기존 adversarial 학습의 필수품이었던 gradient penalty를 완전히 제거함.
    • Discriminator 내부의 LayerNorm을 RMSNorm으로 교체하여 JVP 연산 시 발생하는 gradient norm spike(기울기 폭발) 현상을 억제하고 훈련 안정성을 대폭 끌어올림.
  • 초고효율 Post-training 패러다임 입증
    • FM과 CAFM은 도달하고자 하는 ground-truth flow가 동일하므로, 처음부터 무겁게 학습할 필요 없이 이미 잘 학습된 FM 모델을 CAFM으로 Post-training 하는 것이 극도로 효율적임을 실험적으로 증명함.

Experimental Highlights

  • ImageNet 256px Generation Post-training (단 10 Epochs 추가 학습)
    • Latent-space SiT-XL/2 모델:
      • Guidance-free FID: 8.26 -> 3.63 (절반 이하로 대폭 감소, 압도적 성능 향상)
      • Guided FID: 2.06 -> 1.53 (SOTA급 성능 달성)
    • Pixel-space JiT-H/16 모델:
      • Guidance-free FID: 7.17 -> 3.57
      • Guided FID: 1.86 -> 1.80
  • Text-to-Image Generation (Z-Image 모델 적용)
    • 자체 데이터셋으로 20k iteration Post-training 수행.
    • GenEval benchmark: 0.81 -> 0.85 상승
    • DPG-Bench: 83.7 -> 85.2 상승
  • 대조군 실험 (Control Trial)의 성공
    • 단순히 학습을 더 해서 성능이 오른 것이 아님을 증명하기 위해, 동일한 조건에서 L2 loss (FM objective)로 추가 학습한 통제군을 두었음. FM 통제군은 성능 향상이 거의 없었으며, 오직 CAFM objective만이 극적인 성능 향상을 이끌어냄을 완벽히 검증함.

Limitations and Future Work

  • Limitation 1: 완벽한 분포 일반화 보장 불가
    • CAFM이 모델의 일반화 능력을 크게 개선하여 Guidance-free 성능을 비약적으로 높였으나, 데이터 밀도가 낮은 영역(outliers)에서는 여전히 잘못된 OOD 이미지가 생성될 수 있음. 즉, Classifier-free guidance (CFG) 기술을 완전히 대체하기보다는 함께 직교(orthogonal)하여 사용할 때 최고의 성능을 냄.
  • Limitation 2: 연산 비용 (Computational Overhead) 증가
    • Discriminator 네트워크 추가 및 JVP forward/backward 연산으로 인해 기존 FM 대비 1 epoch 당 학습 시간이 약 4.8배 증가함. 이 때문에 Pre-training보다는 Post-training 용도로 훨씬 적합함.
  • Future Work: Pre-training 스케줄링 최적화
    • 현재 CAFM으로 처음부터 학습(scratch)하면 FM보다 수렴 속도가 느린 문제가 있음. 향후 학습 단계에 맞춰 Discriminator 업데이트 횟수(N)와 optimal transport loss weight를 동적으로 조절하는 스케줄링 기법을 고도화하여 Pre-training 효율을 개선하는 연구가 필요함.

Overall Summary

이 논문은 고정된 L2 loss의 한계로 인해 Out-of-distribution 이미지를 생성하는 Flow Matching 모델의 고질적인 문제를 해결하기 위해, **Continuous Adversarial Flow Models (CAFMs)**를 제안했습니다. 핵심은 JVP 연산을 통해 물리적으로 올바른 Discriminator를 continuous-time에 통합한 것이며, 복잡한 페널티 없이도 학습 안정성을 확보했습니다. 특히, 기존에 훈련된 SOTA급 FM 모델(SiT, JiT)에 단 10 epoch의 Post-training만 적용해도 Guidance 없는 생성 퀄리티(FID)를 두 배 이상 끌어올리는 놀라운 결과를 증명하여, 향후 생성 모델의 충실도(fidelity)를 극대화하는 강력하고 필수적인 Fine-tuning 표준 방법론으로 자리 잡을 잠재력이 매우 큽니다.


쉬운 설명

기존의 Flow Matching(FM) 모델은 노이즈에서 이미지로 가는 길을 찾을 때 단순히 "직선 거리(L2 loss)"만 보고 그림 그리는 법을 배웁니다. 그러다 보니 실제 사진들이 모여있는 복잡한 굴곡(Manifold)을 무시하게 되어, 가이드라인(Guidance)을 주지 않으면 현실에 없는 어색한 그림(OOD)을 맘대로 그려버립니다.

이 논문은 **"미술 선생님(Discriminator)"을 새로 고용하는 방법(CAFM)**을 제안합니다. 이 선생님은 단순히 완성된 그림의 차이를 보는 게 아니라, 그림이 그려지는 **"붓터치의 방향과 흐름(JVP를 활용한 속도의 미분값)"**을 매 순간 감시합니다. "방금 그 붓터치 방향은 실제 사진이 만들어지는 흐름이 아니야!"라고 계속 혼내면서(Adversarial training) 방향을 교정해 줍니다.

가장 대단한 점은, 이미 직선 거리만 보고 대충 그림 그리는 법을 배운 졸업생(기존 FM 모델)이라도, 이 깐깐한 선생님에게 **아주 짧은 단기 특강(Post-training 10 epochs)**만 받으면 가이드라인 없이도 혼자서 기가 막히게 진짜 같은 그림을 그려내게 된다는 것입니다.

 

 

 

 

 

 

 

Abstract

우리는 adversarial objective로 trained 된 continuous-time flow model의 한 종류인 continuous adversarial flow models를 제안합니다.

고정된 mean-squared-error criterion을 사용하는 flow matching과 달리, 우리의 접근 방식은 training을 가이드하기 위해 learned discriminator를 도입합니다. 이러한 objective의 변화는 다른 generalized distribution을 유도하며, 이는 경험적으로 target data distribution과 더 잘 정렬된 samples를 생성합니다.

우리의 방법은 주로 기존의 flow-matching models를 post-training 하기 위해 제안되었지만, 처음부터 models를 train 할 수도 있습니다.

ImageNet 256px generation task에서 우리의 post-training은 latent-space SiT의 guidance-free FID를 8.26에서 3.63으로, pixel-space JiT의 guidance-free FID를 7.17에서 3.57로 상당히 향상시킵니다. 또한 guided generation을 향상시켜, SiT의 FID를 2.06에서 1.53으로, JiT의 FID를 1.86에서 1.80으로 감소시킵니다.

우리는 나아가 text-to-image generation에서 우리의 접근 방식을 평가하며, 여기서 GenEval 및 DPG benchmarks 모두에서 향상된 결과를 달성합니다.

Keywords: Generative models · Adversarial training · Flow models

 

 

 

 

1 Introduction

더보기

Flow matching은 최근 몇 년간 상당한 성공을 거두었지만, 여전히 중요한 문제가 남아있습니다. 이 문제는 image 및 video synthesis와 같은 visual modalities의 generation에서 특히 두드러지며, guidance가 적용되지 않으면 models가 종종 out-of-distribution samples를 생성합니다. guidance는 sample 품질을 향상시키지만, sampling distribution을 변경합니다. original data의 underlying distribution을 어떻게 더 충실하게 model할 것인가는 여전히 open problem으로 남아있습니다.

flow matching이 out-of-distribution samples를 생성하는 한 가지 이유는 manifold-aware criterion 대신 Euclidean distance criterion을 사용하기 때문입니다. 구체적으로, flow matching(FM)은 noise와 data distributions 사이의 probability flow의 velocity field를 learn합니다. Training은 predicted velocities와 target velocities 사이의 squared $L_2$ loss를 최소화합니다. 이론적으로, 이 objective는 infinite-capacity model 하에서 ground-truth flow로 수렴하며, 이는 overfit되어 training samples를 정확히 재현할 것입니다. 실제로는, finite capacity를 가진 models는 generalize해야 하며, 따라서 새로운 data samples의 generation이라는 결과를 낳습니다. 그러나 squared $L_2$ objective는 manifold-aware distance가 아닌 Euclidean distance를 측정하여, underlying data distribution에 비해 부정확한 generalization을 유도합니다.

최근의 연구들은 다양한 각도에서 이 문제를 해결하려고 시도했습니다. Representational autoencoders는 flow matching이 작동하는 data space를 변환하고 경험적으로 generation 품질의 향상을 보고했지만, 이는 original data space 대신 latent space에서 작동해야 합니다. Riemannian flow matching은 flow matching을 non-Euclidean geometries로 확장하지만, 이는 data manifold의 수동 정의를 필요로 하며, 일반적인 datasets에 대해서는 종종 알려져 있지 않습니다. 다른 연구는 deep networks가 더 나은 perceptual metrics 역할을 할 수 있다는 경험적 발견에 동기를 부여받아, Euclidean loss를 frozen feature networks에서 파생된 perceptual distances로 대체합니다. 그러나 fixed criterion network는 generator에 의해 악용될 수 있으며, 이는 generated samples에 artifacts를 초래합니다. generator 해킹을 완화하는 한 가지 방법은 criterion network를 generator와 함께 jointly train하는 것이며, 이는 generative adversarial networks를 연상시키는 dynamic을 산출합니다.

Generative adversarial networks(GANs)는 generative 방법론의 독립적인 class입니다. 이들은 ImageNet benchmarks에서 강력한 성능을 달성하며 선명한 image synthesis를 위한 flow-model distillation에 널리 사용됩니다. 우리는 discriminator networks가 pointwise Euclidean losses보다 texture, sharpness, contour 등과 같은 perceptual details에 더 민감하기 때문에 이러한 이점이 발생한다고 hypothesize합니다. 왜냐하면 그들이 manifold structure를 더 잘 포착하도록 learn했을 수 있기 때문입니다. 최근 연구인 adversarial flow models(AFMs)는 adversarial과 flow modeling을 결합하여, training stability를 향상시키고 adversarial objectives를 multi-step flow training으로 확장합니다. 그러나 AFMs는 discrete time으로 공식화되어 있어, adversarial training을 continuous-time flow modeling에 어떻게 통합할 것인지에 대한 question을 열어두고 있습니다.

본 논문에서는 AFMs를 continuous time으로 확장하는 continuous adversarial flow models(CAFMs)를 소개합니다. CAFMs는 noise에서 data로 ordinary differential equation(ODE)를 통합하여 samples를 generate하는 continuous normalizing flow(CNF)의 한 유형입니다. flow-matching models(FMs)와 마찬가지로, CAFMs 역시 simulation-free objective를 사용하여 predefined probability flow의 velocity field를 learn합니다. FMs와 CAFMs는 동일한 ground-truth flow를 target으로 하지만, CAFMs는 fixed Euclidean criterion 대신 learned discriminator를 사용하기 때문에 finite-capacity generalization에서 차이가 있습니다. 경험적으로, 우리의 실험은 CAFMs가 perceptually뿐만 아니라 다양한 metrics에 의해서도 더 많은 in-distribution samples를 produce한다는 것을 발견했습니다. 우리가 아는 한, 우리의 연구는 continuous-time flow modeling에 adversarial training을 적용한 최초의 연구입니다.

FMs와 CAFMs는 동일한 ground-truth flow를 learn하고 주로 model generalization에서 차이가 나기 때문에, 비록 그 objective가 처음부터 training하는 데 사용될 수도 있지만 우리의 방법은 효율성과 실용성을 위해 기존의 FMs를 post-train하도록 주로 설계되었습니다. class-conditional ImageNet 256px generation에서 CAFM post-training은 단 10 epochs의 finetuning만 사용하여 latent-space SiT의 guidance-free FID를 8.26에서 3.63으로, pixel-space JiT의 경우 7.17에서 3.57로 향상시킵니다. CAFMs는 또한 더 나은 guided generation을 달성하여, SiT의 FID를 2.06에서 1.53으로, JiT의 FID를 1.86에서 1.80으로 향상시킵니다. text-to-image 실험에서 CAFMs는 GenEval score를 0.81에서 0.85로, DPG score를 83.7에서 85.2로 증가시킵니다. 이러한 결과는 few-step generation이 아니라 sample fidelity 및 distribution matching을 향상시키기 위해 continuous-time flow modeling에 adversarial training을 통합하는 것에 대한 유망한 전망을 제시합니다.

 

1 Introduction 핵심 정리 노트

전문적인 AI 연구자들을 위해, 일반적인 배경 설명은 최소화하고 본 논문이 제시하는 고유한 문제의식과 해결 방식, 그리고 그 결과에 집중하여 요약했습니다.

  • 기존 Flow Matching의 근본적 한계 (Problem Formulation)
    • Guidance가 없을 시 out-of-distribution sample이 생성되는 고질적 문제가 존재함.
    • 이는 flow matching(FM)이 manifold-aware distance가 아닌 Euclidean distance 기준(squared $L_2$ loss)을 사용하기 때문임.
    • 무한한 용량의 모델에서는 ground-truth flow로 수렴하겠지만, 실제 finite-capacity model은 일반화(generalization) 과정에서 underlying data distribution을 정확히 반영하지 못하는 부정확한 generalization을 겪게 됨.
  • 기존 접근법들의 한계 (Limitations of Prior Works)
    • Representational autoencoders: Original data space가 아닌 latent space에서 작동해야 하는 제약이 있음.
    • Riemannian flow matching: 일반적인 datasets에서는 알기 힘든 data manifold를 수동으로 정의해야 함.
    • Perceptual loss (Fixed feature networks 활용): Generator에 의한 악용(hacking)에 취약하여 artifact를 유발함. 이를 막기 위해 discriminator를 함께 학습시키는 GAN 방식의 동인이 생김.
  • 본 논문의 핵심 제안: Continuous Adversarial Flow Models (CAFMs)
    • 기존에 discrete time으로만 존재하던 Adversarial Flow Models(AFMs)를 continuous time으로 확장한 최초의 연구임.
    • CAFMODE를 적분하여 샘플을 생성하는 continuous normalizing flow(CNF)의 일종임.
    • 고정된 Euclidean criterion 대신 learned discriminator를 사용하여 training을 가이드함으로써, 모델이 질감이나 윤곽 같은 perceptual detailsmanifold structure를 더 잘 포착하도록 유도함.
    • FMCAFM은 동일한 ground-truth flow를 목표로 하지만, discriminator의 도입으로 인해 finite-capacity generalization 방식에서 결정적인 차이를 만들어냄.
  • 효율적인 활용 방안 및 주요 성과 (Efficiency & Key Results)
    • Train from scratch도 가능하지만, 동일한 ground-truth flow를 타겟으로 한다는 점을 활용하여 기존에 학습된 FMpost-training 하는 데 최적화되어 설계됨 (매우 효율적).
    • ImageNet 256px 생성에서 단 10 epochsfinetuning만으로 극적인 성능 향상 달성:
      • Latent-space SiT: Guidance-free FID 8.26 $\rightarrow$ 3.63 (Guided: 2.06 $\rightarrow$ 1.53)
      • Pixel-space JiT: Guidance-free FID 7.17 $\rightarrow$ 3.57 (Guided: 1.86 $\rightarrow$ 1.80)
    • Text-to-image 모델에서도 GenEvalDPG score를 모두 향상시킴.
    • 즉, 본 방법론은 few-step generation 속도 개선보다는 sample fidelitydistribution matching을 정교하게 다듬는 데 강력한 효과를 발휘함.

쉬운 설명 :

Flow Matching(FM) 모델은 데이터를 생성할 때 시작점(노이즈)과 도착점(실제 데이터) 사이의 경로를 학습합니다. 하지만 기존의 방식은 점과 점 사이의 '직선 거리(유클리디안 거리)'만을 기준으로 정답을 맞히려고 했습니다. 실제 데이터들이 모여있는 복잡하고 구불구불한 지형(Manifold)을 무시하고 직선으로만 가려고 하니, 용량이 제한된 실제 AI 모델들은 길을 잃고 이상한 데이터(out-of-distribution)를 만들어내곤 했습니다.

이전에도 이 경로 이탈을 막으려는 시도들이 있었지만, 데이터를 변환해야 하거나(Latent space), 지형을 인간이 직접 입력해 줘야 하거나, 고정된 채점자를 뒀다가 AI가 꼼수(Generator hacking)를 부려 이미지가 깨지는 부작용이 있었습니다.

그래서 연구진은 GAN 모델에서 널리 쓰이는 눈치 빠른 '감별사(Discriminator)'를 도입하기로 합니다. 감별사는 픽셀 단위의 기계적인 거리가 아니라, 이미지의 질감, 선명도 같은 진짜 데이터의 특성(Perceptual details)을 기준으로 채점하기 때문에 AI가 실제 데이터의 지형을 훨씬 더 잘 따라가게 만듭니다.

이 논문의 가장 큰 장점은 이 감별사를 연속적인 시간(Continuous time)에서 작동하는 Flow 모델에 최초로 적용했다는 것(CAFM)이고, 무엇보다 **'가성비'**가 엄청나다는 것입니다. 처음부터 모델을 다시 학습시킬 필요 없이, 기존에 직선거리만 고집하던 바보 같은 FM 모델들을 가져와서 감별사와 함께 딱 10번(10 epochs)만 과외(Post-training)를 시켜주면, 이미지가 깨지지도 않고 품질이 기존 대비 2배 이상 비약적으로 상승한다는 것을 증명해 냈습니다.

 

 

2 Background

더보기

2.1 Flow Matching

Flow matching(FM)은 generation problem을 $z \sim \mathcal{Z} \in \mathbb{R}^n$(보통 Gaussian $\mathcal{N}(0, I)$)인 prior distribution의 samples를 interpolation function으로 정의된 probability flow를 통해 data distribution $x \sim \mathcal{X} \in \mathbb{R}^n$으로 transport하는 것으로 formulate합니다:

$x_t = A(t) x + B(t)z, \quad (1)$

여기서 $t \in [0, 1]$입니다. Linear interpolation이 흔히 사용되며, 이때 $A(t) = 1 - t$, $B(t) = t$, 그리고 $x_t = (1 - t)x + tz$입니다.

$x$ $z$에 condition된 위치 $x_t$에서의 time derivative를 conditional velocity $\bar{v}_t$라고 부르며, 다음과 같이 도출할 수 있습니다:

$\bar{v}_t = \frac{dA(t)}{dt}x+\frac{dB(t)}{dt}z, \quad (2)$

linear interpolation의 경우, $\frac{dA(t)}{dt} = -1$, $\frac{dB(t)}{dt} = 1$, 그리고 $\bar{v}_t = -x + z$가 됩니다.

Flow matching은 conditional velocity $\bar{v}_t$와 match되도록 generator $G(x_t, t) : \mathbb{R}^n \times [0, 1] \rightarrow \mathbb{R}^n$을 train시킵니다:

$\mathcal{L}_{\mathrm{FM}}= \mathbb{E}_{x,z,t} \left[d(G(x_t, t), \bar{v}_t)\right], \quad (3)$

그리고 이 conditional flow matching objective가 $x, z$의 independent coupling에 대한 expectation 상에서, criterion $d(a, b)$가 다음을 만족할 때 probability flow의 marginal velocity $v_t = \mathbb{E}[\bar{v}_t | x_t]$를 learn한다는 것을 발견했습니다:

$\arg \min_a \mathbb{E}_b[d(a, b)] = \mathbb{E}[b]. \quad (4)$

squared $L_2$ criterion이 채택됩니다. 추가적으로 $\frac{1}{n}$ factor가 있는 mean squared error(MSE) variant가 가장 흔하게 사용됩니다:

$\mathcal{L}_{\mathrm{FM}}= \mathbb{E}_{x,z,t} \left[\frac{1}{n}\|G(x_t, t)-\bar{v}_t\|^2_2\right]. \quad (5)$

결과적인 generator $G(x_t, t)$는 probability flow를 따라 각 state $x_t$에서 marginal velocity field $v_t$를 predict하는 continuous-time flow model을 정의합니다. Samples는 ODE를 integrate하여 noise distribution에서 data distribution으로 transport됩니다:

$x_0 = x_1 + \int_0^1 G(x_t, t)\ dt, \quad x_1\sim \mathcal{Z}, \quad (6)$

여기서 integration은 $t = 1$에서 $t = 0$으로 backward로 실행됩니다.

The limitation of flow matching. Eq. (4)를 만족하는 어떠한 criterion $d(a, b)$를 사용하더라도 이론적으로는 model이 $v_t = \mathbb{E}[\bar{v}_t | x_t]$로 수렴하는 것을 보장하지만, 이는 training samples만을 generation하는 데 overfit됩니다. 실제로는 neural networks로 parameterize된 models는 finite capacity를 가지며 generalized distribution을 learn합니다. 이 경우, loss objective는 generalization의 방식에 영향을 미칩니다. 다음을 고려해 보십시오:

$d(a,b) = (a-b)^\top M (a-b), \quad (7)$

여기서 squared $L_2$ criterion은 특수한 경우인 $M = I$에 해당합니다. 일반적으로 $M$은 어떠한 strictly positive definite matrices라도 될 수 있습니다(Sec. C). 이러한 objectives는 동일한 ground-truth flow로 수렴하지만, 다른 generalizations를 유도할 수 있습니다. Flow matching은 data manifold에 대한 인식 없이 isotropic Euclidean distance를 최소화하여 부정확한 generalization과 out-of-distribution generation을 초래합니다.

Using a manifold criterion. 자연스러운 아이디어는 squared $L_2$ criterion을 data manifold 상의 distance를 측정하는 criterion으로 대체하는 것입니다. 그러나 underlying data manifold는 사전에 알려져 있지 않으며, 제한된 training data로부터 그 자체로 추론되고 generalized되어야 합니다. 우리의 연구는 criterion network가 generator와 함께 동시에 learn되는 adversarial training을 탐구합니다. 이는 deep networks가 Euclidean metric보다 더 나은 perceptual distance 역할을 할 수 있는 능력으로 증명되듯, data manifold를 더 잘 포착할 수 있다는 이전의 경험적 발견에 의해 장려됩니다.

2.2 Adversarial Flow Models

Adversarial flow models(AFMs)는 adversarial objective로 trained된 discrete-time flow model의 한 종류입니다. training은 probability flow 상에서 source $x_s$에서 target $x_t$로 samples를 transport하는 generator $G(x_s, s, t) : \mathbb{R}^n\times[0, 1]\times[0, 1] \rightarrow \mathbb{R}^n$과 real 및 generated $x_t$ samples를 구별하는 discriminator $D(x_t, t) : \mathbb{R}^n\times[0, 1] \rightarrow \mathbb{R}$을 포함합니다.

Adversarial training은 $D$가 discrimination을 최대화하는 반면 $G$ $D$에 의한 discrimination을 최소화하는 것을 목표로 하는 minimax optimization game을 포함합니다. adversarial objective는 다음과 같이 정의됩니다:

$\mathcal{L}^D_\textrm{adv} = \mathbb{E}_{x,z,s,t} \left[ f\big(D(x_t, t), D(G(x_s, s, t), t)\big) \right],$

$\mathcal{L}^G_\textrm{adv} = \mathbb{E}_{x,z,s,t} \left[ f\big(D(G(x_s, s, t), t), D(x_t, t)\big) \right], \quad (9)$

여기서 $f(a, b) = -\log(\text{sigmoid}(a-b))$는 최근 연구에서 사용된 많은 실행 가능한 contrastive functions 중 하나입니다. Training은 $G$ $D$를 교대로 update하며 $G(x_s, s, t)$가 동일한 distribution의 $x_t$를 produce할 때 equilibrium에 도달합니다.

AFMs는 추가적으로 $G$에 optimal transport objective를 도입합니다:

$\mathcal{L}^G_\textrm{ot} = \mathbb{E}_{x,z,s,t} \left[ \frac{1}{n} \cdot \frac{1}{|t-s|}\cdot \| G(x_s, s, t) - x_s \|^2_2 \right]. \quad (10)$

이는 flow matching과 달리 $\bar{v}_s$가 아니라 $x_s$까지의 distance를 최소화합니다. expectation 상에서, 이는 $G$가 target $x_t$가 source $x_s$에 가장 가깝다고 predict하도록 장려하여 $G$가 stable training을 위한 고유한 optimal transport를 learn할 수 있게 합니다.

추가적으로, $D$는 vanishing gradient의 문제를 완화하기 위해 gradient penalties $R_1$ $R_2$, 그리고 logit drifting을 방지하기 위해 centering penalty에 의해 규제됩니다:

$\mathcal{L}_\mathrm{r1}^D = \mathbb{E}_{x,z,s,t} \left[\| \nabla_{x_t} D(x_t,t) \|^2_2\right],$

$\mathcal{L}_\mathrm{r2}^D = \mathbb{E}_{x,z,s,t}\left[ \| \nabla_{G(x_s,s,t)}D(G(x_s,s,t), t)\|^2_2 \right],$

$\mathcal{L}_\mathrm{cp}^D = \mathbb{E}_{x,z,s,t} \left[ ( D(x_t,t) + D(G(x_s, s, t), t) )^2 \right]. \quad (13)$

AFMs의 최종 training objectives는 다음과 같습니다:

$\mathcal{L}^D_\textrm{AFM} = \mathcal{L}^D_\textrm{adv} + \lambda_\textrm{gp}\mathcal{L}^D_\textrm{r1} + \lambda_\textrm{gp}\mathcal{L}^D_\textrm{r2} + \lambda_\textrm{cp}\mathcal{L}_\mathrm{cp}^D,$

$\mathcal{L}^G_\textrm{AFM} = \mathcal{L}^G_\textrm{adv} + \lambda_\textrm{ot}\mathcal{L}^G_\textrm{ot}. \quad (15)$

generate를 위해, AFMs는 difference equation을 풀어 noise distribution에서 data distribution으로 samples를 transport합니다:

$x_0 = x_1 + \sum_{i=1}^S \big(G(x_{\tau_i}, \tau_i, \tau_{i-1})-x_{\tau_i}\big), \quad x_1 \sim \mathcal{Z}, \quad (16)$

여기서 summation은 총 $S$ sampling steps로 $i = S$에서 $i = 1$로 backward로 실행되며, $\tau$ $\tau_0 = 0$, $\tau_S = 1$을 만족하는 discrete timesteps의 list입니다.

The limitation of adversarial flow models. AFMs는 discrete-time flow models의 한 형태입니다. timestep interval $|t-s|$를 임의로 작게 만들 수 있지만, $|t - s| \rightarrow 0$일 때 training이 점점 더 불안정해지고 objective가 무너집니다. adversarial training을 continuous-time flow modeling으로 어떻게 확장할 수 있는지는 명확하지 않습니다. 게다가, AFMs는 여전히 gradient-vanishing problem을 가지고 있습니다. 그들은 이 문제를 완화하기 위해 gradient penalties, discriminator augmentation, 그리고 discriminator reset에 의존합니다.

 

2 Background 핵심 정리 노트

전문적인 AI 연구자들을 위해 본 논문의 문제 제기와 선행 연구의 근본적 한계를 중심으로 요약했습니다.

2.1 Flow Matching (FM)의 일반화 오류와 매니폴드 도입의 필요성

  • 작동 원리: 노이즈 $z$와 데이터 $x$ 사이의 probability flow를 정의($x_t$)하고, 이 경로의 conditional velocity $\bar{v}_t$를 예측하도록 모델 $G(x_t, t)$를 학습함. 생성 시에는 $t=1$에서 $0$으로 ODE를 적분함.
  • 핵심 한계 (Generalization Issue): 이론상 무한한 용량의 모델에서는 ground-truth flow로 수렴하지만, 실제 유한한 용량의 neural network에서는 필연적으로 generalization이 발생함. 이때 현재의 FM은 data manifold를 전혀 인식하지 못하는 isotropic Euclidean distance(squared $L_2$ criterion)를 최적화 목표로 삼음. 이는 실제 데이터 분포의 기하학적 구조를 무시한 잘못된 generalization을 유도하여 out-of-distribution 생성을 초래함.
  • 해결 방향 (Adversarial Training): 유클리디안 거리 대신 data manifold 상의 거리를 측정하는 criterion이 필요함. 하지만 사전 정의된 manifold가 없으므로, deep network가 perceptual distance를 훌륭하게 포착한다는 점에 착안하여 generator와 함께 criterion network(discriminator)를 동시에 학습시키는 adversarial training 도입을 제안함.

2.2 Adversarial Flow Models (AFMs)의 의의와 연속 시간 확장의 한계

  • 작동 원리: Adversarial objective를 도입한 discrete-time flow model. Discriminator $D$와 Generator $G$ 사이의 minimax game을 통해 학습하며, $G$는 목표물 $x_t$가 출발지 $x_s$에 가장 가깝도록 유도하는 optimal transport objective $\mathcal{L}^G_\textrm{ot}$를 추가로 사용하여 안정성을 확보함.
  • 핵심 한계 (Continuous-time Extension Failure): AFM은 기본적으로 discrete-time formulation에 갇혀 있음. 이론적으로 timestep interval $|t-s|$를 $0$에 가깝게 줄일 수 있어야 하지만, 실제로 $|t-s| \rightarrow 0$이 되면 학습이 극도로 불안정해지고 objective 자체가 붕괴됨. 즉, continuous-time flow modeling으로의 확장이 불가능함.
  • 불안정성 문제: 여전히 고질적인 gradient-vanishing 문제가 존재하며, 이를 막기 위해 gradient penalties($R_1, R_2$), discriminator augmentation, discriminator reset 등 다양한 휴리스틱과 규제항에 극도로 의존해야 하는 불안정한 구조를 가짐.

쉬운 설명 :

2.1 Flow Matching

기존의 Flow Matching 기술은 노이즈(아무 의미 없는 픽셀)가 실제 이미지가 되는 과정을 단순히 '점과 점 사이의 직선거리($L_2$ 거리)'로만 계산해서 학습합니다. 하지만 실제 데이터들이 모여 있는 공간은 구불구불하고 복잡한 산맥(Manifold)과 같습니다. 기존 방식은 이 산맥의 지형을 완전히 무시하고 허공을 가로지르는 최단거리만 고집하다 보니, 경로를 이탈해 기괴하거나 깨진 이미지(Out-of-distribution)를 만들어내는 부작용이 있었습니다. 그래서 연구진은 픽셀 단위의 직선거리가 아니라, 진짜 이미지의 형태와 질감 같은 '지형'을 볼 줄 아는 감별사(Discriminator)를 도입해서 AI가 올바른 길을 찾도록 유도하자는 아이디어를 냈습니다.

2.2 Adversarial Flow Models

사실 이 감별사를 Flow 모델에 붙여보려는 시도(AFM)가 과거에 이미 있었습니다. 하지만 이 과거 모델은 징검다리를 건너듯 '띄엄띄엄(Discrete-time)' 이동하는 방식에서만 간신히 작동했습니다. 만약 시간 간격을 아주 잘게 쪼개서 물 흐르듯 자연스러운 연속적인 경로(Continuous-time)를 만들려고 하면, 수식이 붕괴되면서 AI가 학습을 포기해버리는 치명적인 한계가 있었습니다. 게다가 띄엄띄엄 학습하는 것조차 너무 불안정해서, 모델이 망가지는 것을 막기 위해 온갖 덕지덕지 붙인 복잡한 페널티 규칙들이 필요했습니다. 결국, 이 섹션은 "과거 연구가 실패했던 연속 시간(Continuous-time)에서의 감별사 도입을 우리가 해내겠다"는 본 논문의 빌드업 역할을 합니다.

 

 

3 Method

더보기

3.1 Continuous Adversarial Flow Models

우리는 adversarial training을 continuous-time flow modeling으로 확장하기 위해 continuous adversarial flow models(CAFMs)를 제안합니다. 우리의 방법은 flow matching과 동일한 형태의 generator $G(x_t, t) : \mathbb{R}^n \times [0, 1] \rightarrow \mathbb{R}^n$(여기서 $x_t$에서 velocity field $v_t$를 predict함)과 AFMs와 동일한 형태의 discriminator $D(x_t, t) : \mathbb{R}^n \times [0, 1] \rightarrow \mathbb{R}$을 포함합니다. discrete-time adversarial training과 달리, 우리는 $D$의 derivative space에서 $v_t$를 discriminate하여 position $x_t$의 derivative로서 velocity $v_t$의 물리적 특성을 명시적으로 반영합니다.

구체적으로, 우리는 primal $(x_t, t)$와 tangent $(\dot{x}_t, \dot{t})$를 가진 $D$의 Jacobian-Vector Product(JVP)를 다음과 같이 나타냅니다:

$D_\mathrm{jvp}(x_t, t, \dot{x}_t, \dot{t}) = \frac{\partial D(x_t,t)}{\partial x_t}\dot{x}_t + \frac{\partial D(x_t,t)}{\partial t}\dot{t}, \quad (17)$

여기서:

$\frac{\partial D(x_t,t)}{\partial x_t}\in \mathbb{R}^{1\times n} \quad \text{and} \quad \frac{\partial D(x_t,t)}{\partial t}\in \mathbb{R}^{1\times 1} \quad (18)$

는 primal variables $x_t$ $t$에 대한 실제 network $D(x_t, t)$의 Jacobian matrices입니다. 전체 JVP function 또한 scalar를 output하며, 이를 우리는 discrimination logit으로 사용합니다:

$D_\mathrm{jvp}(x_t, t, \dot{x}_t, \dot{t}) : (\mathbb{R}^n\times [0,1]\times \mathbb{R}^n\times [0,1])\rightarrow \mathbb{R}. \quad (19)$

training 동안, $D_\mathrm{jvp}$는 $(x_t, t)$를 primal로 하고 $(\bar{v}_t, T)$를 tangent로 사용하여 evaluate되며, 여기서 $t \in [0, 1]$로 trained된 networks의 경우 $T = 1$입니다. continuous-time adversarial objectives는 다음과 같이 정의됩니다:

$\mathcal{L}_\mathrm{adv'}^D = \mathbb{E}_{x,z,t} [ f(D_\mathrm{jvp}(x_t, t, \bar{v}_t, T), D_\mathrm{jvp}(x_t, t, G(x_t,t), T)) ],$

$\mathcal{L}_\mathrm{adv'}^G = \mathbb{E}_{x,z,t} [ f(D_\mathrm{jvp}(x_t, t, G(x_t,t), T), D_\mathrm{jvp}(x_t, t, \bar{v}_t, T)) ], \quad (21)$

여기서 우리는 이전 연구와 유사하게 bounded contrastive function을 채택합니다:

$f(a, b) = (a - 1)^2 + (b + 1)^2. \quad (22)$

model parameters에 대한 gradients는 JVP를 통해 backpropagate됩니다.

Fig. 2는 우리 방법의 직관과 training dynamics를 시각화합니다. 직관적으로, 우리의 discriminator $D$는 directional derivative가 real과 fake flows를 구분하는 scalar potential을 learn합니다. $D$는 더 현실적인 방향에 더 높은 potential을 할당하도록 learn하고, $G$ $D$의 potential을 최대화하는 방향으로 최적화됩니다. $G$가 ground-truth flow를 learn하고 $D$가 모든 곳에서 평평한 potentials를 output할 때 training은 equilibrium에 도달합니다.

Eqs. (20)과 (21)의 objectives는 derivative에만 패널티를 주는 반면 $D$의 절대값은 자유롭게 drift할 수 있기 때문에, $D$의 절대값을 0을 중심으로 중심을 맞추기 위해 centering penalty도 포함합니다:

$\mathcal{L}_\mathrm{cp'}^D = \mathbb{E}_{x,z,t} \left[ D(x_t, t)^2 \right]. \quad (23)$

$n > 1$인 high-dimensional flows에서 training할 때, $n$-dimensional input을 scalar value로 project하는 discriminator는 다수의 $v_t \in \mathbb{R}^n$이 동일한 value를 산출할 수 있기 때문에 ambiguity를 만듭니다. $G$는 null space를 exploit하는 것을 learn할 수 있고, $D$는 이 행동에 대응하기 위해 update됩니다. 그러나 이는 느린 수렴을 야기합니다. $G$가 optimal transport regularization과 관련된 minimum-norm solution을 선택하도록 장려하기 위해 regularizer를 추가할 수 있습니다. $G$에 대한 continuous-time optimal transport regularization은 $|t - s| \rightarrow 0$의 극한에서 Eq. (10)의 discrete 대응물과 같습니다:

$\mathcal{L}_\mathrm{ot'}^G = \mathbb{E}_{x,z,t} \left[ \frac{1}{n} \|G(x_t, t)\|^2_2 \right]. \quad (24)$

adversarial training을 continuous time으로 확장하는 것은 gradient vanishing 문제 또한 완화합니다(Sec. E). 경험적으로, 우리는 실험에서 gradient penalties 없이 CAFMs가 trained될 수 있음을 발견했습니다. 우리는 또한 $G$의 1회 update 당 $N$ steps 동안 $D$를 update하여 optimality를 향해 $D$를 train하는 것이 유리하다는 것을 발견했습니다.

CAFMs의 최종 objectives는 gradient penalties가 제거된 것을 제외하고 discrete-time 대응물과 유사합니다:

$\mathcal{L}_\mathrm{CAFM}^D = \mathcal{L}_\mathrm{adv'}^D + \lambda_\mathrm{cp}\mathcal{L}_\mathrm{cp'}^D,$

$\mathcal{L}_\mathrm{CAFM}^G = \mathcal{L}_\mathrm{adv'}^G + \lambda_\mathrm{ot}\mathcal{L}_\mathrm{ot'}^G. \quad (26)$

우리는 처음부터 training할 때 점진적으로 $\lambda_\mathrm{ot}$를 줄이기 위해 AFMs를 따릅니다. 기존의 flow-matching models를 post-training하기 위해, 우리는 Euclidean norm의 편향을 완전히 제거하기 위해 $\lambda_\mathrm{ot} = 0$으로 설정합니다. centering penalty는 $\lambda_\mathrm{cp} = 0.001$로 설정됩니다. 부록 Sec. D에서 추가 증명과 논의를 제공합니다.

3.2 Practical and Efficient Implementation

JVP는 forward-mode automatic differentiation을 사용하여 효율적으로 compute될 수 있습니다. 이는 단일 forward pass에서 $D(x_t, t)$와 $D_\mathrm{jvp}(x_t, t, \dot{x}t, \dot{t})$를 모두 compute하여 우리가 adversarial loss $\mathcal{L}^D\mathrm{adv'}$와 centering penalty $\mathcal{L}^D_\mathrm{cp'}$를 효율적으로 함께 도출할 수 있게 해줍니다. 추가적으로, 우리는 $D$를 update할 때 동일한 primal에서 다수의 tangents를 효율적으로 compute하기 위해 vectorizing map(vmap)을 사용합니다. 간결한 PyTorch implementation이 Algorithm 1에 제공됩니다. 더 큰 규모의 training을 위해, JVP와 vmap은 PyTorch의 DDP, FSDP, 그리고 gradient checkpointing과 호환됩니다. Implementation details는 Sec. F에 있습니다.

Algorithm 1 Continuous adversarial flow training

1 from functools import partial

2 from torch import mean , ones_like , stack , unbind

3 from torch . func import jvp , vmap

4

5 def step (G , D , x , z , t , c , mode , cp_scale , ot_scale ):

6 D. requires_grad_ ( mode == " dis")

7 G. requires_grad_ ( mode == " gen")

8

9 D = partial (D , condition =c )

10 G = partial (G , condition = c)

11

12 x_t = (1 - t) * x + t * z

13 v_t = -x + z

14 u_t = G( x_t , t )

15 T = ones_like ( t)

16

17 if mode == " dis ":

18 o , do = vmap ( lambda * tangents : jvp (D , ( x_t , t) , tangents ))(

19 stack ([ v_t , u_t ]) ,

20 stack ([T , T ])

21 )

22 dv , du = unbind ( do )

23 return (

24 mean (( dv - 1) ** 2) +

25 mean (( du + 1) ** 2) +

26 mean (o ** 2) * cp_scale

27 )

28 else :

29 _ , du = jvp (D , ( x_t , t) , ( u_t , T ))

30 return (

31 mean (( du - 1) ** 2) +

32 mean ( u_t ** 2) * ot_scale

33 )

network architecture 측면에서, $G$는 JVP computation을 포함하지 않으며 flow matching과 동일하게 어떠한 architectures도 사용할 수 있기 때문에 아무런 제한이 없습니다. $D$의 경우, 우리는 LayerNorm을 RMSNorm으로 변경하는 것이 JVP computation과 관련된 이전 연구의 발견과 일관되게 training stability를 크게 향상시킨다는 것을 발견했습니다. 이전 연구와 달리, 우리는 modulation에 대한 추가적인 normalization이 필요하다는 것을 발견하지 못했습니다. 우리의 실험은 CAFMs가 $G$ $D$ 모두로 standard transformers에서 잘 작동함을 보여줍니다.

3.3 Pre-training vs. Post-training

비록 CAFMs가 처음부터 trained될 수 있지만, 추가적인 discriminator network의 관여, JVP의 forward 및 backward computation, 그리고 generator update 당 다수 steps의 discriminator learning으로 인해 FMs보다 본질적으로 덜 효율적입니다. FMs와 CAFMs는 모두 동일한 probability flow를 learn하고 오직 model generalization에서만 차이가 나기 때문에, models를 FM objective로 pre-train하고 CAFM objective로 post-train하는 것이 훨씬 더 효율적입니다. 따라서 우리는 주로 post-training을 위해 CAFMs를 제안하지만, 그럼에도 불구하고 비록 덜 효율적일지라도 그 objective가 처음부터 models를 train하는 데 사용될 수 있음을 보여줍니다.

 

 

3 Method 핵심 정리 노트

전문적인 AI 연구자들을 위해 본 논문의 핵심 기여인 Continuous-time으로의 확장 방법론(JVP 활용)과 실용적인 구현 디테일을 중심으로 요약했습니다.

3.1 Continuous Adversarial Flow Models (CAFMs)의 핵심 설계

  • Derivative Space에서의 Discrimination (핵심 아이디어):
    • 기존 discrete-time AFM의 한계를 극복하기 위해, velocity $v_t$가 position $x_t$의 derivative라는 물리적 특성에 착안함.
    • Discriminator $D(x_t, t)$의 절대적인 출력값이 아닌, **Jacobian-Vector Product (JVP)**를 활용하여 derivative space에서 실제 $v_t$와 예측된 $v_t$를 판별함.
    • $D_\mathrm{jvp}(x_t, t, \dot{x}_t, \dot{t}) = \frac{\partial D(x_t,t)}{\partial x_t}\dot{x}_t + \frac{\partial D(x_t,t)}{\partial t}\dot{t}$ (이를 scalar logit으로 사용).
  • 목적 함수 및 학습 동인 (Training Dynamics):
    • Bounded contrastive function을 사용하여 Adversarial objectives ($\mathcal{L}_\mathrm{adv'}^D$, $\mathcal{L}_\mathrm{adv'}^G$)를 구성함.
    • $D$는 실제 flow 방향에 높은 potential을 부여하도록 학습되고, $G$$D$의 potential이 최대화되는 방향으로 최적화됨.
  • 필수 Regularization 및 Penalty:
    • Centering penalty ($\mathcal{L}_\mathrm{cp'}^D$): JVP 기반 목적 함수는 derivative에만 패널티를 주므로 $D$의 절대값이 무한히 표류(drift)하는 것을 막기 위해 0으로 중심을 잡아줌.
    • Continuous-time Optimal Transport (OT) regularization ($\mathcal{L}_\mathrm{ot'}^G$): 고차원($n>1$)에서 $D$가 scalar를 출력할 때 발생하는 ambiguity(다수의 $v_t$가 동일한 potential을 가지는 현상)와 $G$의 null space exploitation을 방지하기 위해, minimum-norm solution을 선택하도록 유도함.
  • 안정성 확보: Continuous-time으로 확장하면서 기존 discrete AFM에서 필수적이었던 gradient vanishing 문제가 자연스럽게 완화됨. 따라서 복잡한 gradient penalties ($R_1, R_2$)를 완전히 제거함.

3.2 실용적이고 효율적인 구현 (Practical Implementation)

  • Forward-mode Automatic Differentiation: PyTorch의 jvpvmap을 사용하여 단일 forward pass만으로 $D$와 $D_\mathrm{jvp}$를 동시에 효율적으로 계산함.
  • Architecture 튜닝:
    • Generator $G$는 기존 flow matching과 완벽히 동일한 구조 사용 가능.
    • Discriminator $D$의 경우, JVP 연산 시의 학습 안정성을 위해 LayerNormRMSNorm으로 교체하는 것이 매우 중요함 (추가적인 modulation normalization은 불필요).

3.3 Pre-training vs. Post-training 전략

  • CAFM은 $D$ 네트워크 추가, JVP 연산, $G$ 대비 $D$의 다중 업데이트(다대일 비율) 요구 등으로 인해 처음부터 학습(train from scratch)하기에는 기존 FM보다 계산 비용이 높음.
  • FM과 CAFM은 동일한 probability flow를 목표로 하고 오직 model generalization 방식에서만 차이가 남.
  • 따라서 가장 효율적인 SOTA 전략은: 기존 FM objectivePre-training된 모델을 가져와서, CAFM objective로 짧게 Post-training (Finetuning)하는 것임.

쉬운 설명 :

"순간적인 움직임(속도와 방향)을 평가하는 날카로운 감별사 도입"

1. 위치가 아니라 '방향과 속도'를 채점하다 (JVP의 도입)

기존의 감별사(Discriminator) 기술들은 주로 "너 지금 이상한 위치에 있네? 가짜!"라고 정적인 상태를 채점했습니다. 하지만 데이터를 생성하는 Flow 모델은 끊임없이 흘러가는 '경로'를 만드는 모델입니다.

그래서 연구진은 수학적인 기법(JVP)을 써서, 감별사가 현재 위치뿐만 아니라 **"지금 어느 방향으로, 얼마나 빠르게 꺾고 있는지(Derivative/Velocity)"**를 순간적으로 캐치해서 진짜 경로와 가짜 경로를 판별하게 만들었습니다. 물리학으로 치면 '위치'가 아니라 '속도 벡터'를 감별하는 셈입니다.

2. 곁가지 쳐내기 (안정화 작업)

  • 닻 내리기 (Centering penalty): 감별사가 '경사(기울기)'만 신경 쓰다 보니 전체적인 기준점이 안드로메다로 날아가는 문제가 생겨서, 기준점을 0으로 꽉 묶어두는 규칙을 추가했습니다.
  • 꼼수 방지 (OT regularization): 고차원의 복잡한 공간에서는 AI(생성자)가 감별사의 눈을 피해 이상한 지름길(Null space)로 빠질 수 있습니다. 이를 막기 위해 "가장 움직임이 적고 깔끔한 최단 경로(Minimum-norm)"를 선택하도록 강제했습니다.
  • 쓸데없는 규칙 폐기: 시간을 물 흐르듯 연속적(Continuous)으로 만들었더니, 과거 징검다리 방식(Discrete)에서 모델이 터지는 걸 막기 위해 쓰던 온갖 복잡한 페널티 규칙들이 더 이상 필요 없어져서 다 버렸습니다. 코드가 훨씬 깔끔해졌죠.

3. 효율적인 학습법: "처음부터 키우지 말고, 잘 큰 애를 과외시켜라"

이 감별사 방식은 성능은 끝내주지만, 연산량이 많아서 처음부터 백지상태의 AI를 학습시키려면 돈과 시간이 너무 많이 듭니다.

하지만 다행히도 기존의 일반 Flow 모델과 이 새로운 모델은 궁극적으로 '가야 할 목적지'가 같습니다. 단지 '가는 길을 얼마나 정교하게 찾느냐'의 차이일 뿐이죠. 그래서 연구진은 **"이미 학습이 끝난 기존 Flow 모델을 가져와서, 이 감별사 방식을 덧붙여 짧게 마무리 훈련(Post-training)만 시켜라"**라고 강력히 권장합니다. 이게 압도적으로 효율적이고 실용적이기 때문입니다.

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

 

더보기

기존 Flow Matching은 velocity를 MSE로 직접 맞추는 방식입니다.

 

문제의식은, 이런 MSE 기반 학습이 finite-capacity model에서는 데이터 manifold에 맞는 generalization을 잘 만들지 못할 수 있다는 것입니다.

 

그래서 velocity supervision을 단순한 MSE 기준으로 주는 대신, adversarial한 기준으로 바꾸자는 아이디어가 나옵니다.

 

구체적으로는 GAN처럼 discriminator를 붙이는데, sample 자체를 직접 구분하는 것이 아니라 현재 위치에서의 flow direction이 얼마나 real한지를 보도록 만듭니다.

 

이때 핵심은 discriminator를 그냥 velocity 입력까지 받는 형태로 두는 것이 아니라, discriminator의 Jacobian을 이용한 JVP, 즉 Jacobian-Vector Product 형태로 direction을 평가한다는 점입니다.

 

즉 현재 상태에서 discriminator가 어떤 방향으로 가장 많이 변하는지를 보고, real conditional velocity 방향과 generator가 예측한 velocity 방향을 비교하게 만드는 구조입니다.

 

그러면 generator는 discriminator가 보기에 더 real한 flow direction처럼 보이도록 velocity를 예측하게 됩니다.

 

결국 핵심은 velocity control을 MSE에서 adversarial signal로 바꾸고, 그 signal을 Jacobian 기반의 local direction 평가로 주자는 것입니다.

 

이 방식은 처음부터 새로 학습한다기보다는, 기존 Flow Matching 모델을 먼저 학습한 뒤에 post-training으로 붙이는 방식이 핵심입니다.