Technische Samenvatting: Contrastive Conditional–Unconditional Alignment voor Long-tailed Diffusion Modellen

Probleemstelling

De trainingsdata voor klasse-conditionele beeldgeneratie vertoont vaak een long-tailed distributie, gekenmerkt door een schaarste aan afbeeldingen voor "tail" klassen vergeleken met "head" klassen. Deze onbalans zorgt ervoor dat diffusiemodellen lijden onder mode collapse en verminderde diversiteit bij het genereren van afbeeldingen voor deze ondervertegenwoordigde klassen. Bestaande methoden, waaronder standaard denoising diffusion probabilistic models (DDPM) en score-based diffusion models (SBDM), slagen er vaak niet in om hoogwaardige en diverse afbeeldingen te genereren voor deze ondervertegenwoordigde klassen zonder de kwaliteit van de generatie van head klassen in gevaar te brengen. Hoewel contrastief leren effectief is gebleken voor long-tailed herkenning en GAN-gebaseerde generatie, vereist de toepassing ervan op diffusiemodellen voor klasse-imbalans generatie een nieuwe adaptatie om de specifieke mechanica van het denoising-proces aan te pakken.

Methodologie: Contrastive Conditional–Unconditional Alignment (CCUA)

De auteurs stellen CCUA voor, een framework bestaande uit twee synergetische verliesfuncties die ontworpen zijn om de diversiteit en getrouwheid (fidelity) van tail klasse afbeeldingen te verbeteren, terwijl de kwaliteit van head klassen behouden blijft. De methode is compatibel met zowel U-Net-gebaseerde architecturen als Diffusion Transformers (DiT/SiT).

1. Unsupervised Contrastive Loss (UCL) voor Unconditional Generatie

Om mode collapse aan te pakken, introduceren de auteurs een unsupervised contrastive loss toegepast op unconditional generatie latents.

• Mechanisme: Het noise estimation netwerk wordt gesplitst in een latent encoder $e(\cdot)$ en een decoder $d(\cdot)$. De encoder brengt ruizige afbeeldingen $x_t$ in kaart naar laag-dimensionale latents $h$.
• Verliesformulering: De loss gebruikt een InfoNCE-stijl formulering met alleen negatieve samples. Binnen een mini-batch dient elke afbeelding als een anchor, en alle andere afbeeldingen in de batch dienen als negatieve samples. Er worden geen positieve paren gecreëerd via data-augmentatie; de anchor wordt behandeld als zijn eigen positieve vector.
• Doelstelling: Dit creëert een "afstotende kracht" in de latent space, waarbij de afstand tussen synthetische afbeeldingen wordt gemaximaliseerd, ongeacht hun klasse. Dit dwingt het model om de unconditional latents te diversifiëren, wat voorkomt dat ze instorten naar een constante vector of overfitten op beperkte trainingssamples.
• Theoretische Rechtvaardiging: Het direct toepassen van contrastieve loss op conditional latents kan leiden tot een triviale oplossing waarbij het model de inter-klasse diversiteit maximaliseert (het scheiden van klassen) terwijl de intra-klasse diversiteit wordt genegeerd. Door de loss toe te passen op unconditional latents, wordt het model gedwongen alle samples te diversifiëren, wat vervolgens wordt gedestilleerd naar de conditional generatie via de tweede loss.

2. Alignment Loss (AL) voor Conditional–Unconditional Alignment

De tweede component brengt de conditional generatie in lijn met de gediversifieerde unconditional generatie.

• Mechanisme: De loss minimaliseert de KL-divergentie tussen de conditional distributie $p_\theta(x_{t-1}|x_t, c)$ en de unconditional distributie $p_\theta(x_{t-1}|x_t)$. In de praktijk wordt dit vereenvoudigd tot een MSE-loss tussen de geschatte noise $\epsilon_\theta(x_t, t, c)$ en $\epsilon_\theta(x_t, t, \emptyset)$.
• Tijdsgewogen (Time-Weighting): Cruciaal is dat deze alignment zwaarder wordt gewogen bij grote tijdstappen (de initiële stappen van het reverse proces).
• Motivatie: De auteurs observeren dat tijdens de initiële denoising-stappen de laagfrequente componenten van afbeeldingen uit verschillende klassen (en tussen conditional en unconditional generatie) zeer vergelijkbaar zijn. Door deze initiële stappen uit te lijnen, faciliteert het model kennisdeling van head klassen (overvloedige data) naar tail klassen (schaarse data). Dit stelt tail klassen in staat om te profiteren van de rijke diversiteit die door head klassen is geleerd tijdens de vroege stadia van generatie.

