Riepilogo Tecnico: GRASP – Adattatori Residui Guidati con Partizionamento per Campione

Enunciato del Problema

I transformer per il flow matching da testo a immagine mostrano un grave degrado delle prestazioni nelle distribuzioni a coda lunga, in particolare in domini come l'imaging medico dove le patologie rare dominano gli scenari reali. Mentre i modelli performano bene sulle classi frequenti "di testa", subiscono un collasso delle modalità sulle classi "di coda", generando immagini con scarsa fedeltà e limitata diversità. Gli autori identificano la causa radice come una patologia a livello di ottimizzazione: durante il fine-tuning, gli aggiornamenti del gradiente per le classi di testa e di coda sono disallineati sia in norma che in direzione. Questi gradienti si mediando in modo distruttivo a livello dei parametri, costringendo il modello a privilegiare le classi frequenti e collassare la rappresentazione di quelle rare. Gli interventi esistenti, come le strategie di campionamento post-hoc o le modifiche alla funzione di perdita durante l'addestramento, non affrontano questo specifico disallineamento dei gradienti, lasciando spesso invariata la distribuzione sottostante o modificando l'obiettivo generativo in modi che non colpiscono direttamente il fallimento dell'ottimizzazione.

Metodologia

Il documento propone GRASP (Guided Residual Adapters with Sample-wise Partitioning), un'intervento architetturale non invasivo progettato per migliorare l'allineamento dei gradienti senza alterare l'obiettivo del flow matching o il processo di campionamento.

1. Partizionamento per Campione tramite Condizionamento

Invece di apprendere una rete di gating (che rischia di essere dominata dai gradienti delle classi di testa), GRASP impiega una partizione statica e deterministica dello spazio di condizionamento.

• Base Teorica: Nel flow matching condizionato, i valori di condizione indicizzano insiemi distinti di percorsi probabilistici. Gli autori sostengono che il partizionamento lungo il vettore di condizione sia la fattorizzazione strutturalmente corretta da utilizzare come proxy per l'allineamento dei gradienti.
• Implementazione: Una funzione di instradamento statica $\pi: \mathcal{C} \to \{1, \dots, K\}$ assegna ogni campione a uno dei $K$ esperti in base alla sua condizione.
  • Contesto Medico (MIMIC-CXR-LT): I campioni sono partizionati in gruppi semantici (Testa, Media, Coda) e ulteriormente suddivisi da un prior specifico del dominio (Sano vs Patologico), risultando in quattro esperti.
  • Contesto Generale (ImageNet-LT): Il partizionamento è ottenuto tramite clustering degli embedding testuali (bisezione K-means) sull'encoder di testo congelato, senza priori specifici del dominio.
• Robustezza: Il metodo si basa su priori "grossolani" (ad esempio, bin di frequenza o clustering standard) piuttosto che su una separazione semantica perfetta. L'obiettivo è garantire che, in media, i campioni all'interno di un esperto condividano dinamiche di gradiente più omogenee rispetto a quelle tra esperti diversi.

2. Adattatori Residui Guidati

Una volta instradati i campioni, l'architettura utilizza adattatori residui specifici per gruppo inseriti negli strati feedforward del backbone pre-addestrato congelato (Flux).

• Meccanismo: Per un campione $x_i$ instradato all'esperto $k = \pi(c_i)$, l'output dello strato è $h_{l+1} = F_l(h_l) + A_{k,l}(h_l)$, dove $F_l$ è la trasformazione di base condivisa e $A_{k,l}$ è l'adattatore specifico per l'esperto.
• Isolamento del Gradiente: I parametri all'interno dell'adattatore $k$ vengono aggiornati solo dai campioni assegnati a quell'esperto. Questo isola i gradienti specifici del gruppo, prevenendo interferenze distruttive derivanti dagli aggiornamenti delle classi di testa, consentendo al contempo al backbone condiviso di accumulare strutture trasferibili.
• Progettazione: Gli adattatori sono trasformazioni leggere a due strati. A differenza di LoRA, non impongono una riduzione rigorosa della dimensionalità, permettendo adattatori più ampi ($d' \approx d$ o $d' > d$) per stabilizzare il fine-tuning attraverso grandi divari di dominio.

