일단계 생성 모델은 디노이징 궤적을 깊이 방향으로 펼친다

무엇인가

일단계 생성 모델은 증류나 학습된 플로우 맵을 통해 다단계 확산의 반복 샘플링 궤적을 한 번의 순전파로 압축해 고품질 이미지를 만든다. 이 논문이 던지는 질문은 그 압축 이후에 원래의 디노이징 궤적이 어디로 갔느냐는 것이다. 다단계 확산에서 이 궤적은 거친 구조에서 세밀한 디테일로 나아가는 중간 상태의 열이고, classifier-free guidance나 강화학습 기반 제어, 이미지 편집 같은 후속 기법들이 개입하는 축이다. 일단계 모델은 이 외부 궤적을 없애는 것처럼 보이지만, 저자들은 그 계산이 사라진 게 아니라 네트워크 깊이 방향으로 재조직된다고 주장한다.

어떻게 동작하나

방법의 핵심은 학습 없이 적용 가능한 프로브다. 잔차 스트림을 쓰는 일단계 생성기는 입력 임베딩 h0에서 시작해 L개의 블록을 순차 적용해 hL을 만들고, 최종 예측은 출력 헤드 W_out(hL; τ_task, c)로 나온다. 저자들은 남은 블록을 실행하지 않은 채 각 중간 은닉 상태 hi에 같은 출력 헤드를 그대로 적용해 층별 속도 vi = W_out(hi; τ_task, c)를 얻는다. 이걸 다시 엔드포인트 예측으로 환산한 것이 x̃_r^i = x_t − (t−r)·vi이며, r=0이면 깨끗한 이미지 추정, 0<r<t이면 중간 노이즈 수준의 추정이 된다. 이 프로브는 학습 파라미터를 추가하지 않고 생성기를 수정하지도 않는다. 언어 모델의 logit lens와 유사한 발상이며, 저자들은 별도로 학습한 선형 프로브로도 같은 결론이 나오지만 학습되지 않은 랜덤 프로브로는 나오지 않는다는 점을 부록에서 확인해 이 현상이 단순 수치적 현상이 아니라 의미 있는 수렴임을 보인다. 층별 진행을 정량화하기 위해 각 중간 예측과 최종 예측의 거리를 궤적 내 최대값으로 나눈 ρ(i)를 쓰고, 나아가 입력 노이즈와 최종 예측의 혼합으로 중간 예측을 설명하는 계수 (α_i, σ_i)를 궤적 256개에 공유해 최소제곱으로 적합해 '깊이별 유효 혼합 스케줄'을 복원한다.

무엇과 다른가

실험은 FLUX-schnell(잠재 적대 증류), MeanFlow B/2(12층, 폭 768), Drifting B/2, Shortcut B/2, SANA-Sprint 0.6B를 대상으로 ImageNet 256×256과 CIFAR-10에서 이뤄졌다. FLUX-schnell과 MeanFlow는 깊이를 따라 점진적으로 정제되는 패턴을 보이고 ρ(i)도 대체로 단조 감소하는데, 시간 인덱스가 붙은 전송 과제 없이 학습된 Drifting은 같은 체계적 정제를 보이지 않는다. 혼합 계수 적합에서도 Drifting은 모든 블록에서 노이즈 계수가 σ_i ≈ 0으로 나와 중간 예측에 입력 노이즈 성분이 식별되지 않는다. 반면 MeanFlow는 블록 4에서 α_4 = 2.18, σ_4 = −6.04처럼 두 계수가 [0,1] 밖으로 나오고 α_i + σ_i = 1 제약도 성립하지 않는데, 블록 평균 재구성 오차는 약 0.10으로 낮다. 즉 깊이는 학습 시의 보간 스케줄을 그대로 재현하는 게 아니라 자체적인 유효 디노이징 좌표를 정의한다.

어떻게 쓰나

이 현상은 임의의 전송 구간으로 일반화된다. 증류된 일단계 모델은 엔드포인트 r을 명시적으로 받지 않기 때문에 외부 샘플러가 t→r 부분 과제를 만들고, 각 샘플링 스텝 안에서 다시 깊이별 정제가 일어나 깊이가 시간 안에 중첩된다. FLUX-schnell을 두 단계로 나눠 생성하면 1→0.75 구간에서도 최종 층의 표현이 이미 t=0의 깨끗한 이미지를 상당 부분 함축하고 있다. 구간 (t,r)을 명시적으로 조건받는 MeanFlow는 더 흥미롭다. 요청된 엔드포인트가 노이즈 상태이면 중간 층이 먼저 요청 대상보다 깨끗한 구조를 드러내고 이후 층에서 다시 노이즈 쪽으로 되돌아가는 비단조 'denoise-then-renoise' 패턴이 나타난다. CIFAR-10 5단계 생성에서 ρ(i)가 각 스텝 안에서 내려갔다 다시 올라가는 것으로 정량화된다. 저자들의 설명은 두 갈래다. 플로우 맵의 합성 일관성 Φ_{t→r} = Φ_{0→r} ∘ Φ_{t→0}이 z_t → z_0 → z_r 분해를 허용하지만, 이는 모든 중간 시점 s에 똑같이 성립하므로 깨끗한 엔드포인트만을 특별히 지목하지는 못한다. 결정적인 것은 학습 목표의 구조로, 선형 경로에서 주변 속도가 사후 깨끗한 추정 E[x_0|z_t]에 아핀하게 의존하고 MeanFlow 항등식 u = v − (t−r)·du/dt의 첫 항이 r과 무관하게 깨끗한 추정으로 결정되며 둘째 항은 r→t일 때 사라지는 구간 의존 보정이라는 점이다. 네트워크는 먼저 z_t가 함축하는 깨끗한 샘플을 풀고, 그다음 r에 특화된 보정을 적용하는 두 단계 계산을 재사용하게 된다.

