Post-Training Augmentation Invariance

Titolo: Invarianza all'Augmentation Post-Addestramento (Post-Training Augmentation Invariance)

1. Il Problema

I modelli pre-addestrati su larga scala (come Vision Transformers, CLIP, DINO, ecc.) hanno rivoluzionato il machine learning fornendo caratteristiche (feature) generali trasferibili a molti compiti. Tuttavia, questi modelli non sono necessariamente invarianti a determinate trasformazioni (augmentation) critiche per applicazioni specifiche, come rotazioni, trasformazioni affini, rumore o perturbazioni avversarie.

Il problema centrale affrontato è: come rendere un modello pre-addestrato invariante a specifiche augmentation senza corrompere le sue capacità esistenti sulla distribuzione originale dei dati?

• Riaddestramento: È computazionalmente proibitivo e rischia di degradare le rappresentazioni apprese.
• Fine-tuning standard: Può distorcere significativamente le feature pre-addestrate, alterando la geometria dello spazio latente originale.
• Obiettivo: Aggiungere proprietà di invarianza mantenendo intatta la struttura metrica dello spazio latente originale (preservando le relazioni tra i punti dati non augmentati).

2. Metodologia

Gli autori propongono un framework che appende una rete "adapter" leggera alla parte latente di una rete pre-addestrata congelata ($F$). L'obiettivo è addestrare un encoder $E_\theta$ tale che la composizione $E_\theta \circ F$ sia invariante alle augmentation.

Concetti Chiave e Definizioni

• Encoder Augmentati: Vengono formalizzati come encoder probabilistici che codificano un input dopo averlo sottoposto a una distribuzione di augmentation casuali.
• Invarianza Post-Addestramento: Definita formalmente come $(t, \mu_X, F, V)$-invarianza. Un encoder $E$ è invariante se:
  1. $E \circ F$ è invariante all'augmentation $t$.
  2. $E$ preserva la struttura dello spazio latente di $F$ (fino a una mappa $V$, tipicamente un'isometria o un'approssimazione locale), garantendo che le performance sui dati originali non vengano degradate.

Funzioni di Perdita Proposte

Gli autori introducono due nuove funzioni di perdita basate sulla teoria del trasporto ottimo (Optimal Transport):

1. Minimizzazione Markov-Wasserstein (MaWa):

   • Basata sulla metrica $L_p$ generalizzata tra kernel Markov-Wasserstein.
   • Obiettivo: Minimizzare la distanza tra la distribuzione delle feature augmentate e quella delle feature originali (ancorate).
   • Meccanismo: Utilizza una perdita MSE ancorata (anchored MSE). Il primo termine mantiene le feature originali vicine all'identità (preservazione della struttura), mentre il secondo termine spinge le feature augmentate verso le feature originali corrispondenti (invarianza).
   • Vantaggio: Evita il collasso della rappresentazione (collapse) tipico delle perdite di invarianza naive, poiché l'encoder è vincolato a non allontanarsi troppo dalle feature originali congelate.

2. Massimizzazione della Correlazione di Wasserstein (WaCo):

   • Basata sulla massimizzazione della correlazione di Wasserstein tra la distribuzione congiunta indotta dall'encoder augmentato e il prodotto delle sue marginali.
   • Obiettivo: Massimizzare la dipendenza statistica normalizzata tra input e output, costringendo l'encoder a mappare le diverse viste augmentate dello stesso input nello stesso punto (o vicino) nello spazio latente, mantenendo al contempo la struttura metrica.
   • Vantaggio: Può essere utilizzata anche per la riduzione della dimensionalità e non richiede necessariamente una rete pre-addestrata congelata come riferimento diretto (anche se può essere usata con essa).

Cosa NON funziona (Baselines)

Gli autori confrontano i loro metodi con:

• SimCLR (Contrastive Learning): Sebbene induca invarianza, altera drasticamente la geometria dello spazio latente originale, rendendo le feature inutilizzabili per compiti downstream originali.
• HSIC (Hilbert-Schmidt Independence Criterion): Simile a SimCLR, porta a una corruzione significativa dello spazio latente e a risultati di classificazione non competitivi.

