디퓨전 없이도 생성이 가능할까?

Drifting Model은 왜 “학습과정 자체”를 생성 경로로 바꾸려 했을까

생성 모델을 처음 배울 때 가장 자연스럽게 떠오르는 그림은 이렇다.
노이즈에서 시작해서, 여러 번의 정제를 거쳐, 조금씩 데이터 분포에 가까워지는 과정. 디퓨전도 그렇고, 플로우 매칭도 그렇다. 결국 추론 과정에서 여러 단계의 변환을 거쳐 복잡한 pushforward를 만들어내는 방식이다.

그런데 이 논문은 질문을 조금 다르게 던진다.

정말 그 복잡한 변환을 추론 시간에 해야만 할까?
어차피 딥러닝의 학습은 본질적으로 반복적이다. 그렇다면 그 반복적인 학습 과정 자체가, 노이즈 분포를 데이터 분포 쪽으로 끌고 가는 “누적된 변환”이 될 수 있지 않을까.

이 논문이 제안하는 Drifting Model은 바로 그 생각에서 출발한다. 매 추론 단계에서 경로를 적분하는 대신, 학습 중에 샘플이 어느 방향으로 움직여야 하는지 알려주는 드리프팅 필드를 정의하고, 네트워크가 그 방향을 학습하도록 만든다. 그 결과 최종 추론은 놀랍게도 1 step이면 충분하다. 실제로 ImageNet 256x256에서 저자들의 최종 모델은 1-NFE 설정으로 FID 1.54를 기록했고, 당시 기준으로 1-step 생성 모델 중 매우 강한 성능을 보여주었다.

이 글에서는 논문의 전체 흐름을 따라가되, 단순한 요약보다는
“이 논문이 무엇을 하려는지”,
“왜 그런 설계를 했는지”,
“흥미로운 지점들”
을 중심으로 설명해보려 한다.

생성은 결국 pushforward를 배우는 일이다

생성 모델의 출발점은 간단하다.
사전 분포 $p_\epsilon$ 에서 노이즈 $\epsilon$ 를 뽑고, 신경망 $f_\theta$ 를 통과시켜 샘플 $x = f_\theta(\epsilon)$ 을 만든다. 이때 모델이 만들어내는 전체 샘플 분포를 $q_\theta$ 라고 쓰면,

라고 쓸 수 있다. 말 그대로, $f_\theta$ 가 prior를 밀어낸(pushforward) 결과가 생성 분포라는 뜻이다.

그러니 생성 모델의 목적은 결국 하나다.

즉, 사전 분포를 데이터 분포처럼 보이게 만드는 변환을 학습하는 것.

디퓨전과 플로우 계열 모델은 이 변환을 “한 번에” 하지 않는다. 대신 아주 많은 작은 변환으로 나누어 추론 과정에서 점진적으로 적용한다. 반면 Drifting Model은 이 누적 변환을 추론 시점이 아니라 학습 시점으로 옮기려는 시도에 가깝다.

추론 경로를 학습과정에 밀어 넣기

이 논문에서 가장 흥미로운 점은, 생성 샘플을 정제하는 연산을 별도의 ODE solver나 diffusion trajectory로 두지 않았다는데 있다.

저자들은 다음과 같은 관점을 취한다.

학습 도중 파라미터가 계속 바뀌면, 같은 노이즈 $\epsilon$ 을 넣어도 모델이 내는 샘플은 조금씩 달라진다.
즉 학습 단계 $i$ 에서의 샘플 $x_i = f_{\theta_i}(\epsilon)$ 은, 다음 단계에서

처럼 움직인다.

그렇다면 생성 모델을 학습한다는 것은, 본질적으로 샘플을 데이터 분포 쪽으로 조금씩 이동시키는 규칙을 학습하는 것과도 같다.

이 “조금씩 이동시키는 규칙”이 정해지는 곳이 바로 드리프팅 필드다.

드리프팅 필드는 “다음에 어디로 가야하는지”를 말해주는 벡터장이다

논문에서 드리프팅 필드 $\mathbf{V}_{p, q}$ 는 현재 샘플 $x$ 가

• 실제 데이터 분포 $p$ 쪽으로는 끌리고,
• 현재 생성 분포 $q$ 쪽에서는 밀려나도록

정의되는 벡터장이다.

즉 샘플 업데이트는 다음처럼 생각할 수 있다.

