Effective and Efficient Masked Image Generation Models

Probleemstelling

Bestaande modellen voor beeldgeneratie kampen met verschillende beperkingen:

• Autoregressieve modellen (AR) zijn traag bij het genereren van beelden.
• Diffusiemodellen (vooral continue diffusie) bieden hoge kwaliteit maar vereisen vaak veel rekenkracht en stappen (NFE - Number of Function Evaluations) voor inferentie.
• Masked Image Generation (MIG) modellen zoals MaskGIT zijn efficiënter, maar lijden vaak onder kwaliteitsverlies door informatieverlies bij discrete tokenisatie.
• Masked Diffusion Models (MDM) hebben succesvol getoond in tekstgeneratie, maar hun toepassing op beeldgeneratie en de optimale trainings- en samplingstrategieën waren nog niet volledig onderzocht.

De auteurs stellen dat er een kloof bestaat tussen MIG en MDM, en dat er behoefte is aan een unificatie die zowel de efficiëntie van MIG als de kwaliteit van diffusiemodellen combineert, zonder de nadelen van discrete tokenisatie.

Methodologie

De kern van het paper is het voorstellen van een unificerend raamwerk dat Masked Image Generation en Masked Diffusion Models samenvoegt.

1. Unificatie van MIG en MDM:
   De auteurs tonen aan dat de trainingsdoelstellingen van MaskGIT en MDM wiskundig equivalent zijn binnen één formule. Het verschil zit hem in drie componenten:

   • De verdeling van het maskeren ($q(x_t|x_0)$).
   • De weegfunctie ($w(t)$).
   • De conditionele verdeling ($p_\theta(x_0|x_t)$).
     Ze kiezen voor de MDM-benadering voor het maskeren (onafhankelijk maskeren) en gebruiken een diffusiemodel om de conditionele verdeling te modelleren. Dit elimineert de noodzaak voor discrete tokenizers en voorkomt informatieverlies.

2. Optimalisatie van het Trainingsontwerp:
   Door het ontwerpruimte te verkennen, identificeren de auteurs de volgende optimale strategieën:

   • Masker-schema: Een exponentiële schedule ($\gamma_t = 1 - \exp(-5t)$) presteert beter dan lineaire of cosinus-schema's, omdat deze een hoger maskerpercentage tijdens het trainen mogelijk maakt.
   • Weegfunctie: In tegenstelling tot standaard MDM, werkt een constante weegfunctie ($w(t)=1$) beter in combinatie met de exponentiële schedule, wat training stabiliseert.
   • Architectuur: De MAE-architectuur (Masked Autoencoder) wordt gebruikt, waarbij de encoder alleen ongemaskerde tokens verwerkt. Dit verbetert de prestaties ten opzichte van een standaard single-encoder transformer.
   • Tijdsvertraging (Time Truncation): Het beperken van het trainingsinterval tot $t_{min} = 0.2$ versnelt de convergentie zonder prestatieverlies.
   • CFG met Mask: Voor conditionele generatie wordt in plaats van een "fake class token" een mask token gebruikt voor het onvoorwaardelijke model. Dit verbetert de prestaties aanzienlijk.

3. Optimalisatie van het Sampling-ontwerp:

   • Sampling-schema: Tijdens het genereren (sampling) wordt ook de exponentiële schedule gebruikt, omdat deze in de vroege stadia minder tokens voorspelt, wat de kwaliteit bij een laag aantal stappen verbetert.
   • DPM-Solver: In plaats van DDPM wordt de DPM-Solver gebruikt. Dit is een snellere ODE-solver die minder stappen nodig heeft voor goede resultaten.
   • Tijdsinterval voor Classifier-Free Guidance (CFG): Dit is een cruciale innovatie. De auteurs stellen dat sterke guidance in de vroege stadia de variatie vermindert en de FID (Fréchet Inception Distance) verslechtert omdat tokens in MDM onomkeerbaar zijn. Ze introduceren een tijdsinterval-strategie waarbij CFG alleen wordt toegepast in de latere stadia van het samplingproces (bijv. tussen $t=0.1$ en $t=0.3$). Dit behoudt de kwaliteit terwijl het aantal rekenstappen (NFE) drastisch wordt verlaagd.

Belangrijkste Bijdragen

• Unificerend Raamwerk: Een theoretische basis die MIG en MDM verenigt, waardoor een systematische exploratie van trainings- en samplingstrategieën mogelijk wordt.
• eMIGM Model: De ontwikkeling van een nieuw model (eMIGM) dat voortbouwt op deze inzichten en state-of-the-art prestaties levert.
• Tijdsinterval-strategie voor CFG: Een nieuwe methode om guidance toe te passen die de sampling-efficiëntie verhoogt zonder kwaliteitsverlies, specifiek ontworpen voor de irreversibele aard van token-generatie.
• Schalingswetten: Het aantonen dat eMIGM profiteert van schaling; grotere modellen zijn niet alleen beter, maar ook efficiënter in termen van training-FLOPs en sampling-tijd.

Resultaten

Het model (eMIGM) werd geëvalueerd op ImageNet bij resoluties van 256x256 en 512x512, gemeten aan de hand van FID.

• ImageNet 256x256:

  • eMIGM presteert beter dan de seminale VAR-modellen (Visual Autoregressive) met een vergelijkbaar aantal parameters en NFE.
  • De grootste variant (eMIGM-H) bereikt een FID van 1.57 met slechts 180 NFE. Dit is vergelijkbaar met de state-of-the-art continue diffusiemodel REPA (FID 1.42), maar REPA vereist meer stappen en zelf-supervisie.
  • Het presteert ook beter dan sterke GAN-modellen zoals StyleGAN-XL.

• ImageNet 512x512:

  • eMIGM-L (478M parameters) bereikt een FID van 1.77 met 80 NFE.
  • Dit is beter dan het sterke continue diffusiemodel EDM2 (FID 1.81), ondanks dat EDM2 meer parameters heeft.
  • Het model overtreft ook VAR-d36 (2.3B parameters) met een veel kleiner model.

• Efficiëntie:

  • eMIGM vereist minder dan 45% van de NFE van de beste continue diffusiemodellen om vergelijkbare resultaten te bereiken.
  • Grotere modellen (eMIGM-L, eMIGM-H) behalen betere kwaliteit bij gelijke inferentietijd, wat aantoont dat schaling de efficiëntie verbetert.

Betekenis

Dit paper is significant omdat het de kloof tussen efficiënte masked modellen en hoogwaardige diffusiemodellen dicht. Door een unificerend raamwerk te bieden en specifieke strategieën (zoals de tijdsinterval-CFG en het gebruik van MAE-architecturen) te introduceren, toont het aan dat masked image generation een zeer competitief alternatief is voor continue diffusie.

eMIGM biedt een nieuwe richting voor snelle, hoogwaardige beeldgeneratie die minder rekenkracht vereist, wat essentieel is voor de schaalbaarheid en praktische toepasbaarheid van generatieve AI. De code is open-source beschikbaar gesteld, wat verdere onderzoek en ontwikkeling in dit domein zal stimuleren.