3. Synergie

Het framework combineert de standaard DDPM loss met de voorgestelde CCUA loss:
$$\mathcal{L}_{total} = \mathcal{L}_{ddpm} + \alpha \cdot \mathcal{L}_{ucl} + \gamma \cdot \mathcal{L}_{al}$$
De synergie ligt in het feit dat UCL de unconditional latent space diversifieert, en AL de conditional latent space uitlijnt met deze gediversifieerde unconditional space. Deze indirecte diversificatie van conditional latents is empirisch effectiever gebleken dan het direct toepassen van contrastieve loss op conditional latents.

Belangrijkste Bijdragen

1. CCUA Framework: Een nieuwe methode voor klasse-imbalante diffusiemodellen die principes van long-tailed GANs (specifiek unconditional training op lagere resoluties/vroege stadia) combineert met een conditional-unconditional alignment loss.
2. Alignment Strategie: Het voorstel om conditional en unconditional generatie specifiek uit te lijnen bij grote tijdstappen (initiële denoising-stappen) om de gelijkenis van laagfrequente componenten te benutten en kennisoverdracht van head naar tail klassen mogelijk te maken.
3. Unsupervised Contrastive Regularisatie: Het gebruik van een unsupervised contrastive loss met alleen negatieve samples op unconditional latents om zowel intra-klasse als inter-klasse diversiteit te maximaliseren zonder klasse-labels nodig te hebben voor de contrastieve paren.
4. Architectuur-agnostisch: De methode is aangetoond effectief te zijn op zowel U-Net-gebaseerde architecturen als Diffusion Transformers (SiT).

Experimentele Resultaten

De methode werd geëvalueerd op de ImageNet-LT (256×256), TinyImageNet-LT (64×64), Places-LT (64×64) en CIFAR-LT datasets.

• Kwantitatieve Prestaties:
  • Op ImageNet-LT presteert CCUA significant beter dan baselines inclusief SiT, CBDM, REPA en Dispersive Loss.
  • Bij 900k trainingsstappen bereikte CCUA een FID van 15.1 (vs. 19.9 voor SiT) en een FIDtail van 22.5 (vs. 33.9 voor SiT), wat een verbetering van ~11.4 punten betekent voor de getrouwheid van de tail klasse.
  • De methode verbeterde ook de Recall (diversiteit) en Precision (getrouwheid) metrieken over de hele linie.
  • Op TinyImageNet-LT behaalde CCUA de beste FID (12.0) en FIDtail (20.8), waarmee het DDPM, CBDM en OCLT versloeg.
• Kwalitatieve Resultaten: Visualisaties tonen aan dat CCUA afbeeldingen genereert met een hogere diversiteit en getrouwheid voor tail klassen (bijv. 'red wine', 'espresso', 'window shade') vergeleken met standaard SiT, dat vaak repetitieve of kwalitatief minder goede afbeeldingen produceert voor deze klassen.
• Ablatie-studies:
  • De combinatie van beide losses is noodzakelijk; het gebruik van slechts één van beide levert suboptimale resultaten op.
  • Het toepassen van de alignment loss bij grote tijdstappen (lineaire weging) is superieur aan constante of kwadratische weging.
  • Het gebruik van latent representaties van de N/4-de block van de transformer biedt over het algemeen de beste balans tussen algemene FID en ruimtelijke FID/Recall.

Betekenis en Claims

Het artikel claimt dat CCUA effectief het long-tailed generatieprobleem aanpakt door principes van long-tailed GANs (specifiek unconditional training op lagere resoluties/vroege stadia) aan te passen aan het diffusieparadigma. Door de diversificatietaak (behandeld door unconditional contrastive learning) te ontkoppelen van de conditioningtaak (behandeld door alignment), vermijdt de methode de valkuilen van directe conditional contrastive regularisatie.

De auteurs benadrukken dat hun aanpak:

• De generatie van tail klassen verbetert zonder de kwaliteit van head klassen te verslechteren.
• Eenvoudig te implementeren is binnen bestaande DDPM en SBDM pipelines.
• Presteert beter dan gelijktijdige werken zoals Dispersive Loss en DiffROP, met name in scenario's met extreme klasse-imbalans (bijv. ImageNet-LT met een imbalansfactor van 0.001).
• Een theoretische en empirische rechtvaardiging biedt voor waarom het uitlijnen van conditional en unconditional generatie tijdens de initiële stappen van het reverse proces gunstig is voor kennisdeling in long-tailed settings.