Riepilogo Tecnico: T3R - Adattamento Profondo al Tempo di Test per le Reti Neurali a Grafo tramite Rotazione del Gradiente

1. Definizione del Problema

Le Reti Neurali a Grafo (GNN) impiegate in sistemi del mondo reale, come reti di distribuzione idrica, infrastrutture 5G e reti informatiche, operano tipicamente con pesi fissi dopo l'addestramento. Questo paradigma statico porta a un degrado delle prestazioni quando si incontrano spostamenti di distribuzione (concept drift), un fenomeno comune nei dati dei sensori dove le condizioni cambiano nel tempo.

Sebbene il Test-Time Training (TTT) offra una soluzione adattando i modelli utilizzando dati di test non etichettati tramite un compito ausiliario di Apprendimento Auto-Supervisionato (SSL), gli approcci esistenti affrontano limitazioni significative nel dominio dei grafi:

1. Adattamento Superficiale: I metodi TTT convenzionali (ad es., Sun et al., 2020) aggiornano tipicamente solo l'encoder condiviso e il decoder SSL durante l'inferenza, lasciando fisso il decoder del compito principale. Ciò limita la capacità del modello di adattarsi completamente a nuove distribuzioni.
2. Disallineamento dei Compiti: Selezionare un compito SSL ausiliario che si allinei significativamente con un compito specifico basato su grafi è difficile. I conflitti di gradiente tra i compiti possono ostacolare l'ottimizzazione congiunta.
3. Scarsità di Dati: In domini come i sistemi idrici, i dati etichettati sono scarsi e il riaddestramento o il fine-tuning tramite simulazioni è spesso impraticabile a causa della mancanza di una conoscenza completa del sistema.

2. Metodologia: T3R

Gli autori propongono T3R (Test-Time Training con Rotograd per livello), un nuovo framework che sfrutta la rotazione del gradiente per consentire un adattamento profondo attraverso quasi l'intera architettura GNN senza etichette di verità fondamentale (ground-truth).

Architettura Core

T3R utilizza un'architettura a forma di Y composta da:

• Un encoder condiviso ($\theta_e$).
• Una testa principale ($\theta_m$) per il compito target.
• Una testa SSL ausiliaria ($\theta_s$) per un obiettivo privo di etichette (ad es., mascheramento dei nodi).

Meccanismi Chiave

1. Allineamento del Gradiente e Affinità del Compito:
Il metodo misura l'affinità del compito come la similarità del coseno tra i vettori di gradiente delle teste principale e SSL. Un alto allineamento positivo indica che la minimizzazione di un obiettivo riduce anche l'altro. Gli autori selezionano empiricamente i compiti SSL (specificamente il mascheramento dei nodi) che mostrano il più alto allineamento positivo del gradiente con il compito principale.

2. Rotograd per Livello:
Per mitigare il conflitto di gradiente e migliorare l'affinità del compito, T3R integra Rotograd (Javaloy & Valera, 2022) nel framework TTT.

• Fase di Addestramento: Matrici di rotazione ortogonali ($R^{(l)}_m, R^{(l)}_s$) vengono preposte a ciascuno strato dei decoder principale e SSL. Una perdita di rotazione ($L_{rot}$) viene ottimizzata ciclicamente per massimizzare la similarità del coseno tra i gradienti originali e i gradienti ruotati, specifici per il compito. Ciò incoraggia i gradienti di entrambi i compiti ad allinearsi in una direzione comune.
• Fase di Test-Time: Durante l'inferenza, la perdita del compito principale non è disponibile. T3R utilizza le matrici di rotazione apprese per generare gradienti surrogati per il compito principale. Nello specifico, il gradiente SSL viene ruotato verso la direzione media del gradiente originale e reindirizzato per aggiornare gli strati del decoder principale.
  • Formula: $\tilde{g}^{(l+1)}_m = R^{(l)\top}_s g^{(l)}_s R^{(l)}_m$.
  • Questo meccanismo consente di aggiornare l'intera architettura (encoder condiviso, decoder SSL e quasi tutti gli strati del decoder principale), invece di limitarsi all'encoder.

