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

Sintesi Tecnica: Guida dei Modelli di Diffusione tramite l'Influenza Class-Contrastive per la Classificazione Medica Few-Shot

Definizione del Problema
Nella classificazione di immagini mediche in ambito few-shot, i dati etichettati sono scarsi, rendendo essenziale l'aumento dei dati (data augmentation). Sebbene i modelli di diffusione pronti all'uso possano generare campioni di addestramento sintetici, gli approcci esistenti si concentrano principalmente sul miglioramento del realismo, della diversità o dell'adattamento al dominio di queste immagini. Gli autori sostengono che tali metriche non garantiscono l'utilità per il compito di classificazione a valle. Esiste una discrepanza significativa in cui immagini generate con una qualità visiva comparabile producono prestazioni di classificazione vastamente differenti. Il problema centrale identificato è la mancanza di un criterio fondato per misurare e ottimizzare l' "utilità per il compito" (task usefulness) dei campioni sintetici. Il paper pone come limite principale delle pratiche attuali non la sola qualità della generazione, ma il fallimento nel generare campioni che aiutino specificamente il classificatore a raffinare il proprio confine di decisione.

Metodologia
La soluzione proposta, la Class-Contrastive Influence (C2I), sposta l'obiettivo di ottimizzazione dalla qualità dell'immagine all'influenza basata sul gradiente dei campioni sintetici sul classificatore.

1. Criterio di Class-Contrastive Influence (C2I):

   • Il metodo sfrutta la stima dell'influenza basata sul gradiente per quantificare come un campione di addestramento influenzi la perdita di validazione.
   • Gli autori osservano che i punteggi di influenza grezzi, agnostici rispetto alla classe, non correlano con le prestazioni di classificazione perché i set di validazione contengono segnali contrastanti provenienti da classi diverse.
   • Invezione, definiscono la C2I basandosi sul gap di influenza: un campione è considerato utile se il suo gradiente di perdita si allinea positivamente con i gradienti di validazione della stessa classe e si oppone (si allinea negativamente con) ai gradienti di altre classi.
   • Per la classificazione binaria, la C2I è formulata come la differenza al quadrato tra la media dell'influenza sui set di validazione positivi e negativi, normalizzata dalla varianza. Per i compiti multi-classe, è definita come lo score softmax della media dell'influenza sulla classe corrispondente rispetto a tutte le classi.

2. Connessione Teorica con la Difficoltà del Campione (Sample Hardness):

   • L'analisi teorica (Teorema 1) dimostra che massimizzare la C2I in un contesto di regressione logistica spinge il campione generato verso la media globale del set di validazione.
   • Poiché la media globale di due cluster di classe si trova tipicamente vicino al confine di decisione, massimizzare la C2I seleziona intrinsecamente esempi "difficili" (hard)—campioni che sono difficili da classificare e si trovano vicino al confine.
   • L'addestramento su questi esempi difficili, prossimali al confine, aiuta a raffinare il confine di decisione, migliorando la generalizzazione e la robustezza.

3. Fine-Tuning tramite Reinforcement Learning (RL):

   • Per operazionalizzare questo processo, gli autori propongono uno schema di fine-tuning per modelli di diffusione Text-to-Image (T2I) pre-addestrati (specificamente Stable Diffusion 2.1).
   • Funzione di Reward: Il punteggio C2I funge da segnale di reward in un framework di Reinforcement Learning. Il generatore viene sottoposto a fine-tuning per massimizzare l'aspettativa del reward C2I.
   • Ottimizzazione: Gli autori utilizzano il Denoising Diffusion Policy Optimization (DDPO) per aggiornare i parametri del modello di diffusione (specificamente i pesi LoRA) basandosi sul reward C2I. Questo guida il processo di generazione verso le regioni del manifold dei dati "class-informative" senza richiedere modifiche architettoniche al generatore o al classificatore.
   • Workflow: Un classificatore (ViT-B/16) viene prima addestrato tramite fine-tuning sul set di addestramento few-shot. I suoi gradienti su un set di validazione tenuto separato vengono utilizzati per calcolare il reward C2I per i batch generati, che poi guidano il fine-tuning RL del modello di diffusione.

Contributi Chiave

• Formalizzazione dell'Efficacia: Il paper formalizza l'efficacia del campione attraverso la C2I, dimostrando che i campioni ad alta C2I corrispondono a esempi difficili e prossimali al confine, piuttosto che solo realistici.
• Schema di Generazione basato su RL: Introduce un framework RL "plug-and-play" che utilizza la C2I come reward per guidare i modelli di diffusione off-the-shelf verso la generazione di campioni utili al compito e class-informative.
• Validazione Empirica: Il metodo è validato attraverso molteplici benchmark di imaging medico few-shot (BreastMNIST, DermaMNIST, PneumoniaMNIST), mostrando miglioramenti costanti rispetto all'augmentation standard e ad altri baseline basati sulla diffusione.

Risultati Sperimentali
Il metodo è stato valutato su tre benchmark MedMNIST in setting few-shot (16 o 32 campioni etichettati per classe).

• Performance: Il generatore guidato da C2I ha superato costantemente i baseline, inclusi le trasformazioni standard (RandAugment, RandomErasing, Mixup) e altri metodi di augmentation basati sulla diffusione (DataDream, Dataset Expansion, DistDiff).
  • In media, il metodo proposto ha ottenuto l'AUC più elevato attraverso diversi backbone (ViT-B/16 e ResNet-18).
  • Notevolmente, il metodo ha migliorato le prestazioni su ResNet-18, indicando che i campioni sintetizzati sono ampiamente informativi e trasferibili tra diverse architetture di modelli.
• Robustezza: I modelli addestrati con l'augmentation guidata da C2I hanno dimostrato una robustezza superiore sotto varie condizioni di rumore (sale e pepe, sfocatura gaussiana, compressione JPEG), suggerendo la generazione di un confine di decisione più stabile.
• Correlazione della Difficoltà (Hardness): Gli esperimenti hanno confermato che man mano che il fine-tuning RL progredisce, il reward C2I aumenta e le caratteristiche delle immagini generate si avvicinano ai centroidi delle feature di validazione, confermando la generazione di esempi più difficili.
• Ablation: L'aumento del numero di epoche di addestramento RL e del numero di immagini sintetizzate ha generalmente migliorato l'AUC a valle, mentre alcuni baseline (ad esempio, Dataset Expansion) hanno mostrato un degrado delle prestazioni con l'aumento dei dati sintetici.

Significatività e Rivendicazioni
Il paper sostiene che l'augmentation sintetica sia più efficace quando guidata dall' utilità per il compito piuttosto che dalla sola qualità dell'immagine. Allineando i dati sintetici con i gradienti di validazione, la C2I trasforma i modelli di diffusione generici in generatori mirati che rafforzano i confini decisionali dove i dati etichettati sono più scarsi.

Gli autori riconoscono i limiti, notando che il loro metodo introduce un carico computazionale aggiuntivo rispetto all'uso dei soli dati di addestramento originali, principalmente a causa del processo di fine-tuning RL (che richiede circa 5 ore di GPU su una A100). Tuttavia, affermano che questo costo è giustificato dai guadagni costanti in accuratezza e robustezza nei regimi di dati ridotti. Lo studio suggerisce un cambio di paradigma nell'uso dei modelli di diffusione per l'augmentation: ottimizzare la generazione per la sua influenza a valle sul classificatore.