Beyond Generative Priors: Minority Sampling with JEPA-Guided Diffusion

Probleemstelling

Minority sampling heeft tot doel instances met een lage dichtheid op een data-maand te genereren, wat cruciaal is voor toepassingen zoals medische diagnose, anomaliedetectie en creatieve AI. Bestaande benaderingen definiëren "minderheid" echter relatief tot de generatieve prior die door het model zelf is geleerd (d.w.z. de verdeling van de trainingsdata). Dit creëert een fundamentele beperking: monsters die als "zeldzaam" worden geïdentificeerd, zijn vaak slechts artefacten van de specifieke trainingsset of de inductieve bias van het model, in plaats van echt atypische instances in de realiteit. Een model dat is getraind op standaard datasets kan bijvoorbeeld een hond op een witte achtergrond als "zeldzaam" beschouwen, simpelweg omdat deze ondervertegenwoordigd is in de trainingsdata, terwijl een stealth-vliegtuig in de echte wereld echt zeldzaam kan zijn, maar niet als zodanig wordt vastgelegd als de prior van het model te smal is. Het artikel betoogt dat huidige methoden falen in het afstemmen op priors uit de realiteit, die noodzakelijk zijn om echte semantische zeldzaamheid vast te leggen.

Methodologie: JEPA Guidance

De auteurs stellen JEPA Guidance voor, een kader dat minority sampling verschuift van een generator-gecentreerd perspectief naar een wereld-gecentreerd perspectief. Dit wordt bereikt door diffusiemodellen te sturen met behulp van de impliciete dichtheid van een Joint-Embedding Predictive Architecture (JEPA), een klasse van wereldmodellen (zoals DINOv2) die bekend staat om het coderen van brede, semantisch rijke representaties.

Kernmechanisme:

1. Wereld-gecentreerde definitie: In plaats van de waarschijnlijkheid onder de dichtheid $p_\theta(x)$ van het diffusiemodel te minimaliseren, streeft de methode ernaar de waarschijnlijkheid onder de door JEPA geïnduceerde dichtheid $p_\phi(x)$ te minimaliseren. Minority-samples worden gedefinieerd als die welke zich bevinden in gebieden met een lage dichtheid in de representatieruimte van JEPA.
2. JEPA-SCORE: De methode maakt gebruik van JEPA-SCORE, een dichtheidsmaat afgeleid van de Jacobiaan van een vooraf getrainde JEPA-encoder $f_\phi$. Specifiek is dit de som van de logaritmen van de singuliere waarden van de Jacobiaan $J_f(x)$. Lagere JEPA-SCORE-waarden duiden op een lagere dichtheid (hogere atypiciteit) in de wereldprior.
3. Benadering via Randomized SVD: Het berekenen van de volledige Jacobiaan en de Singuliere Waarde Decompositie (SVD) daarvan bij elke diffusiestap is computationeel onhaalbaar. De auteurs maken gebruik van Randomized SVD om de top-$k$ singuliere waarden van de gecomprimeerde Jacobiaan efficiënt te benaderen. Zij bieden een theoretische bovengrens voor de benaderingsfout, die wordt ontbonden in de schattingsfout van de randomized SVD en de truncatiefout.
4. Envelope Theorem voor efficiëntie: Om backpropagation door de innerlijke optimalisatielus van de randomized SVD (die afhankelijk is van de invoer $x_t$) te vermijden, roepen de auteurs het Envelope Theorem in. Dit stelt hen in staat om de optimale projectiematrix $Q^*$ als een constante te behandelen tijdens de gradiëntberekening, wat het geheugengebruik en de computationele overhead aanzienlijk vermindert terwijl de correcte eerste-orde gradiënten behouden blijven.
5. Deferred Guidance: Om het domeinverschil aan te pakken tussen de ruwe tussenliggende latents van het diffusiemodel en de schone invoer die door de JEPA-encoder wordt verwacht, introduceren de auteurs deferred guidance. JEPA guidance wordt pas toegepast na een tussenliggende tijdstap $\tau T$ (bijvoorbeeld $t=0.8T$), zodat de gedenoiseerde schattingen $\hat{x}_{0|t}$ voldoende helder zijn voor de JEPA-encoder om effectief te evalueren. Dit stelt de methode ook in staat om uit te breiden naar conditionele generatie (bijvoorbeeld tekst-naar-afbeelding), aangezien de conditie vroeg in het proces wordt vastgesteld voordat de wereld-gecentreerde guidance begint.

