IMS3: Breaking Distributional Aggregation in Diffusion-Based Dataset Distillation

1. Het Probleem: Distributie-Aggregatie

Dataset-distillatie (Dataset Distillation) heeft als doel een compact synthetisch dataset te creëren dat de trainingsprestaties van een groot origineel dataset benadert. Recentelijk hebben methoden op basis van diffusiemodellen (zoals Minimax Diffusion) veelbelovende resultaten geboekt door gebruik te maken van de sterke generatieve capaciteit van deze modellen.

Echter, de auteurs identificeren een fundamenteel doel-misalignment (doelconflict):

• Generatieve vs. Discriminerende Doelen: Diffusiemodellen zijn geoptimaliseerd om de waarschijnlijkheid van data te maximaliseren (generatieve likelihood). Dit leidt ertoe dat gegenereerde samples zich overmatig concentreren in hoog-dichtheidsgebieden van de data-manifold.
• Het Aggregatie-probleem: Voor een effectieve dataset-distillatie (die dient voor classificatie) zijn echter samples in laag-dichtheidsgebieden, specifiek nabij de beslissingsgrenzen (decision boundaries), cruciaal. Bestaande methoden missen deze randgebieden, wat resulteert in een gebrek aan diversiteit en een slechtere generalisatie voor downstream-classificatoren. Dit fenomeen noemen de auteurs distributional aggregation.

2. Methodologie: ImS3 Framework

Om dit probleem op te lossen, stellen de auteurs ImS3 voor, een tweestapsframework dat bestaat uit twee complementaire strategieën: Inversion-Matching (IM) en Selective Subgroup Sampling (S3).

A. Inversion-Matching (IM) voor Fine-Tuning

Deze fase richt zich op het uitbreiden van de distributie-coverage door het diffusiemodel te fine-tunen.

• Inzicht: De auteurs benutten de inherente instabiliteit van het diffusie-inversieproces. Wanneer een real image wordt "teruggevoerd" naar ruis (inversie), wijken de trajecten vaak af van hoog-dichtheidsgebieden en bewegen ze zich naar laag-dichtheidsgebieden vanwege cumulatieve benaderingsfouten.
• Mechanisme: In plaats van deze instabiliteit als een fout te zien, gebruiken ze deze als een signaal.
  1. Ze berekenen een inversie-ruis-latent ($z^{inv}_t$) voor een gegeven tijdstip $t$.
  2. Ze trainen het model om de standaard ruis-latent ($z_t$) af te stemmen op de inversie-latent ($z^{inv}_t$) via een cosine-similariteitsverlies ($\mathcal{L}_{IM}$).
  3. Dit verlies wordt gecombineerd met het standaard diffusieverlies ($\mathcal{L}_{Diff}$) om de beeldkwaliteit (fidelity) te behouden terwijl de coverage van laag-dichtheidsgebieden wordt vergroot.
• Efficiëntie: Ze gebruiken Parameter-Efficient Fine-Tuning (PEFT) via Difffit, waarbij alleen lichte adapters worden bijgewerkt, wat de rekencosts verlaagt zonder stabiliteit te verliezen.

B. Selective Subgroup Sampling (S3) voor Sampling

Deze fase vindt plaats tijdens het genereren van het synthetische dataset en is trainingsvrij.

• Doel: Het verbeteren van de inter-class scheiding (discriminatie) zonder de representativiteit te verliezen.
• Mechanisme:
  1. Voor elke klasse worden meerdere kandidaat-subgroepen (subgroups) gegenereerd.
  2. Voor elke subgroep wordt het centroid (middelpunt) in de feature-ruimte berekend.
  3. Een selectie-optimisatie kiest de beste subgroep per klasse op basis van twee criteria:
     • Representativiteit: De subgroep moet dicht bij het centroid van de real data liggen.
     • Scheiding: De subgroep moet ver verwijderd zijn van de centroids van andere klassen.
  4. Dit wordt geoptimaliseerd via een doelwitfunctie ($\mathcal{L}_{S^3}$) die een balans zoekt tussen deze twee factoren.

3. Belangrijkste Bijdragen

1. Identificatie van Distributional Aggregation: Het paper benadrukt en kwantificeert het probleem dat bestaande diffusie-methoden te veel focus leggen op hoog-dichtheidsgebieden, ten koste van de cruciale beslissingsgrenzen.
2. Inversion-Matching (IM): Een nieuwe fine-tuning-strategie die de inherente instabiliteit van diffusie-inversie omarmt om het model te sturen naar ondervertegenwoordigde, laag-dichtheidsgebieden.
3. Selective Subgroup Sampling (S3): Een trainingsvrije sampling-strategie die de discriminatieve structuur van het dataset verbetert door bewuste selectie van subgroepen die zowel representatief als goed gescheiden zijn.
4. State-of-the-Art Prestaties: ImS3 bereikt nieuwe topprestaties op meerdere benchmarks, wat aantoont dat het combineren van distributie-uitbreiding en bewuste sampling effectief is.

4. Resultaten

De auteurs hebben hun methode getest op diverse datasets, waaronder ImageWoof (fijne granulariteit, hoge inter-class gelijkenis), ImageNette, ImageIDC en ImageNet-100.

• Prestaties: ImS3 overtreft consistent bestaande methoden (zoals Minimax, D4M, D3HR, en optimalisatie-based methoden zoals SRe2L) in termen van Top-1 nauwkeurigheid.
  • Op ImageWoof (IPC=10) verbetert ImS3 de nauwkeurigheid voor ResNetAP-10 van 35.7% (Minimax) naar 41.8%.
  • Op ImageNette (IPC=50) bereikt het 84.2%, wat een verbetering is ten opzichte van de sterkste baseline.
• Visualisatie: t-SNE visualisaties tonen aan dat ImS3 een bredere dekking van de feature-ruimte biedt en duidelijkere scheiding tussen klassen creëert vergeleken bij de geaggregeerde clusters van Minimax.
• Robuustheid: De methode werkt goed over verschillende modelarchitecturen (ConvNet, ResNet, ViT) en verschillende "Images Per Class" (IPC) budgetten.

5. Significantie

Dit onderzoek is significant omdat het een fundamentele beperking in diffusie-based dataset distillatie aanpakt. Door de focus te verschuiven van puur generatieve kwaliteit naar discriminatieve nut, biedt ImS3 een praktische oplossing voor het trainen van modellen op extreem kleine datasets zonder in te leveren op generalisatievermogen. Het paper toont aan dat het bewust manipuleren van de distributie-coverage (via inversie) en de sampling-strategie (via subgroepselectie) essentieel is voor de volgende generatie dataset-distillatie technieken. De code en modellen worden openbaar gemaakt om toekomstig onderzoek te faciliteren.