Riepilogo Tecnico: Inversione di Trasformazioni di Dati tramite Campionamento di Diffusione

Definizione del Problema

Il documento affronta il problema dell'inversione di trasformazione su gruppi di Lie generali. In molti scenari di apprendimento automatico e modellazione scientifica, i dati osservati $\tilde{x}$ sono generati applicando una trasformazione ignota $g$ appartenente a un gruppo di Lie $G$ a un campione in-distribuzione $x$ (ovvero, $\tilde{x} = g \cdot x$). L'obiettivo è recuperare una trasformazione inversa $g^{-1}$ che mappi l'osservazione fuori distribuzione alla distribuzione originale dei dati.

Questo viene formulato come un problema inverso cieco (blind inverse problem) poiché la trasformazione di disturbo è ignota e la soluzione è generalmente non univoca. Mentre i metodi esistenti operano spesso nello spazio dei dati o sono specializzati per specifiche famiglie di trasformazioni (ad esempio, gruppi affini o di rotazione), questo lavoro mira alle azioni di gruppi di Lie generali, inclusi gruppi non compatti e non Abeliani privi di metriche bi-invarianti (ad esempio, omografie di immagini $PGL(3, \mathbb{R})$ e simmetrie di punto di Lie per PDE).

Una chiave applicazione evidenziata è l'equivarianza al tempo di test (test-time equivariance). Le reti neurali pre-addestrate spesso mancano di robustezza rispetto alle trasformazioni geometriche. Gli autori propongono di invertire queste trasformazioni al tempo di test per riportare gli input nella distribuzione di addestramento del modello, migliorando così la robustezza senza richiedere architetture equivarianti specializzate o nuovi addestramenti.

Metodologia: Trasformazione-Invertente Diffusion Energy (TIED)

Il metodo proposto, TIED, modella la distribuzione a posteriori sulle trasformazioni $p(g | \tilde{x})$ come una distribuzione di Boltzmann definita da una funzione di energia $E$ sullo spazio dei dati. L'energia è derivata da un prior dei dati $p_x(x)$, dove $E(x) = -\log p_x(x)$. La densità a posteriori è data da:
$$p(g | \tilde{x}) \propto p_x(g^{-1} \cdot \tilde{x}) |\det J_{g^{-1}}(\tilde{x})|$$

Per campionare da questa a posteriori, TIED introduce un processo di diffusione direttamente sul gruppo di Lie. La metodologia supera due sfide principali: la natura accidentata e multimodale dei paesaggi energetici (spesso derivati da reti neurali) e i vincoli geometrici del campionamento su varietà (manifolds).

1. Diffusione Trivializzata sui Gruppi di Lie

Invece di eseguire aggiornamenti negli spazi tangenti di arbitrarie componenti del gruppo, TIED utilizza la trivializzazione. Gli aggiornamenti vengono eseguiti nell'algebra di Lie $\mathfrak{g}$ (lo spazio tangente l'identità) e trasportati all'elemento del gruppo corrente tramite la mappa esponenziale.

• Processo Forward: Un SDE (Stochastic Differential Equation) a varianza esplosiva viene definito sul gruppo disegnando perturbazioni infinitesimali nell'algebra di Lie e trasportandole tramite moltiplicazione a destra:
  $$dg_t = d(L_{g_t})_e [\gamma(t) dw^{\mathfrak{g}}_t]$$
  dove $w^{\mathfrak{g}}_t$ è il moto browniano nell'algebra di Lie.
• Processo Reverse: L'SDE nel tempo inverso è derivato per recuperare campioni dalla distribuzione a posteriori. Fondamentalmente, il termine di drift coinvolge lo score trivializzato $\nabla_g \log p_t(g_t)$, che è un elemento dell'algebra di Lie.

2. Identità dello Score Target Trivializzato

Un contributo teorico centrale è una nuova identità dello score target trivializzato. Lo standard score matching richiede l'accesso a campioni puliti, che non sono disponibili in questo contesto di inversione. Invece, TIED esprime lo score della distribuzione rumorosa in termini del gradiente della funzione di energia pulita.
Gli autori dimostrano che per gruppi di Lie generali (senza richiedere una metrica bi-invariante), lo score trivializzato può essere stimato tramite un'aspettativa Monte Carlo sulla kernel di rumore $k_t$:
$$\nabla_g \log p_t(g_t) = \nabla_g \log \int_G k_t(w) \exp\left(-E(g_t w^{-1}) - \log \lambda(w)\right) d\mu(w)$$
dove $\lambda(w)$ tiene conto del cambiamento della misura di Haar sotto inversione. Ciò consente di stimare lo score utilizzando solo la funzione di energia $E$ e i suoi gradienti, calcolati tramite differenziazione automatica.

3. Implementazione Pratica

L'algoritmo di campionamento (Algoritmo 1) discretizza l'SDE inversa. In ogni step, esso:

