정규화 궤적 모델: 정확한 우도 기반의 효율적인 소수 단계 이미지 생성

기존의 소수 단계 이미지 생성 기법들은 속도를 위해 확률적 우도 프레임워크를 포기하여 학습 안정성과 품질에 한계가 있었다. NTM은 각 역과정을 가역적인 Normalizing Flow로 모델링하여 정확한 우도 학습을 유지하면서도 단 4번의 샘플링만으로 고품질 이미지를 생성한다.

Diffusion 모델은 노이즈에서 데이터를 복원할 때 수많은 작은 가우시안 단계를 거친다고 가정한다. 하지만 샘플링 단계를 줄이면 각 단계가 커지면서 실제 역확률 분포는 가우시안이 아닌 복잡하고 다봉형(multimodal)인 형태를 띠게 되어 기존 방식으로는 품질이 급격히 저하된다.

NTM은 이 문제를 해결하기 위해 각 단계를 단순한 가우시안이 아닌 'Normalizing Flow'라는 가역적 변환으로 정의한다. 이는 복잡한 분포를 단순한 가우시안 분포로 매핑하거나 그 반대로 변환할 수 있는 수학적 도구이다. 구체적으로 Transporter라는 가역 블록이 데이터 공간을 잠재 공간으로 변환하면, 그 공간에서 Predictor가 다음 단계의 분포를 예측한다.

결과적으로 NTM은 단계가 커져도 비가우시안(non-Gaussian) 특성을 정확하게 포착할 수 있다. 이는 기존의 Flow Matching 모델이 수백 단계가 필요했던 품질을 단 4단계만으로 구현할 수 있게 하며, 수학적으로 정확한 확률 밀도 계산이 가능해 학습이 매우 안정적이다.

NTM은 역과정 p(xs | xt)를 조건부 Normalizing Flow로 정형화한다. 전체 구조는 공유되는 Transporter(fT)와 Predictor(fP)로 구성된다. Transporter는 입력 이미지 xs와 xt를 잠재 공간 u로 매핑하며, 이때 가역성을 보장하기 위해 TarFlow 스타일의 가역 블록을 사용한다.

Predictor는 노이즈가 섞인 표현 ut로부터 깨끗한 표현 us를 예측하는 역할을 수행한다. Predictor는 궤적 차원에서 작동하는 비인과적(non-causal) Transformer 구조를 채택하여 모든 시간 단계의 상관관계를 병렬로 처리한다. 학습 시에는 NTM 손실 함수를 최소화하는데, 이는 Predictor의 예측 오차와 Transporter 블록의 Jacobian 행렬식을 포함한 정확한 네거티브 로그 우도(NLL)이다.

[입력 이미지 x → Transporter 가역 변환 → 잠재 변수 u 생성] 과정을 거치며, [u와 시간 t → Predictor 연산 → 다음 단계 u_hat 예측]을 수행한다. 최종적으로 [u_hat → Transporter 역변환 → 생성 이미지 x_hat]을 얻는다. 이 과정에서 계산된 Jacobian 행렬식은 확률 밀도의 변화량을 보정하여 모델이 정확한 데이터 분포를 학습하게 한다.

NTM은 256x256 해상도의 ImageNet 클래스 조건부 생성 실험에서 16단계 샘플링 시 2.80 FID를 기록하며, 수백 단계가 필요한 기존 Normalizing Flow 모델인 STARFlow(5.56 FID)를 크게 앞질렀다. 4단계 샘플링만으로도 3.83 FID를 달성하여 효율성을 입증했다.

텍스트-이미지 생성 벤치마크인 GenEval에서 NTM은 4단계 샘플링만으로 0.82점을 기록했다. 이는 SDXL(0.55)이나 FLUX.1-dev(0.66)와 같은 강력한 Diffusion 기반 모델들의 소수 단계 결과보다 높은 수치이며, 복잡한 프롬프트에 대한 충실도와 이미지 품질이 우수함을 보여주었다.

NTM 아키텍처는 3072 차원의 Hidden Dimension을 가진 3개의 블록으로 구성된다. 각 블록 내에서 Transporter는 4개의 레이어를 가진 2개의 가역 블록을 사용하며, Predictor는 24개 레이어의 Transformer를 사용한다. Predictor는 궤적 전체를 입력으로 받아 시간 단계 간의 정보를 교환한다.

학습 시에는 사전 학습된 FLUX.2-klein(4B) 모델을 초기값으로 사용하는 파인튜닝 레시피를 제공한다. 이때 Mean-alignment Auxiliary Loss를 추가하여 학습 초기 단계에서 모델이 기존 Diffusion 솔루션에서 너무 멀어지지 않도록 고정(anchor)한다. 샘플링 시에는 Classifier-free Guidance를 적용하며, 이를 위해 가우시안 커플링 파라미터인 평균과 분산을 직접 조정하는 Logits-guided formulation을 사용한다.

NTM은 단일 단계(T=1) 생성 시에는 성능이 급격히 저하되는 한계를 보인다. 이는 얕은 Transporter 용량이 데이터 분포 전체를 한 번에 캡처하기에 부족하기 때문이며, 4~8단계가 성능과 속도의 최적 지점(sweet spot)으로 분석되었다.