여기서 중요한 건, 이 벡터장이 단순히 “그럴듯한 방향”이 아니라, 최종적으로 평형 상태를 갖도록 설계되어야 한다는 점이다.

만약 현재 생성 분포 $q$ 가 이미 데이터 분포 $p$ 와 같아졌다면, 더 이상 움직일 이유가 없어야 한다. 그러므로 이상적인 드리프팅 필드는

을 만족해야 한다. 논문은 이를 보장하기 위한 충분조건으로 반대칭성(anti-symmetry)

을 둔다. 그러면 $p=q$ 일 때, $\mathbf{V}_{p, p}=0$ 이 자동으로 따라온다.

“움직이지 않는다”는 것이 정말 “정답에 도달했다”는 뜻일까

이 논문에서 흥미로운 지점 중 하나가 바로 여기다.

$p=q$ 이면 drift가 0이 되어야 한다는 것은 이해하기 쉽다. 하지만 그 역, 즉

는 일반적으로 자동으로 성립하지 않는다. 논문도 이 점을 분명히 인정한다. 임의의 벡터장에서는 zero-drift가 곧 distribution matching을 의미하지 않을 수 있다.

이 부분을 나는 줄다리기 비유로 이해하는 편이 좋았다.

완전히 동일한 힘이 양쪽에서 작용하면 줄은 움직이지 않는다.
하지만 줄이 움직이지 않는다고 해서, 양쪽 사람이 동일하다고 결론 낼 수는 없다.

즉, 정지 상태와 정답 상태는 다르다.

논문 후반부와 부록에서는 바로 이 문제,
“왜 zero-drift가 단순한 정지가 아니라 distribution matching에 가까운 제약이 되는가”
를 설명한다. 완전히 일반적인 정리는 아니지만, 자신들이 제안한 kernelized drifting field에서는 그 조건이 꽤 강한 제약이 된다고 주장한다.

Drift의 크기를 줄이도록 하는 손실함수

이상적인 fixed point는

형태다. 즉 이미 평형에 도달했다면 drift를 한 번 적용해도 샘플이 바뀌지 않는다. 논문은 이 업데이트를 loss로 바꾸어 다음과 같이 쓴다.

핵심은 오른쪽 항이 frozen target이라는 점이다.

즉 현재 샘플을 drift를 한 번 적용한 위치까지 이동시켜 놓고, 그 위치를 이번 iteration에서는 고정된 목표점처럼 취급한다. 논문도 stop-gradient가 “frozen state”를 제공하고, 네트워크 예측을 그 frozen target 쪽으로 이동시키는 방식이라고 설명한다.

여기서 stopgrad는 단순한 구현 트릭이 아니다.
오히려 이 논문의 주요 장치로 작용한다.

왜냐하면 drift \(\mathbf{V}_{p, q_\theta}(x)\) 는 샘플 하나 \(x\) 만의 함수가 아니라, 현재 generator 전체가 유도한 분포 \(q_\theta\) 에도 의존하기 때문이다. 즉 이를 그대로 미분하려면 샘플 하나가 아니라 분포 자체의 변화를 함께 따라가야 한다. 논문은 바로 이 점 때문에 \(\mathbf{V}\) 를 직접 backpropagate 하지 않고, drifted target을 고정해 간접적으로 최적화한다고 말한다.

조금 추상적으로 말하면,
“분포를 직접 미분하는 대신, 분포가 가리키는 다음 위치를 정답처럼 만들어 회귀한다”
라고 이해하면 된다.

드리프팅 필드는 어떻게 계산할까

이제 진짜 중요한 질문이 남는다.
도대체 $\mathbf{V}_{p, q}(x)$ 는 어떻게 정의해야 할까.

논문은 가장 일반적인 형태로 다음을 둔다.

즉

• positive 쪽으로 가는 인력(attraction)
• negative 쪽으로 가는 척력(repulsion)

의 차로 drift를 만든다. Figure 2가 바로 이 해석을 시각적으로 보여준다.

수식으로 쓰면,

이다.

겉보기에는 단순하지만, 이 구조 덕분에 드리프팅 필드는 “실제 데이터에는 끌리고, 현재 생성이 몰린 곳에서는 밀려나는” 매우 직관적인 mean shift 형태를 갖게 된다.

Generated sample 의 정의

이 논문을 처음 읽을 때 헷갈렸던 부분이다.
Negative sample $y^-$ 를 generated sample에서 뽑는다면, 그럼 $x$ 는 뭐지?

