Steering Diffusion Models via Class-Contrastive Influence for Few-Shot Medical Classification

Technische Samenvatting: Sturen van Diffusion Modellen via Class-Contrastive Influence voor Few-Shot Medische Classificatie

Probleemstelling
Bij few-shot medische beeldclassificatie is gelabelde data schaars, waardoor data-augmentatie essentieel is. Hoewel off-the-shelf diffusion modellen synthetische trainingsvoorbeelden kunnen genereren, richten bestaande benaderingen zich primair op het verbeteren van de realisme, diversiteit of domeinadaptatie van deze beelden. De auteurs stellen dat deze metrieken geen garantie bieden voor de bruikbaarheid voor de downstream classificatietaak. Er bestaat een significante discrepantie waarbij gegenereerde beelden met een vergelijkbare visuele kwaliteit zeer verschillende classificatieprestaties opleveren. Het kernprobleem dat wordt geïdentificeerd, is het gebrek aan een principieel criterium om de "taakbruikbaarheid" van synthetische monsters te meten en te optimaliseren. De paper stelt dat de centrale beperking van de huidige praktijken niet alleen de generatiekwaliteit is, maar het falen om monsters te genereren die specifiek helpen bij het verfijnen van de beslissingsgrens van de classifier.

Methodologie
De voorgestelde oplossing, Class-Contrastive Influence (C2I), verschuift het optimalisatiedoel van beeldkwaliteit naar de gradiënt-gebaseerde invloed van synthetische monsters op de classifier.

1. Class-Contrastive Influence (C2I) Criterium:

   • De methode maakt gebruik van gradiënt-gebaseerde invloestschatting om te kwantificeren hoe een trainingsmonster de validatieverlies beïnvloedt.
   • De auteurs observeren dat rauwe, klasse-agnostische invloedsscores niet correleren met classificatieprestaties omdat validatiesets conflicterende signalen bevatten van verschillende klassen.
   • In plaats daarvan definiëren zij C2I op basis van de invloedsgap: een monster wordt als nuttig beschouwd als de verliesgradiënt positief uitlijnt met de validatiegradiënten van de zelfde klasse en negatief uitlijnt met gradiënten van andere klassen.
   • Voor binaire classificatie wordt C2I geformuleerd als het kwadraat van het verschil tussen het gemiddelde de invloed op de positieve en negatieve validatiesets, genormaliseerd door de variantie. Voor multi-class taken wordt het gedefinieerd als de softmax-score van de gemiddelde invloed op de overeenkomstige klasse relatief aan alle klassen.

2. Theoretische Connectie met Sample Hardness:

   • Theoretische analyse (Theorem 1) demonstreert dat het maximaliseren van C2I in een logistische regressie setting de gegenereerde monsters naar het globale gemiddelde van de validatieset drijft.
   • Aangezien het globale gemiddelde van twee klasse-clusters doorgaans nabij de beslissingsgrens ligt, leidt het maximaliseren van C2I inherent tot het selecteren van "harde" voorbeelden—monsters die moeilijk te classificeren zijn en dicht bij de grens liggen.
   • Trainen op deze harde, grens-nabije voorbeelden helpt bij het verfijnen van de beslissingsgrens, wat generalisatie en robuustheid verbetert.

3. Reinforcement Learning (RL) Fine-Tuning:

   • Om dit operationeel te maken, stellen de auteurs een fine-tuning schema voor bestaande pre-trained Text-to-Image (T2I) diffusion modellen (specifiek Stable Diffusion 2.1).
   • Reward Functie: De C2I-score dient als het reward-signaal in een Reinforcement Learning framework. De generator wordt gefinetuned om de verwachte C2I-reward te maximaliseren.
   • Optimalisatie: De auteurs gebruiken Denoising Diffusion Policy Optimization (DDPO) om de parameters van het diffusion model (specifiek de LoRA-gewichten) bij te werken op basis van de C2I-reward. Dit stuurt het generatieproces naar de "klasse-informatieve" regio's van de datamanifold zonder architecturale wijzigingen te vereisen in de generator of de classifier.
   • Workflow: Een classifier (ViT-B/16) wordt eerst gefinetuned op de few-shot trainingsset. De gradiënten van deze classifier op een uitgesloten validatieset worden gebruikt om de C2I-reward te berekenen voor gegenereerde batches, die vervolgens de RL fine-tuning van het diffusion model sturen.