Contributi Chiave

1. Trasferimento della Coerenza del Gradiente al Flow Matching: Gli autori adattano l'obiettivo di coerenza del gradiente dal linguaggio modellato a coda lunga (Dong et al.) al flow matching da testo a immagine. Crucialmente, sostituiscono il meccanismo di gating appreso con una partizione deterministica guidata dal condizionamento. Sostengono che nel flow matching condizionato, il condizionamento indicizza naturalmente la geometria del percorso probabilistico, rendendo il partizionamento statico una fattorizzazione strutturalmente corretta piuttosto che un euristico generico.
2. Progettazione Non Invasiva e Componibile: GRASP lascia intatti la funzione di perdita del flow matching e il campionatore, modificando solo quali parametri assorbono quali gradienti. Ciò permette di combinarlo con interventi ortogonali. Gli autori dimostrano che combinare GRASP con il campionamento minoritario auto-guidato al momento dell'inferenza produce risultati superiori rispetto all'uso di ciascun metodo singolarmente.
3. Validazione Empirica del Condizionamento come Proxy del Gradiente: Attraverso studi di ablazione e analisi dei gradienti, il documento valida che l'omogeneità del condizionamento si correla fortemente con la coerenza dei gradienti per esperto, confermando che le partizioni statiche approssimano efficacemente il raggruppamento ideale dei gradienti anche con priori rumorosi.

Risultati

Il metodo è stato valutato su MIMIC-CXR-LT (radiografie toraciche mediche), NIH-CXR-LT (classificazione a valle) e ImageNet-LT (generazione generale a coda lunga).

• Qualità e Diversità della Generazione di Immagini:
  • Su MIMIC-CXR-LT, GRASP ha ridotto l'FID complessivo fino all'80% rispetto al fine-tuning completo e ha aumentato la copertura delle classi di coda fino al 44%.
  • Ha superato le linee di base dello stato dell'arte, inclusi il fine-tuning completo, adattatori generici e varianti di Mixture-of-Experts (MoE) con instradamento appreso come MoCLE.
  • L'analisi qualitativa ha mostrato che GRASP ha recuperato caratteristiche specifiche della classe per patologie rare (ad esempio, Pneumoperitoneo) che le linee di base avevano appianato o fatto collassare.
• Classificazione a Valle:
  • Quando le immagini sintetiche generate da GRASP sono state utilizzate per addestrare un classificatore DenseNet-121 su NIH-CXR-LT, il modello ha raggiunto un punteggio F1 macro pari a quello dell'addestramento su dati reali.
  • Crucialmente, GRASP ha permesso punteggi F1 non nulli su 9 classi su 13, mentre il fine-tuning completo ha raggiunto questo risultato solo su 3 classi.
• Generalizzazione: I risultati su ImageNet-LT hanno confermato che il meccanismo non è legato a bias induttivi medici e si generalizza a compiti di generazione a coda lunga di ampia portata.

Significato e Affermazioni

Il documento posiziona GRASP come uno strumento di augmentazione controllata e di analisi, affermando esplicitamente che le immagini sintetiche non sono un sostituto di coorti clinicamente fondate o di revisioni esperte. Gli autori sottolineano che i dati sintetici possono preservare o amplificare i bias se utilizzati senza provenienza e validazione.

Il significato principale rivendicato è che l'instradamento statico è la scelta predefinita per la generazione da testo a immagine a coda lunga ogni volta che sono disponibili anche priori grossolani sul condizionamento (un caso comune nell'imaging medico). A differenza del gating appreso, che fallisce in regimi a coda lunga estrema a causa del dominio delle classi di testa, il partizionamento statico garantisce che ogni campione di coda aggiorni il proprio esperto assegnato. Ciò fornisce un meccanismo stabile, efficiente e non invasivo per mitigare le patologie a livello di ottimizzazione senza modificare l'obiettivo generativo fondamentale.