정확히 말하면, $x$ 역시 현재 배치에서 generator가 만든 샘플 중 하나다.
구현에서는 보통 현재 배치 전체를 generated samples로 만들고, 그 배치 자체를 negative pool로 재사용한다. 즉 각 샘플은 어떤 순간에는 “지금 drift를 계산할 기준점 $x$”가 되고, 다른 샘플의 관점에서는 negative reference $y^-$ 가 된다. 논문의 의사코드도 바로 이런 형태를 쓴다.

e = randn([N, C])
x = f(e)          # generated samples
y_neg = x         # reuse generated samples as negatives

V = compute_V(x, y_pos, y_neg)
x_drifted = stopgrad(x + V)
loss = mse_loss(x, x_drifted)

즉 각 샘플마다 $\mathbf{V}(x_i)$ 를 따로 계산해야 하지만, 실제 구현은 for-loop가 아니라 batch 전체에 대한 pairwise distance matrix와 softmax, 행렬곱을 통해 한 번에 처리한다. 논문 부록의 Algorithm 2가 바로 그 계산 방식을 정확히 따르고 있다.

커널함수로 사용한 RBF, softmax를 통한 구현

논문은 샘플간 유사도를 재는 커널로 다음을 사용한다.

즉 $\ell_2$ 거리가 가까운 샘플일수록 큰 값을 주는 RBF류 커널이다. 여기서 $\tau$ 는 얼마나 local 하게 볼지를 정하는 temperature 다.

그런데 논문은 이를 구현할 때, 정규화된 커널 $\tilde{k}$ 를 softmax로 계산한다고 말한다.
이건 커널을 다른 함수로 바꾼다는 뜻이 아니라,

처럼 각 $x$ 에 대해 후보 $y$ 들 사이에서 정규화한 가중치를 쓴다는 뜻이다. 즉 softmax의 logit이 단지 $-|x-y_i|/\tau$ 일 뿐이다. 논문도 이를 InfoNCE의 softmax와 유사하다고 설명한다.

왜 $y$ 축으로 정규화하느냐도 자연스럽다.
Drift는 “하나의 $x$ 가 여러 후보 $y$ 중 누구에게 얼마나 끌리거나 밀릴지”를 정하는 것이므로, 각 $x$ 에 대해 $y$ 후보들 사이에서 weight를 나누는 것이 맞다.

논문은 여기에 더해 $x$ 축 정규화도 한 번 더 넣으면 약간의 성능 향상이 있었다고 보고한다. 부록의 ablation을 보면 정규화가 없는 경우보다 $y$-softmax를 넣은 경우가 확실히 좋고, $x+y$ 이중 정규화를 넣으면 조금 더 좋아진다. 결국 이 방법은 단순히 “가깝다/멀다”를 보는 것을 넘어서, batch 수준에서 더 안정적인 similarity geometry를 만들려는 시도에 가깝다.

Drifting model이 잘 동작하기 위한 임베딩 공간의 설정

그리고 이 지점에서 feature encoder 이야기가 등장한다.

논문은 드리프팅을 꼭 픽셀 공간에서 계산할 필요가 없다고 말한다. 오히려 고차원 데이터ㅓ에서는 feature spaceㅔㅇ서 drift를 계산하는 것이 더 유리할 수 있다고 본다. 실제로 저자들은 여러 scale/location의 feature에 대해서 drifting loss를 계산하고, 그게 richer gradient information을 준다고 설명한다.

Feature-space loss는 다음처럼 작성할 수 있다.

이 식은 얼핏 perceptual loss와 비슷해보이지만, 논문은 둘이 개념적으로 다르다고 선을 긋는다.

Perceptual loss는 $x$ 와 짝지어진 정답 $x_\text{target}$ 의 feature $\phi(x_\text{target})$ 을 회귀 목표로 삼는다. 반면 drifting의 target은 $\phi(x) + V(\phi(x))$, 즉 현재 feature를 distribution-level drift 방향으로 한 번 이동시킨 값이다. 그래서 이 loss는 paired target이 없이도 되고, 원리적으로는 feature space의 pushforward 분포 $\phi_\# q$ 와 $\phi_\# p$ 를 맞추려는 목적을 갖는다.

이 대목은 이 논문을 단순한 perceptual regression과 구분해주는 중요한 문장이다.

왜 SSL feature encoder가 필요할까

이 논문에서 feature encoder는 선택사항처럼 보이지만, 실제로는 거의 핵심 부품에 가깝다.