Belangrijkste Bijdragen

• Formalisering van Effectiviteit: De paper formaliseert monster-effectiviteit via C2I, en demonstreert dat hoge-C2I monsters overeenkomen met harde, grens-nabije voorbeelden in plaats van enkel realistische beelden.
• RL-gebaseerd Generatieschema: Het introduceert een plug-and-play RL-framework dat off-the-shelf diffusion modellen gebruikt met C2I als reward om te sturen naar het genereren van taak-nuttige, klasse-informatieve monsters.
• Empirische Validatie: De methode wordt gevalideerd over meerdere medische beeldvorming benchmarks (BreastMNIST, DermaMNIST, PneumoniaMNIST), waarbij consistente verbeteringen worden getoond ten opzichte van standaard augmentatie en andere diffusion-gebaseerde baselines.

Experimentele Resultaten
De methode werd geëvalueerd op drie MedMNIST benchmarks onder few-shot settings (16 of 32 gelabelde monsters per klasse).

• Prestaties: De door C2I gestuurde generator presteerde consequent beter dan baselines, inclus\n bij standaard transformaties (RandAugment, RandomErasing, Mixup) en andere diffusion-gebaseerde augmentatiemethoden (DataDream, Dataset Expansion, DistDiff).
  • Gemiddeld bereikte de voorgestelde methode de hoogste AUC over verschillende backbones (ViT-B/16 en ResNet-18).
  • Opvallend genoeg verbeterde de methode de prestaties op ResNet-18, wat aangeeft dat de gesynthetiseerde monsters breed informatief zijn en overdraagbaar zijn naar verschillende modelarchitecturen.
• Robuustheid: Modellen getraind met C2I-gestuurde augmentatie vertoonden superieure robuustheid onder diverse ruiscondities (salt-and-pepper, Gaussian blur, JPEG compressie), wat duidt op de generatie van een stabielere beslissingsgrens.
• Hardness Correlatie: Experimenten bevestigden dat naarmate de RL fine-tuning vordert, de C2I-reward toeneemt en de kenmerken van de gegenereerde beelden dichter bij de validatie-feature centroids komen te liggen, wat de generatie van hardere voorbeelden bevestigt.
• Ablatie: Het verhogen van het aantal RL training epochs en het aantal gesynthetiseerde beelden verbeterde over het algemeen de downstream AUC, terwijl sommige baselines (bijv. Dataset Expansion) een prestatieverslechtering vertoonden bij een toename van synthetische data.

Betekenis en Claims
De paper claimt dat synthetische augmentatie het meest effectief is wanneer deze wordt gestuurd door taakbruikbaarheid in plaats van enkel beeldkwaliteit. Door de synthetische data af te stemmen op de validatiegradiënten, transformeert C2I generieke diffusion modellen in gerichte generatoren die de beslissingsgrenzen versterken waar gelabelde data het meest schaars is.

De auteurs erkennen beperkingen en merken op dat hun methode extra computationele overhead introduceert vergeleken met het gebruik van alleen de originele trainingsdata, voornamelijk door het RL fine-tuning proces (dat ongeveer 5 GPU-uren op een A100 vereist). Ze stellen echter dat deze kosten gerechtvaardigd zijn door de consistente winst in nauwkeurigheid en robuustheid in omgevingen met weinig data. De studie suggereert een paradigmaverschuiving in het gebruik van diffusion modellen voor augmentatie: het optimaliseren van generatie voor de downstream invloed op de classifier.