Analyseren en verbeteren van de trainingsdynamiek van diffusiemodellen

Analyzing and Improving the Training Dynamics of Diffusion Models

December 5, 2023

Auteurs: Tero Karras, Miika Aittala, Jaakko Lehtinen, Janne Hellsten, Timo Aila, Samuli Laine

cs.AI

Samenvatting

Diffusiemodellen domineren momenteel het veld van data-gedreven beeldgeneratie met hun ongeëvenaarde schaalbaarheid naar grote datasets. In dit artikel identificeren en corrigeren we verschillende oorzaken van ongelijkmatige en inefficiënte training in de populaire ADM-diffusiemodelarchitectuur, zonder de hoog-niveau structuur ervan aan te passen. Door ongecontroleerde grootteveranderingen en onevenwichtigheden in zowel de netwerkactivaties als de gewichten gedurende de training te observeren, herontwerpen we de netwerklagen om de grootte van activaties, gewichten en updates op verwachting te behouden. We constateren dat systematische toepassing van deze filosofie de waargenomen verschuivingen en onevenwichtigheden elimineert, wat resulteert in aanzienlijk betere netwerken bij gelijke rekencomplexiteit. Onze aanpassingen verbeteren het vorige FID-record van 2.41 in ImageNet-512-synthese naar 1.81, bereikt met snelle deterministische steekproefname. Als een onafhankelijke bijdrage presenteren we een methode voor het instellen van de exponentiële voortschrijdende gemiddelde (EMA) parameters achteraf, d.w.z. na het voltooien van de trainingsrun. Dit maakt een precieze afstemming van de EMA-lengte mogelijk zonder de kosten van het uitvoeren van meerdere trainingsruns, en onthult de verrassende interacties ervan met netwerkarchitectuur, trainingsduur en begeleiding.

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.