Improving Diversity in Black-box Few-shot Knowledge Distillation

1. Probleemstelling

Het artikel behandelt de uitdaging van Black-box Few-shot Knowledge Distillation (KD).

• Context: Standaard KD overdraagt kennis van een groot "docent"-model naar een kleiner "student"-model. Echter, traditionele KD vereist een groot trainingsdataset en volledige toegang tot de interne parameters van de docent (white-box).
• Beperkingen: In real-world scenario's (bijvoorbeeld het distilleren van propriëtaire modellen zoals Google's FaceNet of ChatGPT) heeft de student-ontwikkelaar vaak:
  1. Black-box toegang: Alleen de uiteindelijke voorspellingssannihilaties van de docent zijn beschikbaar; interne gradiënten, logits of parameters zijn verborgen.
  2. Few-shot data: Er is slechts een zeer klein aantal gelabelde afbeeldingen ($N$) beschikbaar voor het trainen van de student, wat onvoldoende is om het ware data-manifold te representeren.
• Beperkingen van bestaande oplossingen: Huidige few-shot KD-methoden vertrouwen vaak op white-box docenten of gebruiken eenvoudige data-augmentatie (zoals MixUp), wat synthetische afbeeldingen produceert die diversiteit of realistische semantiek missen. Andere methoden die Variational Autoencoders (CVAE) gebruiken, lijden aan wazige artefacten en slagen er niet in om diversiteit expliciet te bevorderen.

2. Methodologie: DivBFKD

De auteurs stellen DivBFKD (Diverse Black-box Few-Shot Knowledge Distillation) voor, een tweefasig raamwerk ontworpen om diverse synthetische afbeeldingen te genereren die de oorspronkelijke trainingsverdeling van de docent nabootsen zonder toegang tot de interne werking van de docent.

Fase 1: Generatie (WGAN met High-Confidence Selectie)

De kerninnovatie ligt in het trainen van een Wasserstein Generative Adversarial Network (WGAN) om diverse en hoogwaardige synthetische afbeeldingen te genereren.

• High-Confidence Afbeeldingen: In plaats van alle gegenereerde afbeeldingen te gebruiken, filtert de methode synthetische afbeeldingen op basis van het vertrouwen van de docent. Een afbeelding $\tilde{x}$ wordt als "high-confidence" beschouwd als de maximale voorspelde waarschijnlijkheid $c_{\tilde{x}}$ van de docent een drempel $\tau$ overschrijdt.
• Adaptieve Drempels ($\tau$): Een vaste drempel is inefficiënt omdat docenten klas-specifieke vertekeningen vertonen (bijvoorbeeld het voorspellen van "broeken" met een hoger vertrouwen dan "overhemden"). De auteurs stellen adaptieve drempels voor die per klas worden berekend met behulp van de $q$-quantiel (bijvoorbeeld $q=0.1$) van vertrouwensscores uit de few-shot real-afbeeldingen. Dit zorgt ervoor dat de drempel relatief is tot het gedrag van de docent voor elke specifieke klas.
• Adversarial Training On-the-Fly:
  • De WGAN-generator ($G$) creëert synthetische afbeeldingen.
  • De discriminator ($D$) wordt getraind om te onderscheiden tussen real-afbeeldingen (uit de few-shot set $D$) en high-confidence synthetische afbeeldingen (geselecteerd uit de vorige iteratie).
  • Feedbacklus: Door high-confidence synthetische afbeeldingen als "real" te behandelen voor de discriminator, wordt de generator gestimuleerd om meer diverse afbeeldingen te produceren die de docent met vertrouwen classificeert. Dit creëert een positieve feedbacklus die de diversiteit van de gegenereerde set uitbreidt.

Fase 2: Distillatie

• Dataset Constructie: Een distillatiedataset ($D_{KD}$) wordt samengesteld door de originele few-shot real-afbeeldingen ($N$) te combineren met een budget aan synthetische afbeeldingen ($M$) gegenereerd door de getrainde WGAN.
• Klasbalancering: Om vertekening te voorkomen, worden synthetische afbeeldingen gegenereerd met klasbalancering via rejection sampling.
• Training: Het studentnetwerk wordt getraind op $D_{KD}$ met behulp van standaard Cross-Entropy-verlies tegen de voorspellingen van de docent (soft labels).

