Calibri: Enhancing Diffusion Transformers via Parameter-Efficient Calibration

Probleemstelling

Diffusion Transformers (DiTs) zijn de nieuwe standaard in generatieve AI-modellen voor tekst-naar-beeld taken (zoals Stable Diffusion 3 en FLUX). Hoewel deze modellen krachtig zijn, blijkt uit recente analyses dat hun architectuur suboptimaal is afgesteld.

• Ongelijke bijdrage: Niet alle lagen binnen een DiT dragen evenveel bij aan de uiteindelijke kwaliteit. Sommige lagen kunnen zelfs schadelijke artefacten introduceren.
• Inefficiëntie: Bestaande methoden om generatieve modellen te aligneren met menselijke voorkeuren (zoals RLHF, DPO of DDPO) vereisen vaak volledige fine-tuning van het model, wat computatief zeer duur en tijdrovend is.
• Aantal inferentiestappen: Huidige modellen vereisen vaak veel stappen (NFE - Number of Function Evaluations) om hoogwaardige beelden te genereren, wat de inferentie vertraagt.

Het paper stelt de hypothese dat de standaard DiT-architectuur een eenvoudige, post-hoc kalibratie kan ondergaan om de prestaties aanzienlijk te verbeteren zonder de zware kosten van volledige hertraining.

Methodologie: Calibri

Calibri is een parameter-efficiënte methode die de bijdrage van DiT-componenten optimaliseert door slechts een zeer klein aantal parameters aan te passen (ongeveer $10^2$ parameters).

1. Kerninzicht en Motivatie
De auteurs ontdekten via experimenten dat:

• Het uitschakelen van bepaalde blokken soms de beeldkwaliteit verbetert.
• Het vermenigvuldigen van de output van elk DiT-blok met een enkele, geleerde scalair (schaalfactor) de prestaties consistent verbetert.

2. Het Optimalisatieprobleem
Calibri vormt de kalibratie als een black-box optimalisatieprobleem. Het doel is om een vector van kalibratiecoëfficiënten $c$ te vinden die een beloningsfunctie $R(c)$ maximaliseert (bijv. gemeten door een reward model zoals HPSv3 of ImageReward).

• Zoekruimte: De parameters kunnen worden toegepast op drie niveaus van granulariteit:
  • Block Scaling: Eén schaalfactor per architecturaal blok (Attention + MLP).
  • Layer Scaling: Afzonderlijke schaalfactoren voor elke laag binnen een blok.
  • Gate Scaling: Specifieke factoren voor multimodale interacties (bijv. visuele vs. tekstuele tokens in MM-DiT).

3. Zoekalgoritme
Om de optimale coëfficiënten te vinden, gebruikt Calibri de Covariance Matrix Adaptation Evolution Strategy (CMA-ES).

• Dit is een gradient-free (gegradiëntvrije) evolutionaire strategie.
• Het werkt door iteratief een verdeling van kandidaat-oplossingen te verfijnen op basis van een multivariate Gaussische verdeling.
• Dit maakt het mogelijk om de optimale schaalfactoren te vinden zonder backpropagation door het hele zware model.

4. Calibri Ensemble
De methode introduceert ook een ensemble-aanpak waarbij meerdere gekalibreerde modellen tegelijkertijd worden gebruikt. Dit kan worden gezien als een generalisatie van Skip Layer Guidance en integreert naadloos met Classifier-Free Guidance (CFG). Door de output van meerdere gekalibreerde instanties te combineren, wordt de robuustheid en kwaliteit verder verhoogd.

Belangrijkste Bijdragen

1. Parameter-efficiëntie: Calibri verbetert de prestaties door slechts ongeveer 100 parameters aan te passen, in tegenstelling tot de miljoenen parameters die bij traditionele fine-tuning worden bijgewerkt.
2. Verbeterde Kwaliteit en Efficiëntie: De methode levert niet alleen hogere kwaliteitscores op, maar verlaagt ook drastisch het aantal inferentiestappen dat nodig is voor het genereren van beelden.
3. Combinatie met Alignment: Calibri kan worden gecombineerd met bestaande alignment-methoden (zoals Flow-GRPO) om de prestaties van reeds getrainde modellen verder te verbeteren.
4. Generieke Toepasbaarheid: De methode werkt effectief op verschillende state-of-the-art DiT-modellen (FLUX, Stable Diffusion 3.5, Qwen-Image).

Resultaten

De experimenten tonen overtuigende resultaten aan op diverse benchmarks:

• Kwaliteitsverbetering:
  • Op het FLUX-model steeg de HPSv3-score van 11,41 naar 13,48.
  • Op Stable Diffusion 3.5 Medium steeg de HPSv3-score van 11,15 naar 14,10.
  • Menselijke evaluaties tonen aan dat gebruikers Calibri-gegenereerde beelden significant verkiezen boven de baselines (ongeveer 52-54% winst in algemene voorkeur).
• Versnelde Inferentie (NFE):
  • Calibri vermindert het aantal benodigde stappen drastisch. Waar FLUX normaal 30 stappen nodig heeft, volstaat Calibri met 15 stappen voor een betere kwaliteit.
  • SD-3.5M daalt van 80 naar 30 stappen, en Qwen-Image van 100 naar 30 stappen.
• Kosten:
  • De kalibratie is een eenmalige, offline kostenpost. Het kalibreren van FLUX kost bijvoorbeeld slechts 32 H100 GPU-uren, wat resulteert in een permanente 2-voudige versnelling tijdens de inferentie.
• Vergelijking met Fine-tuning:
  • Calibri bereikt vergelijkbare of betere resultaten dan Flow-GRPO (een methode die 18,78 miljoen parameters bijwerkt), terwijl Calibri slechts 216 parameters aanpast.

Betekenis en Impact

Calibri biedt een praktische en kosteneffectieve oplossing voor het optimaliseren van moderne generatieve AI-modellen. Het bewijst dat grote prestatieverbeteringen niet per se vereisen dat het hele model opnieuw wordt getraind of zwaar wordt gefinetuned.

• Toepassing in de praktijk: Door de reductie in inferentiestappen en de lage trainingskosten, maakt Calibri hoogwaardige beeldgeneratie toegankelijker en sneller voor real-world toepassingen.
• Paradigmaverschuiving: Het paper benadrukt dat de "suboptimale weging" van bestaande DiT-architecturen een grote bron van potentieel is die vaak over het hoofd wordt gezien, en dat eenvoudige post-hoc kalibratie een krachtig alternatief is voor zware trainingstechnieken.