1. Campiona il rumore $w$ dalla kernel di rumore $k_t$.
2. Calcola l'energia $E$ per i candidati trasformati.
3. Stima lo score trivializzato usando l'identità Monte Carlo (impiegando il trucco log-sum-exp per la stabilità).
4. Aggiorna l'elemento del gruppo usando la mappa esponenziale.

Il metodo supporta varie definizioni di energia, incluse quelle derivate da Autoencoder Variazionali (VAE) o direttamente dalla confidenza di predizione di un classificatore pre-addestrato.

Contributi Chiave

1. Campionatore TIED: Un nuovo campionatore di diffusione per invertire le trasformazioni dei dati su gruppi di Lie generali. Garantisce che gli aggiornamenti rimangano sulla varietà e gestisce azioni non compatte e non lineari utilizzando solo la conoscenza dell'algebra di Lie.
2. Framework Teorico: Una prova che l'inversione della trasformazione può essere utilizzata per l'equivarianza al tempo di test dei modelli pre-addestrati. Gli autori dimostrano che la distribuzione a posteriori dell'inversione della trasformazione è $G$-equivariante, garantendo la solidità dell'approccio.
3. Nuova Identità dello Score: Una generalizzazione dell'identità dello score target ai gruppi di Lie che funziona senza metriche bi-invarianti, consentendo un campionamento efficiente basato sullo score tramite i gradienti dell'energia anziché campioni di dati puliti.
4. Validazione Empirica: Dimostrazioni su omografie di immagini ($PGL(3, \mathbb{R})$) e simmetrie di punti di Lie per PDE ($SL(2, \mathbb{R}) \ltimes H(1, \mathbb{R})$), mostrando miglioramenti costanti rispetto a forti baseline.

Risultati Sperimentali

Il documento valuta TIED su tre compiti principali:

1. Campionamento Sintetico su $SO(10)$: TIED campiona con successo da una distribuzione di energia multimodale su un gruppo di Lie ad alta dimensione ($dim=45$), convergendo significativamente più velocemente di un campionatore Langevin cinetico trivializzato.
2. Classificazione di Immagini (MNIST):
   • Setup: Un ResNet18 pre-addestrato viene testato su immagini MNIST sottoposte a trasformazioni affini ($Aff(2, \mathbb{R})$) e di omografia ($PGL(3, \mathbb{R})$) ignote.
   • Baseline: Confrontato con metodi basati su ottimizzazione (LieLAC, FoCal, ITS) e metodi di campionamento (Langevin Cinetico).
   • Risultati: TIED supera significativamente tutte le baseline. Ad esempio, sotto trasformazioni affini con energia basata sul classificatore, TIED raggiunge un'accuratezza dell'82,64% rispetto al 55,48% del ResNet18 grezzo e al 74,72% del Langevin Cinetico. Supera anche reti equivarianti specializzate (affConv) in alcuni contesti.
3. Risoluzione di PDE:
   • Setup: Un operatore neurale DeepONet viene testato su equazioni del calore e di Burgers 1D sotto trasformazioni di simmetria di punto di Lie ignote.
   • Risultati: TIED raggiunge l'errore relativo $L_2$ più basso in tutti i casi di test, superando LieLAC e FoCal. Per l'equazione del calore 1D con simmetrie $SL(2, \mathbb{R}) \ltimes H(1, \mathbb{R})$, TIED riduce l'errore a 0,043, rispetto allo 0,066 di LieLAC e allo 0,690 del DeepONet non corretto.

Significato e Rivendicazioni

Gli autori affermano che TIED fornisce una pipeline senza addestramento (training-free pipeline) per rendere qualsiasi modello pre-addestrato equivariante alle trasformazioni di gruppi di Lie al tempo di test. Sfruttando i modelli basati sull'energia (che possono essere approssimati dal modello stesso pre-addestrato), il metodo evita la necessità di architetture equivarianti specializzate o costosi ri-addestramenti.

La significatività risiede in:

• Generalità: Si applica a una vasta classe di gruppi di Lie, inclusi quelli non compatti e privi di metriche bi-invarianti, che hanno rappresentato una sfida per precedenti metodi di diffusione o ottimizzazione.
• Robustezza: Restaura efficacemente gli input trasformati alla distribuzione di addestramento, migliorando significativamente la generalizzazione delle reti pre-addestrate alle trasformazioni geometriche fuori distribuzione.
• Efficienza: L'approccio basato sulla diffusione, combinato con l'identità dello score trivializzato, offre un processo di campionamento più efficiente e stabile rispetto ai tradizionali MCMC o all'ottimizzazione iterativa su paesaggi energetici accidentati.

Il documento conclude che questo approccio rappresenta una forma di scaling dell'equivarianza al tempo di inferenza, offrendo una soluzione pratica per la robustezza nell'era dei grandi modelli di fondazione dove la gestione delle trasformazioni fuori distribuzione è critica.