저자들은 자기 방법이 커널 $k(\cdot, \cdot)$ 로 샘플 유사도를 재기 때문에, feature space에서 의미론적으로 비슷한 샘플이 실제로 가까이 위치하는 것이 중요하다고 설명한다. 그리고 이 목표가 SSL과 잘 맞는다고 말한다. 그래서 pre-trained SSL 모델을 feature extractor로 사용한다.

왜 SSL 이 적합한지는 생각보다 직관적이다.

SSL 은 본질적으로

• 같은 입력의 다른 view는 비슷한 representation으로 보내고,
• 다른 샘플들은 구분되는 방향으로

feature space를 학습한다. 그러면 $\phi(x)$ 공간에서의 거리 $| \phi(x) - \phi(y) |$ 가 단순한 픽셀 차이보다 더 “의미 있는 유사도”가 되기 쉽다. 그리고 drifting은 바로 그 거리 위에 RBF 커널을 얹어 작동한다.

즉, 좋은 SSL encoder는 drifting에 필요한 좋은 거리 공간을 제공하는 셈이다.

SimCLR, MoCo-v2 도 좋지만, latent-MAE가 더 좋았다

논문 7페이지의 표3은 꽤 중요한 실험을 보여준다.

저자들은 public pre-trained SSL 인코더인 SimCLR과 MoCo-v2도 비교했고, 둘다 나쁘지 않은 결과를 냈다. 하지만 자신들이 설계한 latent-MAE가 더 좋았고, width와 pretrain epoch을 늘릴수록 성능이 계속 좋아졌다. 특히 latent-MAE를 640 width, 1280 epoch로 pretrain하고 classification fine-tuning 까지 붙인 설정은 FID 3.36 까지 내려간다.

여기서 중요한 건, “SimCLR과 MoCo-v2가 SSL 이 아니어서 안 좋았다”는 게 아니라는 점이다.
둘 다 SSL 인코더가 맞다. 문제는 어느 공간에서 feature를 뽑느냐다.

SimCLR과 MoCo-v2는 기본적으로 pixel-domain 인코더다.
그런데 이 논문의 generator는 SD-VAE latent 공간 (32, 32, 4) 에서 작동한다. 그러므로 pixel-space SSL encoder를 사용하려면, 학습 중 매번 VAE decoder 를 돌려 latent를 이미지로 복원한 뒤 feature 를 뽑아야 한다. 논문은 바로 이 점을 지적하고, 이를 우회하기 위해 latent space 위에서 직접 pretrain한 ResNet-style latent-MAE를 도입했다고 설명한다.

결국 latent MAE가 더 좋았던 이유는 2가지로 요약된다.

첫째, 생성 공간과 feature 공간의 mismatch를 줄였다.
둘째, 좋은 feature encoder일수록 커널이 너무 flat해지지 않고, 가까운 샘플들 사이에서 더 강한 drift를 만들어 richer training signal을 준다.

논문도 표3 아래에서 바로 이 해석을 제시한다. 강한 feature encoder는 거의 0이 되어버리는 flat kernel의 발생을 줄여주며, 샘플 간 유사도를 더 잘 반영한다고 말한다. 심지어 저자들은 feature encoder 없이 ImageNet에서 모델을 제대로 작동시키지 못했다고도 적는다.

이 문장은, 이 논문이 단순히 “drift라는 아이디어만 좋다”로 끝나지 않는다는 걸 잘 보여준다.
좋은 drift를 만들려면, 그 전에 좋은 similarity geometry가 필요하다.

Classifier Free Guidance의 재해석

이 논문은 Classifier-Free Guidance(CFG)도 drifting 관점으로 다시 쓴다.

조건 클래스 $c$ 가 주어지면, positive는 $p_\text{data}(\cdot \mid c)$ 에서 뽑으면 된다.
그런데 guidance를 얻으려면 negative도 단순히 생성 샘플만으로는 부족할 수 있다. 그래서 저자들은 negative 분포를

처럼 정의한다. 논문 본문은 $p_\text{data}(\cdot \mid \varnothing)$ 를 unconditional data distribution으로 쓰지만, 각주에서는 엄밀히는 조건 크래스 $c$ 를 제외한 다른 크래스들의 분포를 뜻한다고 설명한다.

이 식을 풀면,

가 나오고, 여기서 $\alpha = 1/(1-\gamma)$ 다.