3. Belangrijkste Bijdragen

1. Nieuw Trainingschema voor WGAN: Introduceert een mechanisme om high-confidence afbeeldingen on-the-fly in de adversarial learning-lus te injecteren. Dit gebruikt de docent effectief als supervisor om de generator te leiden naar de onbekende trainingsverdeling van de docent.
2. Adaptieve Drempels: Stelt een klas-specifieke, dataset-agnostische methode voor om "high-confidence" te definiëren met behulp van op quantiel gebaseerde drempels, wat docent-vertekening mitigeert en de selectie van abnormale afbeeldingen voorkomt.
3. Diverse Black-box Few-shot KD (DivBFKD): Een compleet raamwerk dat het gebrek aan diversiteit in synthetische data aanpakt, een kritieke bottleneck in few-shot KD.
4. State-of-the-Art Prestaties: Toont significante verbeteringen ten opzichte van bestaande few-shot KD-methoden over diverse datasets en architecturen.

4. Experimentele Resultaten

De methode werd geëvalueerd op zeven datasets (MNIST, FMNIST, SVHN, CIFAR10, CIFAR100, Tiny-ImageNet, Imagenette) met diverse architecturen (LeNet, AlexNet, ResNet, VGG).

• Prestatiewinst:
  • Eenvoudige Datasets: DivBFKD presteerde beter dan de tweede beste black-box methode (FS-BBT) met ongeveer 2% op FMNIST en CIFAR10. Het verbeterde Standard-KD met 17% op CIFAR10.
  • Complexe Datasets: Op CIFAR100 bereikte DivBFKD een nauwkeurigheid van 59,91%, wat FS-BBT (56,28%) met ongeveer 3,6% overtrof. Op Imagenette (high-res) bereikte het 86,48%, waarmee de kloof met de "Student-Full" bovengrens aanzienlijk werd verkleind.
  • Cross-Architectuur: De methode bleef robuust, zelfs wanneer docent en student verschillende architecturen hadden (bijvoorbeeld ResNet-docent naar VGG-student), met een verbetering van 14–22% in nauwkeurigheid ten opzichte van Standard-KD.
• Diversiteits- en Kwaliteitsmetrieken:
  • Coverage: Synthetische afbeeldingen dekten 42% van de buurten van trainingsafbeeldingen van de docent, vergeleken met slechts 18% voor few-shot afbeeldingen en 29% voor standaard WGAN-outputs.
  • FID/IS: Bereikte de beste Fréchet Inception Distance (13,96) en Inception Score (4,42) onder de baselines, wat wijst op hoge diversiteit en kwaliteit.
• Ablatie Studies: Bevestigde dat zowel Adaptieve Drempels (AT) als Klasbalancering (CB) essentiële componenten zijn, waarbij AT de meest significante boost in diversiteit biedt.
• Vergelijking met Data-Free KD: DivBFKD presteerde aanzienlijk beter dan data-free (zero-shot) KD-methoden, wat bewijst dat zelfs een kleine hoeveelheid real-data (2K afbeeldingen) gecombineerd met diverse synthetische data superieur is aan het hebben van geen real-data.

5. Betekenis

Dit artikel levert een significante bijdrage aan het veld van modelcompressie en privacybehoudende AI:

• Praktische Toepasbaarheid: Het lost een kritiek real-world probleem op waarbij organisaties geen grote datasets of modelgewichten kunnen delen vanwege privacy- of intellectuele-eigendomszorgen.
• Focus op Diversiteit: Het verschuift de focus van het genereren van enige synthetische data naar het genereren van diverse en docent-gealigneerde data, wat de sleutel is tot effectieve distillatie in dataskarce omgevingen.
• Efficiëntie: De voorgestelde methode voegt verwaarloosbare computere last toe (eenvoudige drempelstelling en on-the-fly selectie) terwijl het state-of-the-art resultaten levert.
• Generaliseerbaarheid: De aanpak is architectuur-agnostisch en werkt over eenvoudige en complexe afbeeldingsclassificatietaken, waardoor het een veelzijdig hulpmiddel is voor het implementeren van lichtgewicht modellen in beperkte omgevingen.