3. Spostamenti degli Iperparametri:
Gli autori osservano che l'aggiustamento degli iperparametri durante la fase di test migliora le prestazioni. Nello specifico, aumentano il coefficiente del compito ausiliario ($\lambda_{test}$) e riducono l'intensità dell'augmentation (ad es., tassi di mascheramento dei nodi più bassi) per amplificare i segnali SSL e prevenire il fallimento su piccoli batch di test.

3. Contributi Principali

Il documento delinea tre contributi primari:

1. Metodo T3R: Un nuovo approccio di adattamento che utilizza matrici di rotazione per livello per migliorare l'affinità del compito durante l'addestramento e generare gradienti surrogati per un adattamento profondo e su tutta l'architettura durante il test.
2. Selezione Empirica SSL: Un'indagine sulla selezione di obiettivi SSL basati su grafi che si allineano con i compiti principali, dimostrando che l'allineamento positivo del gradiente correla con una migliore adattabilità.
3. Valutazione Estesa: Esperimenti completi su dataset di regressione (reti di distribuzione idrica) e classificazione (grafi molecolari) in scenari Out-of-Distribution (OOD) e cross-domain.

4. Risultati Sperimentali

Gli autori hanno valutato T3R sul dataset DiTEC-WDN (DWD) (regressione) e sui dataset Open Graph Benchmark (OGB) (classificazione).

Regressione (Reti di Distribuzione Idrica)

• Prestazioni: T3R ha ottenuto un miglioramento relativo del 38,38% nell'RMSE rispetto alla standard Empirical Risk Minimization (ERM) e ha superato altri baseline di adattamento (Tent, TTT standard, TTT+Rotograd) attraverso molteplici metriche (MAE, R2, PCC).
• Osservazione: Il TTT standard e TTT+Rotograd (che aggiorna solo l'encoder) hanno mostrato guadagni limitati o sono stati inferiori a ERM in questo contesto, evidenziando la necessità di aggiornare il decoder principale.
• Cross-Validazione: Nei test cross-domain (addestramento su una topologia di rete, test su un'altra), T3R con un singolo step di adattamento ($t=1$) ha prodotto i migliori risultati. Un adattamento più profondo (più step) ha generalmente degradato le prestazioni nei compiti di regressione.

Classificazione (Grafi Molecolari)

• Adattamento Cross-Domain: Negli esperimenti di adattamento tra ogbg-molbace e ogbg-molbbbp, T3R ha ottenuto miglioramenti relativi significativi nella ROC-AUC, raggiungendo fino al 69,56% di miglioramento relativo rispetto all'inferenza standard senza adattamento.
• Confronto: Mentre altri metodi mostravano lievi miglioramenti o una convergenza precoce, T3R ha dimostrato guadagni robusti, specialmente quando gli step di adattamento erano accuratamente calibrati.

Studi di Ablazione

• Strati di Rotazione: La rimozione degli strati di rotazione ha aumentato l'incertezza e degradato le prestazioni, confermando la necessità dell'adattamento per livello.
• Coefficiente di Perdita: La calibrazione del coefficiente della perdita ausiliaria del tempo di test ($\lambda_{test}$) per essere superiore a quello di addestramento ha migliorato le prestazioni del compito principale, mentre valori eccessivi hanno portato a un sovra-adattamento.

5. Significato e Rivendicazioni

Il documento sostiene che T3R affronta il limite critico dei metodi TTT esistenti abilitando un adattamento più profondo attraverso l'intera architettura del modello, non solo l'encoder. Sfruttando la rotazione del gradiente per creare gradienti surrogati, T3R consente ai modelli di adattarsi efficacemente in scenari in cui:

• I dati etichettati non sono disponibili o sono costosi da raccogliere.
• Gli spostamenti di distribuzione sono frequenti (ad es., deriva dei sensori nelle infrastrutture).
• Il fine-tuning o il riaddestamento convenzionale è impraticabile.

Gli autori posizionano T3R come un passo verso lo sviluppo di pipeline di adattamento robuste e prive di etichette per sistemi basati su grafi in tempo reale, in particolare in domini di infrastrutture critiche come le reti idriche. Riconoscono che, sebbene il metodo sia promettente, la relazione tra l'allineamento del gradiente e l'adattabilità richiede una ulteriore giustificazione teorica, e il costo computazionale degli ulteriori aggiornamenti del gradiente è un compromesso per l'aumento dell'accuratezza.