시각적 생성을 위한 표현 프레셰 손실

이미지 생성 모델의 표준 평가 지표인 FID를 평가용이 아닌 학습용 손실 함수(Loss)로 직접 사용하여 모델 성능을 개선하는 방법론을 제시한다. 이를 통해 복잡한 증류 과정 없이도 다단계 생성 모델을 고성능 단일 단계(One-step) 생성 모델로 전환할 수 있어 추론 효율성을 획기적으로 높인다.

이미지 생성 모델의 품질을 측정할 때 가장 널리 쓰이는 FID는 실제 이미지와 생성 이미지의 특징 분포 사이의 거리를 계산한다. 하지만 이 계산에는 수만 장의 샘플이 필요하기 때문에, 수십~수백 장 단위로 업데이트되는 딥러닝의 Gradient Descent 과정에서 손실 함수로 직접 쓰기에는 통계적 노이즈가 너무 크다는 한계가 있었다.

이 논문은 '디커플링(Decoupling)' 원리를 통해 이 문제를 해결한다. 모델을 업데이트할 때 사용하는 현재 배치의 데이터뿐만 아니라, 이전 단계에서 생성된 데이터들의 특징값들을 큐(Queue)에 저장하거나 지수 이동 평균(EMA) 방식으로 누적하여 대규모 샘플 통계를 유지한다. 이렇게 확보된 안정적인 분포 정보를 바탕으로 현재 배치의 그래디언트만 역전파함으로써 FID를 직접 줄이는 학습이 가능해진다.

결과적으로 모델은 특정 이미지 한 장의 픽셀을 맞추는 것이 아니라, 전체적인 데이터 분포 자체를 실제 데이터와 일치시키도록 학습된다. 이는 모델이 단 한 번의 추론(1-NFE)만으로도 실제와 유사한 고품질 이미지를 생성할 수 있게 만드는 강력한 가이드라인이 된다.

FD-loss의 핵심은 FD 추정을 위한 모집단 크기 N(예: 50k)과 그래디언트 계산을 위한 배치 크기 B(예: 1024)를 분리하는 것이다. 두 가지 구현 방식을 제안하는데, 첫 번째는 최근 생성된 샘플의 특징을 저장하는 온라인 큐(Queue) 방식이고, 두 번째는 특징의 1차 및 2차 모멘트를 업데이트하는 지수 이동 평균(EMA) 방식이다.

[생성 이미지 x_hat 입력 → 특징 추출기 phi 연산 → 특징값 phi(x_hat) 획득 → EMA 또는 큐를 통한 통계량 mu_g, Sigma_g 계산] 과정을 거친다. 이후 미리 계산된 실제 데이터의 통계량(mu_r, Sigma_r)과의 Fréchet Distance를 계산하여 손실값을 도출한다.

다중 표현 공간(Multi-representation) 최적화를 위해 Inception-v3 외에도 DINOv2, MAE, SigLIP2 등 다양한 백본의 특징 공간에서 계산된 FD를 결합한다. 각 공간에서의 FD 값은 스케일이 다르므로, [각 공간의 FD 값 → 정지된 FD 값으로 나누는 정규화 연산 → 가중치 합산 → 최종 손실값] 순으로 처리하여 안정적인 학습을 보장한다.

ImageNet 256x256 벤치마크에서 FD-loss로 포스트 트레이닝된 단일 단계 생성 모델은 Inception 특징 공간 기준 0.72 FID를 달성했다. 이는 기존 베이스 모델의 성능을 크게 상회하는 수치이다. 또한, 다단계 모델인 JiT-H를 단일 단계로 전환했을 때 FID가 291.59에서 0.75로 급격히 개선됨을 확인했다.

Ablation study 결과, EMA 방식이 큐 방식보다 메모리 효율적이면서도 더 나은 성능(0.81 FID)을 보였다. 특히 단일 Inception 공간만 최적화할 경우 FID 수치는 낮아지지만 시각적 아티팩트가 발생할 수 있는데, MAE나 SigLIP 등 현대적인 표현 공간을 함께 최적화(FD-SIM)할 경우 시각적 품질과 물체 구조가 훨씬 더 정교해지는 것으로 나타났다.

FD-loss는 미분 가능한 Fréchet Distance 수식을 직접 최적화 목표로 삼는다. 핵심은 Sigma_g(생성 데이터 공분산)를 계산할 때 현재 배치의 그래디언트만 흐르게 하고, 통계적 추정치는 대규모 윈도우를 통해 확보하여 편향(Bias)을 줄이는 것이다. 연구팀은 Inception-v3, ConvNeXt-v2, DINOv2, MAE, SigLIP2, CLIP 등 6개의 상이한 표현 공간을 활용하여 모델의 사각지대를 최소화했다. 또한 행렬 제곱근 연산의 효율성을 위해 torch.linalg.eigvalsh를 활용한 고유값 분해 방식을 채택하여 학습 속도를 최적화했다.

FD-loss의 성능은 선택한 표현 공간(Representation space)과 참조 통계의 품질에 의존한다. 또한 6개의 표현 공간을 사용하더라도 인간의 지각적 품질을 완벽하게 대체하지는 못하며, 특정 공간에 대한 과도한 최적화가 다른 공간에서의 성능 저하를 야기할 수 있는 '보상 해킹(Reward hacking)' 가능성이 존재한다.