Belangrijkste Bijdragen

• Wereld-gecentreerde Minority Sampling: Een paradigma-shift die minority-samples herdefinieert op basis van priors uit de realiteit (via JEPAs) in plaats van de smalle priors van het generatieve model.
• JEPA Guidance Framework: Een principieel diffusie-samplingmethode die trajecten stuurt naar gebieden met een lage dichtheid onder door JEPA geïnduceerde priors.
• Efficiële Berekening: De ontwikkeling van principieel benaderingsstrategieën (Randomized SVD) en de toepassing van het Envelope Theorem om JEPA guidance computationeel praktisch te maken, vergezeld van theoretische foutgrenzen.
• Training-Free Integratie: De aanpak is volledig training-free, vereist alleen vooraf getrainde diffusie- en JEPA-modellen, en kan tijdens de inferentie worden geïntegreerd zonder de onderliggende modellen te wijzigen.

Experimentele Resultaten

De auteurs hebben JEPA guidance geëvalueerd voor onvoorwaardelijke, klasse-voorwaardelijke en tekst-naar-afbeelding generatietaken (met behulp van CelebA, ImageNet, Stable Diffusion 1.5 en SDXL-Lightning).

• Atypiciteit: JEPA guidance bereikt consistent de laagste JEPA-SCORE (wat wijst op hogere atypiciteit onder wereldpriors) in vergelijking met baselines zoals Sehwag et al. (2022), MG, SGMS en MinorityPrompt.
• Fideliteit en Validiteit: Ondanks het richten op zeldzame gebieden, behoudt de methode een hoge fideliteit. Bij tekst-naar-afbeelding taken bereikt het concurrerende tekst-afstemmingsscores (CLIPScore, PickScore, ImageReward) terwijl het semantisch onderscheidende samples genereert (bijvoorbeeld oudere militair personeel, Afro-Amerikaanse honkballers) die andere methoden niet kunnen produceren.
• Kwalitatieve Vergelijking: Visualisaties tonen aan dat generator-gecentreerde methoden vaak samples produceren met atypische contexten (bijvoorbeeld een adelaar in de buurt van een kraan), terwijl JEPA guidance samples produceert met atypische semantiek of gedragingen (bijvoorbeeld parenende adelaars, ongewone dierlijke vormen), wat beter overeenkomt met zeldzaamheid in de realiteit.
• Downstream Utility: In een taak voor data-augmentatie voor classifier-training leverde het aanvullen met JEPA-geleide minority-samples een superieure of vergelijkbare prestatie op dan andere methoden met minder samples, wat suggereert dat deze samples informatiever trainingsignalen bieden.

Betekenis en Claims

Het artikel claimt dat JEPA guidance een robuuste oplossing biedt voor de beperking dat bestaande minority-samplers gevangen zitten in de eigen verdeling van het generatieve model. Door gebruik te maken van wereldmodellen (JEPAs), genereert de methode samples die echt zeldzaam zijn in de realiteit, en niet alleen zeldzaam relatief tot een specifieke dataset.

De auteurs benadrukken dat deze aanpak:

1. Het begrip minderheid ontkoppelt van de generatieve prior.
2. De kloof overbrugt tussen generatieve modellering en wereldmodellering, een sleutelcomponent in de zoektocht naar kunstmatige algemene intelligentie.
3. Een praktische, training-free tool biedt voor toepassingen die diverse en atypische data vereisen, zoals robuustheidstesten, verbetering van eerlijkheid en creatieve contentgeneratie.

De auteurs erkennen ook een potentieel dual-use probleem: het omkeren van de richting van de guidance bestaande biases in wereldpriors kan versterken, en zij moedigen verantwoord gebruik aan. Zij merken op dat hoewel de methode in vergelijking met standaard samplers computationele overhead introduceert, het sneller blijft dan sommige complexe baselines (zoals MinorityPrompt) en een redelijke afweging biedt voor de significante winst in wereld-gecentreerde atypiciteit.