확산 모델의 학습 역학 분석 및 개선
Analyzing and Improving the Training Dynamics of Diffusion Models
December 5, 2023
저자: Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, Samuli Laine
cs.AI
초록
확산 모델(Diffusion models)은 현재 대규모 데이터셋에 대한 비할 데 없는 확장성으로 데이터 기반 이미지 합성 분야를 주도하고 있습니다. 본 논문에서는 널리 사용되는 ADM(ADM diffusion model) 아키텍처에서 발생하는 불균일하고 비효율적인 학습의 여러 원인을 파악하고 이를 수정하는 방법을 제안합니다. 이때, 고수준 구조를 변경하지 않고도 이러한 개선을 달성합니다. 학습 과정에서 네트워크 활성화 및 가중치의 통제되지 않은 크기 변화와 불균형을 관찰한 후, 우리는 네트워크 계층을 재설계하여 활성화, 가중치, 업데이트 크기를 기대값 기준으로 보존하도록 했습니다. 이러한 철학을 체계적으로 적용함으로써 관찰된 드리프트(drift)와 불균형을 제거할 수 있었으며, 동일한 계산 복잡도에서 훨씬 더 나은 네트워크를 얻을 수 있었습니다. 우리의 수정 사항은 ImageNet-512 합성에서 이전 기록인 FID 2.41을 빠른 결정론적 샘플링을 사용하여 1.81로 개선했습니다.
독립적인 기여로서, 우리는 학습 완료 후 사후적으로(ex post-hoc) 지수 이동 평균(EMA, Exponential Moving Average) 매개변수를 설정하는 방법을 제시합니다. 이를 통해 여러 번의 학습 실행 비용 없이 EMA 길이를 정밀하게 조정할 수 있으며, 네트워크 아키텍처, 학습 시간, 가이던스(guidance)와의 놀라운 상호작용을 밝힐 수 있습니다.
English
Diffusion models currently dominate the field of data-driven image synthesis
with their unparalleled scaling to large datasets. In this paper, we identify
and rectify several causes for uneven and ineffective training in the popular
ADM diffusion model architecture, without altering its high-level structure.
Observing uncontrolled magnitude changes and imbalances in both the network
activations and weights over the course of training, we redesign the network
layers to preserve activation, weight, and update magnitudes on expectation. We
find that systematic application of this philosophy eliminates the observed
drifts and imbalances, resulting in considerably better networks at equal
computational complexity. Our modifications improve the previous record FID of
2.41 in ImageNet-512 synthesis to 1.81, achieved using fast deterministic
sampling.
As an independent contribution, we present a method for setting the
exponential moving average (EMA) parameters post-hoc, i.e., after completing
the training run. This allows precise tuning of EMA length without the cost of
performing several training runs, and reveals its surprising interactions with
network architecture, training time, and guidance.