전제와 한계

5장에서는 이 깊이별 구조를 명시적으로 모델링한다. 층마다 다른 블록 Bi를 쓰는 대신, 구간 조건 (τ_i, τ_{i+1})을 받는 단일 가중치 공유 블록 B_θ를 반복 적용하고 τ_i = i/L로 두면 L개 층에 대응하는 L단계가 된다. 학습은 데이터 없이 랜덤 노이즈, 샘플링한 클래스 레이블, 일단계 교사 모델의 은닉 표현만으로 이뤄지며 구간 조건은 AdaLN 변조로 주입된다. 결과적으로 ImageNet 256×256의 MeanFlow SiT-B/2를 131M에서 16.0M(8.2배), SiT-L/2를 459M에서 27.6M(16.6배)으로 압축해 FID-50k 15.6과 11.7을 얻고, FD-loss 후학습을 거치면 8.2와 4.9까지 개선된다. 원본 모델의 FID는 각각 6.2와 4.0이다. 고유 블록을 두 개 쓰면 압축률을 일부 포기하는 대신 SiT-B/2가 26.6M(4.9배)에서 FID 5.84, SiT-L/2가 46.5M(9.9배)에서 후학습 후 4.04를 기록한다. 깊이 방향도 확산의 샘플링 스텝처럼 다룰 수 있어서, 여러 층을 아우르는 거친 구간으로 학습하면 블록 적용 횟수를 줄일 수 있다. SiT-L/2에 고유 블록 두 개를 쓰고 적용 횟수를 절반으로 줄이면 FID 5.25, 3분의 1로 줄이면 6.11이 나오며, 교사 모델은 후학습 후 4.01이다. 반대로 같은 압축 절차를 Drifting에 적용하면 FID-50k 47.09로, 더 강한 교사(FID 1.75)에서 증류했음에도 MeanFlow SiT-B/2 학생(15.61)보다 약 3배 나쁘다. CIFAR-10에서 구간 조건을 제거한 절제 실험에서는 FID-50k가 두 배 이상 나빠진다.

개발자 관점에서 이 논문이 주는 실무적 시사점은 두 가지다. 첫째, 일단계 생성 모델의 중간 레이어를 자체 출력 헤드로 디코딩하는 것만으로 별도 학습 없이 생성 품질이 어느 깊이에서 어느 정도로 수렴하는지 진단할 수 있다. 둘째, 시간 인덱스가 붙은 전송 과제로 학습된 모델이라면 층별 계산을 가중치 공유 블록으로 접어 파라미터와 추론 연산을 함께 줄일 수 있고, 깊이 스텝 수를 확산의 샘플링 스텝처럼 품질-연산 트레이드오프 노브로 쓸 수 있다. 다만 이 압축이 통하는지는 모델 계열에 달려 있다. 시간 인덱스 전송 과제가 없는 드리프팅 계열은 깊이를 따라 예측이 날카로워지더라도 압축 가능한 디노이징 구조를 형성하지 않는다.

저자들이 밝힌 전제와 한계도 분명하다. 프로브는 파라미터가 없고 생성기를 수정하지 않지만, 학습된 선형 프로브로는 같은 결론이 나오는 반면 학습되지 않은 랜덤 프로브로는 나오지 않으므로 관찰된 수렴의 의미를 프로브 형태에 의존하지 않는다는 점을 별도로 확인해야 한다. 복원된 깊이별 혼합 스케줄은 α_i, σ_i가 [0,1]을 벗어나고 α_i + σ_i = 1도 성립하지 않아, 학습 시 사용한 플로우 매칭 보간 스케줄과 동일하지 않다. 합성 일관성만으로는 denoise-then-renoise를 완전히 설명하지 못하며, 깨끗한 엔드포인트를 특별히 지목하는 것은 학습 목표의 구조라는 것이 저자들의 해석이다. 또한 CIFAR-10 실험은 대부분의 확산·일단계 코드베이스가 DDPM++/NCSN++ 계열 U-Net을 쓰기 때문에 저자들이 직접 DiT-B/2를 학습해 수행했으며, 드리프팅 모델은 시간 인덱스 전송 과제가 없는 대조군으로 사용되었다.