Layer by layer, module by module: Choose both for optimal OOD probing of ViT

Questo studio dimostra che per l'analisi OOD dei Vision Transformer è fondamentale scegliere strategicamente sia il livello intermedio che il modulo specifico da sondare, rivelando che l'attivazione interna della rete feed-forward è ottimale in caso di forte shift distributivo, mentre l'output normalizzato dell'attenzione multi-testa lo è quando lo shift è debole.

L'essenziale

Ecco una spiegazione semplice e creativa del paper, pensata per chiunque, anche senza un background tecnico.

Immagina che un Vision Transformer (ViT) sia come un grande chef che ha passato anni a imparare a cucinare (pre-addestramento) guardando milioni di ricette e piatti diversi su Internet (il dataset ImageNet). Questo chef è diventato bravissimo a riconoscere ingredienti e piatti classici.

Ora, il problema è: cosa succede se portiamo questo chef in un ristorante nuovo, con ingredienti strani o un menu completamente diverso (i dati "Out-of-Distribution" o OOD)?

1. Il Problema: L'Ultimo Passo è il Più Fragile

Fino a poco tempo fa, tutti pensavano che per capire un nuovo piatto, bisognasse guardare il piatto finito appena uscito dalla cucina (l'ultimo strato della rete neurale).

• La scoperta: Gli autori di questo studio hanno notato che, se il nuovo ristorante è molto diverso da quello dove lo chef ha imparato, il piatto finale spesso viene rovinato. È come se lo chef, nel tentativo di essere perfetto, si fosse "fissato" troppo sulle regole vecchie e avesse perso la flessibilità.
• L'analogia: Immagina che lo chef, mentre cucina, abbia un'idea brillante a metà strada (uno strato intermedio), ma poi, quando aggiunge gli ultimi condimenti (l'ultimo strato), si confonda perché gli ingredienti sono diversi dal solito.
• La soluzione: Se il menu è molto diverso, è meglio guardare cosa c'è nel piatto a metà cottura (gli strati intermedi). Lì, la "sostanza" del piatto è ancora pura e riconoscibile, prima che lo chef provi a forzare una ricetta sbagliata.

2. L'Esperimento: "Layer by Layer" (Strato per Strato)

Gli autori hanno fatto un esperimento: hanno chiesto a un assistente (un semplice classificatore lineare) di indovinare cosa stava cucinando lo chef guardando il piatto in ogni singolo momento della preparazione.

• Risultato: Se il ristorante è simile a quello originale (In-Distribution), l'ultimo strato vince sempre. Ma se il ristorante è strano (Out-of-Distribution), gli strati intermedi sono molto più robusti e affidabili.

3. Il Dettaglio: "Module by Module" (Modulo per Modulo)

Qui la cosa diventa ancora più interessante. All'interno di ogni "stazione" di cucina (un blocco del transformer), ci sono diversi passaggi:

1. LN (LayerNorm): Come mettere il grembiule e prepararsi.
2. MHA (Self-Attention): Lo chef che guarda gli ingredienti e decide cosa combinare.
3. FFN (Feed-Forward Network): La parte dove lo chef effettivamente lavora gli ingredienti (taglia, mescola, cuoce). Questo modulo ha due fasi:
   • FC1: Allarga gli ingredienti (come se prendesse un pezzo di carne e lo sminuzzasse in mille pezzetti per analizzarlo meglio).
   • Act (Attivazione): Il momento in cui si decide cosa è importante tra tutti quei pezzetti (il "filtro").
   • FC2: Ricompatta tutto per tornare alla dimensione originale.

La scoperta chiave:

• Se il menu è molto strano (forte spostamento): Non guardare il piatto finito (FC2) né la preparazione generale (RC2). Guarda mentre lo chef sta analizzando gli ingredienti sminuzzati (l'attivazione Act). È lì che si trova l'informazione più pura e utile, prima che venga "schiacciata" di nuovo.
• Se il menu è quasi uguale (debole spostamento): Va bene guardare il risultato della preparazione standard (LN2), che è più stabile.

In Sintesi: Cosa dobbiamo imparare?

Immagina di dover giudicare un'opera d'arte fatta da un artista che ha cambiato stile.

1. Non guardare solo il quadro finito: Se l'artista è stato confuso dal nuovo stile, il quadro finale potrebbe essere un disastro.
2. Guarda il processo: Spesso, il momento più chiaro e intelligente è stato a metà del lavoro, quando l'artista aveva ancora le idee chiare prima di cercare di adattarsi forzatamente.
3. Scegli il punto giusto:
   • Se il cambiamento è gigantesco, guarda il momento in cui l'artista sta "pensando" (l'attivazione nel mezzo del processo).
   • Se il cambiamento è leggero, puoi guardare il risultato standard.

Perché è importante?
Questo studio ci dice che quando usiamo l'Intelligenza Artificiale nel mondo reale (dove le cose cambiano sempre), non dobbiamo fidarci ciecamente dell'ultima risposta che ci dà il modello. Dobbiamo essere più curiosi e guardare "dentro" il modello, scegliendo il momento giusto per fare domande, per ottenere risultati più precisi e affidabili.

È come dire: "Non chiedere allo chef cosa ha cucinato solo quando ha finito di servire. Chiedigli cosa stava pensando mentre tagliava le verdure, specialmente se gli ingredienti erano strani!"

Sommario tecnico

Ecco un riassunto tecnico dettagliato del paper "LAYER BY LAYER, MODULE BY MODULE: CHOOSE BOTH FOR OPTIMAL OOD PROBING OF VIT", presentato al workshop CAO di ICLR 2026.

1. Il Problema

I modelli fondazione basati su architetture Transformer (in particolare i Vision Transformers o ViT) hanno ottenuto risultati eccezionali, ma la loro affidabilità in ambienti reali è minacciata dagli spostamenti di distribuzione (distribution shifts). Quando i dati di downstream (test) differiscono dai dati di pre-addestramento (Out-of-Distribution o OOD), le rappresentazioni estratte dagli strati finali del modello tendono a degradare drasticamente.

Studi precedenti (es. Skean et al., 2025) hanno suggerito che negli LLM gli strati intermedi sono superiori a quelli finali, ma per i ViT le conclusioni sono state contrastanti: alcuni lavori indicano che lo strato finale rimane ottimale, mentre altri mostrano il contrario. Inoltre, la maggior parte delle analisi si è limitata a scenari In-Distribution (ID) o non ha indagato la struttura interna dei blocchi del Transformer, concentrandosi solo sull'output del blocco completo.

2. Metodologia

Gli autori conducono uno studio empirico approfondito su un ViT pre-addestrato da 86M parametri (su ImageNet-21k) utilizzando un approccio di linear probing (sondaggio lineare).

• Dataset: Sono stati utilizzati 11 benchmark di classificazione diversificati, inclusi dataset ID (CIFAR-10, CIFAR-100, Flowers102, Pets) e dataset OOD con vari gradi di shift (varianti di CIFAR-10-C come Rumore, Sfocatura, Neve; e domini di DomainNet come Clipart e Sketch).
• Analisi "Layer by Layer": Invece di sondare solo l'ultimo strato, gli autori valutano le prestazioni di linear probing su ogni strato intermedio del ViT, confrontando modelli pre-addestrati (frozen) con modelli fine-tuned.
• Analisi "Module by Module": All'interno di ogni blocco Transformer, gli autori sondano i punti di uscita specifici di ogni componente:
  • Normalizzazione (LN1, LN2)
  • Multi-Head Attention (MHA)
  • Connessioni Residuali (RC1, RC2)
  • Feed-Forward Network (FFN): specificamente gli input/uscite dei layer FC1, FC2 e l'attivazione (Act/GeLU) intermedia.
• Protocollo: Per ogni rappresentazione nascosta, viene estratto il token CLS, fatto pooling e utilizzato per addestrare un classificatore di regressione logistica (L-BFGS).

3. Contributi Chiave e Risultati

A. Lo Spostamento di Distribuzione è la Causa Principale

L'analisi conferma che il degrado delle prestazioni negli strati finali è direttamente correlato all'entità dello spostamento di distribuzione tra i dati di pre-addestramento e quelli di downstream.

• Scenario ID: Quando i dati sono simili a quelli di pre-addestramento, lo strato finale offre le migliori prestazioni.
• Scenario OOD: All'aumentare dello shift, le prestazioni degli strati finali crollano, mentre gli strati intermedi si rivelano significativamente più robusti. Questo smentisce l'idea che il beneficio degli strati intermedi sia solo un sottoprodotto dell'obiettivo di pre-addestramento autoregressivo; è invece una conseguenza della necessità di robustezza allo shift.

B. Non Tutti i Moduli sono Uguali (Module-Level Analysis)

L'analisi granulare rivela che sondare l'output standard del blocco Transformer (RC2, la connessione residua finale) è subottimale nella maggior parte dei casi OOD. Le prestazioni variano drasticamente a seconda del modulo sondato:

1. FC2 (Output del FFN): È il modulo peggiore da sondare, specialmente sotto forte shift di distribuzione, poiché comprime le informazioni riducendo la separabilità lineare.
2. Act (Attivazione GeLU intermedia del FFN): Si rivela il modulo migliore in scenari OOD con shift significativo. Operando in uno spazio di dimensionalità più alta (4d) e filtrando il rumore prima della compressione finale, cattura informazioni semantiche più ricche e disaccoppiate.
3. LN2 (LayerNorm prima del FFN): Offre risultati stabili e leggermente superiori a RC2, risultando una scelta sicura quando lo shift è difficile da rilevare o debole.
4. RC2 (Output standard): Funziona bene solo in scenari ID o con shift molto lievi.

C. Takeaway Pratici

Il paper fornisce due linee guida operative:

1. In setting ID: Sondare sempre gli strati finali.
2. In setting OOD: Sondare le attivazioni intermedie (Act) dei layer Feed-Forward negli strati intermedi della rete. Se lo shift è incerto, sondare LN2 è un approccio più sicuro rispetto allo standard RC2.

4. Significato e Implicazioni

Questo lavoro offre una nuova prospettiva fondamentale sull'uso delle rappresentazioni nascoste nei Vision Transformers:

• Robustezza: Dimostra che la "ricchezza" delle rappresentazioni non è monotona rispetto alla profondità della rete quando si affrontano dati OOD.
• Efficienza: Fornisce una strategia a basso costo computazionale (linear probing su moduli specifici) per migliorare le prestazioni senza bisogno di un costoso fine-tuning completo, cruciale per modelli fondazione dove i dati di pre-addestramento non sono accessibili.
• Guida per l'Industria: Suggerisce che per applicazioni reali soggette a drift di distribuzione, le strategie di estrazione delle feature devono essere dinamiche e basate sulla comprensione della struttura interna del Transformer, spostando l'attenzione dagli output finali alle attivazioni interne dei layer Feed-Forward.

In sintesi, il paper stabilisce che per ottenere la massima robustezza OOD nei ViT, non basta scegliere l'layer giusto, ma bisogna scegliere il modulo giusto all'interno di quel layer: l'attivazione del Feed-Forward Network è la chiave per la resilienza allo shift.