3. Risultati Sperimentali

Gli esperimenti sono stati condotti su diversi dataset (STL10, TinyImageNet) e modelli pre-addestrati (DINOv2, SwAV, CLIP, ResNet50 supervisionato).

• Performance di Classificazione:

  • Su STL10 con DINOv2, l'approccio proposto (MaWa) raggiunge il 94% di accuratezza su immagini arbitrariamente ruotate, contro il 71% della rete base senza adapter.
  • Per l'invarianza al rumore, l'accuratezza sale dal 58% al 86%.
  • I metodi MaWa e WaCo superano costantemente SimCLR e HSIC, specialmente su augmentation distruttive come rotazioni e rumore.

• Preservazione della Struttura:

  • Le reti adapter addestrate con MaWa e WaCo agiscono quasi come isometrie sullo spazio latente originale. Le statistiche di struttura (coefficiente di determinazione $R^2$, Lipschitz constants, RMSD normalizzato) mostrano che le distanze tra i punti non augmentati vengono preservate quasi perfettamente.
  • Al contrario, SimCLR e HSIC causano un crollo della struttura metrica ($R^2 \approx 0$, distanze corrotte), rendendo le feature inadatte per compiti che richiedono la struttura originale.

• Analisi dei Tassi di Collisione:

  • L'efficacia del metodo MaWa dipende dal "tasso di collisione" (collision rate) delle feature pre-addestrate. Se le augmentation spostano le feature attraverso i confini di classe in modo coerente (basso tasso di collisione allineato), MaWa funziona bene.
  • Su modelli supervisionati come ResNet50, dove le collisioni residue sono alte, MaWa fallisce (collassa), mentre WaCo mostra una maggiore robustezza, sebbene con risultati inferiori rispetto ai modelli self-supervised.

• Riduzione della Dimensionalità:

  • WaCo è stato dimostrato efficace anche per ridurre la dimensionalità (es. da 384 a 96 dimensioni) mantenendo un'alta accuratezza e invarianza.

4. Contributi Chiave

1. Definizione Formale: Formalizzazione precisa del problema dell'invarianza post-addestramento e introduzione del concetto di "encoder augmentato" probabilistico.
2. Nuovi Obiettivi di Apprendimento: Proposta di due nuove funzioni di perdita (MaWa e WaCo) basate sul trasporto ottimo, specifiche per preservare la struttura mentre si aggiunge invarianza.
3. Evidenza Empirica: Dimostrazione che semplici MLP a uno strato nascosto, addestrati su queste perdite e agganciati a reti pre-addestrate congelate, possono ottenere invarianza quasi perfetta senza fine-tuning delle pesi originali.
4. Confronto Critico: Dimostrazione che metodi contrastivi standard (SimCLR) e basati su indipendenza (HSIC) sono inadeguati per questo scopo specifico perché corrompono lo spazio latente pre-esistente.

5. Significato e Implicazioni

Questo lavoro offre una soluzione pratica ed efficiente per adattare modelli foundation a requisiti di robustezza specifici (es. satelliti, medicina, robotica) senza il costo computazionale del riaddestramento o il rischio di degradazione delle performance.

• Efficienza: Non richiede il fine-tuning dei pesi della rete principale, rendendo il processo leggero e scalabile.
• Versatilità: Funziona su diverse architetture (Transformer, CNN) e tipi di augmentation.
• Robustezza Adversariale: I risultati sull'invarianza al rumore suggeriscono che questi metodi potrebbero essere promettenti per la robustezza avversaria, trattando gli attacchi come augmentation non invertibili.

In sintesi, il paper stabilisce che l'invarianza può essere "aggiunta" a un modello pre-addestrato in modo sicuro e strutturale, utilizzando principi di trasporto ottimo, superando i limiti dei metodi di apprendimento rappresentazionale tradizionali in contesti di adattamento post-addestramento.