이 식은 CFG가 왜 “조건부 분포를 강화하고 비조건부 분포를 빼는 방식”인지 잘 보여준다. 논문도 이 formulation이 original CFG의 spirit과 맞닿아 있다고 설명한다.

개인적으로는 이 부분을 읽고, CFG를 다시 분포의 외삽(extrapolation)으로 이해하게 되었다.
고양이를 만들 때 단지 진짜 고양이 쪽으로만 가는 것이 아니라, 고양이가 아닌 실제 데이터 manifold와도 멀어지게 만드는 것. Drifting 관점은 이 구조를 자연스럽게 설명한다.

배치를 구성하는 방법

논문의 5페이지는 class label을 기준으로 배치를 구성한다고 설명한다.

먼저 $N_c$ 개의 클래스 라벨을 샘플링하고, 각 라벨마다 Algorithm 1을 독립적으로 수행한다. 클래스 하나당 $N = N_\text{neg}$ 개의 generated negatives와 $N_\text{pos}$ 개의 positive를 준비하며, 전체 effective batch size는

가 된다. 즉 하나의 배치는 사실 여러 개의 “클래스별 작은 배치”들을 묶은 형태에 가깝다.

이 구조는 drifting이 단순히 전체 생성분포를 맞추는 것에 그치지 않고, 조건부 분포별로도 작동하도록 만들어준다.

Toy 실험으로 보는 분포의 변화

논문은 단순한 toy example도 보여준다.

이봉(binomial) 분포를 목표로 했을 때, 초기 생성분포가 한쪽으로 쏠려 있어도 다른 mode 쪽 positive 들이 drift를 유도하면서 mode collapse를 완화하는 모습을 시각화한다. 또 학습이 진행될수록 생성 분포가 실제로 데이터 분포 쪽으로 옮겨가고, 어느 순간 평형 상태에 가까워지는 모습을 보여준다. 이 그림들은 이 논문이 단순히 숫자 경쟁만 하는 것이 아니라, 정말로 “분포를 이동시키는 관점”을 갖고 있음을 보여준다.

그리고 정량 실험에서는

• positive/negative 샘플 수가 늘어날 수록 $\mathbf{V}$ 추정이 더 정확해져 품질이 좋아지고,
• drifting loss를 계산하는 feature space의 선택이 성능에 큰 영향을 미치며,
• 더 긴 학습과 더 강한 하이퍼파라미터 recipe, 더 큰 모델로 갈수록 FID가 계속 좋아지는 흐름을 보여준다.

특히 표4에서는 baseline 3.36 에서 시작해, 더 긴 학습과 recipe tuning으로 1.75, 더 큰 L/2 모델에서 1.54 까지 내려가는 과정을 보여준다. 여기서 흥미로운 점 하나는, 최종 최고 성능이 CFG scale 1.0, 즉 diffusion 문맥에서는 사실상 “no CFG”에 해당하는 설정에서 나왔다는 것이다. 논문은 이 CFG formulation이 FID와 IS 사이의 trade-off를 보인다고 설명한다.

논문이 던지는 가장 좋은 질문

개인적으로 이 논문의 가장 큰 장점은
“우리는 왜 생성 경로를 추론 시간에만 두고 생각해왔을까?”
라는 질문을 던진다는 점이다.

디퓨전과 플로우 매칭은 매우 강력하지만, 결국 추론 시점에 trajectory를 따라가야 한다.
Drifting Model은 그 trajectory를 학습 과정에 흡수할 수 있다고 주장한다. 그리고 그 아이디어를 추상적인 철학이 아니라,

• anti-symmetric drifting field
• stop-gradient fixed-target loss
• kernelized attraction/repulstion
• SSL feature-space drifting

이라는 꽤 구체적인 설계로 연결했다.

물론 논문도 한계를 인정한다.
평형상태의 역, 즉 $\mathbf{V}=0$ 이 언제 정말 $p=q$ 를 보장하는지는 완전히 일반적인 정리로 끝내지는 못했고, kernel이나 feature encoder의 선택도 아직 최적이라고 보기 어렵다. Feature encoder 없이 ImageNet에서 제대로 작동하지 못했다는 보고 역시, 이 방법이 “아무 공간에서나 통하는 보편적 해답”은 아니라는 점을 보여준다.

마치며

Drifting Model은 한 문장으로 요약하면 이렇다.

추론 사건의 정제를 학습 시간의 반복으로 밀어 넣고, 생성 샘플이 데이터 분포 쪽으로 스스로 drift 하도록 만드는 모델.