Technische Samenvatting: Inverteren van Datatransformaties via Diffusie-sampling

Probleemstelling

Het artikel behandelt het probleem van transformatie-inversie op algemene Lie-groepen. In veel scenario's van machine learning en wetenschappelijke modellering wordt geobserveerde data $\tilde{x}$ gegenereerd door een onbekende transformatie $g$ uit een Lie-groep $G$ toe te passen op een in-distributie sample $x$ (d.w.z. $\tilde{x} = g \cdot x$). Het doel is om een inverse transformatie $g^{-1}$ te herstellen die de out-of-distribution observatie terugbrengt naar de oorspronkelijke datadistributie.

Dit wordt geformuleerd als een blind invers probleem omdat de nuisance-transformatie onbekend is en de oplossing over het algemeen niet uniek is. Terwijl bestaande methoden vaak opereren in de dataruimte of gespecialiseerd zijn voor specifieke transformatiefamilies (bijv. affiene of rotatiegroepen), richt dit werk zich op algemene Lie-groepacties, inclusief niet-compacte, niet-Abelse groepen zonder bi-invariante metrieken (bijv. beeldhomografieën $PGL(3, \mathbb{R})$ en Lie-puntensymmetrieën voor PDE's).

Een belangrijke toepassing die wordt belicht is test-time equivariantie. Pretrained neurale netwerken missen vaak robuustheid tegen geometrische transformaties. De auteurs stellen voor om deze transformaties tijdens de testtijd te inverteren om inputs terug te brengen naar de trainingsdistributie van het model, waardoor de robuustheid wordt verbeteren zonder dat gespecialiseerde equivariante architecturen of hertraining vereist zijn.

Methodologie: Transformation-Inverting Energy Diffusion (TIED)

De voorgestelde methode, TIED, modelleert de posterior over transformaties $p(g | \tilde{x})$ als een Boltzmann-distributie gedefinieerd door een energiefunctie $E$ op de dataruimte. De energie is afgeleid van een data-prior $p_x(x)$, waarbij $E(x) = -\log p_x(x)$. De posterior dichtheid wordt gegeven door:
$$p(g | \tilde{x}) \propto p_x(g^{-1} \cdot \tilde{x}) |\det J_{g^{-1}}(\tilde{x})|$$

Om te samplen uit deze posterior introduceert TIED een diffusieproces direct op de Lie-groep. De methodologie overwint twee belangrijke uitdagingen: het ruige, multimodale karakter van energielandschappen (vaak afgeleid van neurale netwerken) en de geometrische beperkingen van sampling op manifolds.

1. Trivialized Diffusion op Lie-groepen

In plaats van updates uit te voeren in de tangentieruimten van willekeurige groepselementen, maakt TIED gebruik van trivialisatie. Updates worden uitgevoerd in de Lie-algebra $\mathfrak{g}$ (de tangentieruimte bij de identiteit) en getransporteerd naar het huidige groepselement via de exponentiële afbeelding.

• Forward Proces: Een variance-exploding Stochastic Differential Equation (SDE) wordt gedefinieerd op de groep door infinitesimale perturbaties in de Lie-algebra te trekken en deze via recht vermenigvuldigen te transporteren:
  $$dg_t = d(L_{g_t})_e [\gamma(t) dw^{\mathfrak{g}}_t]$$
  waarbij $w^{\mathfrak{g}}_t$ Brownse beweging is in de Lie-algebra.
• Reverse Proces: De reverse-time SDE is afgeleid om samples te herstellen uit de posterior. Cruciaal is dat de driftterm de trivialized score $\nabla_g \log p_t(g_t)$ bevat, wat een element van de Lie-algebra is.

2. Trivialized Target Score Identiteit

Een kernbijdrage is een nieuwe trivialized target score identiteit. Standaard score matching vereist toegang tot schone samples, die in deze inversie-setting niet beschikbaar zijn. In plaats daarvan drukt TIED de score van de ruisige distributie uit in termen van de gradiënt van de schone energiefunctie.
De auteurs bewijzen dat voor algemene Lie-groepen (zonder een bi-invariante metriek te vereisen), de trivialized score geschat kan worden via een Monte Carlo-verwachting over de ruiskern $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)$$
waarbij $\lambda(w)$ de verandering van de Haar-maat onder inversie account. Dit maakt het mogelijk om de score te schatten met behulp van enkel de energiefunctie $E$ en de bijbehorende gradiënten, berekend via automatische differentiatie.

3. Praktische Implementatie

Het sampling algoritme (Algoritme 1) discretiseert de reverse SDE. In elke stap:

1. Sample ruis $w$ uit de ruiskern $k_t$.
2. Bereken de energie $E$ voor getransformeerde kandidaten.
3. Schat de trivialized score met behulp van de Monte Carlo-identiteit (waarbij de log-sum-exp truc wordt gebruikt voor stabiliteit).
4. Update het groepselement met behulp van de exponentiële afbeelding.

De methode ondersteunt diverse energie-definities, waaronder die afgeleid van Variational Autoencoders (VAEs) of direct van de voorspellingszekerheid van een pretrained classifier.

Belangrijkste Bijdragen

1. TIED Sampler: Een nieuwe diffusie-sampler voor het inverteren van datatransformaties op algemene Lie-groepen. Het garandeert dat updates op de manifold blijven en handelt niet-compacte, niet-lineaire acties af met enkel kennis van de Lie-algebra.
2. Theoretisch Kader: Een bewijs dat transformatie-inversie gebruikt kan worden voor test-time equivariantie van pretrained modellen. De auteurs demonstreren dat de transformatie-inversie posterior $G$-equivariant is, wat de validiteit van de aanpak waarborgt.
3. Nieuwe Score Identiteit: Een generalisatie van de target score identiteit naar Lie-groepen die werkt zonder bi-invariante metrieken, wat efficiënte score-based sampling mogelijk maakt met energie-gradiënten in plaats van schone datastatistieken.
4. Empirische Validatie: Demonstraties op beeldhomografieën ($PGL(3, \mathbb{R})$) en PDE-symmetrieën ($SL(2, \mathbb{R}) \ltimes H(1, \mathbb{R})$), die consistente prestatieverbeteringen laten zien ten opzichte van sterke baselines.

Experimentele Resultaten

Het artikel evalueert TIED op drie hoofdtaken:

1. Synthetische Sampling op $SO(10)$: TIED samplet succesvol uit een multimodale energiedistributie op een hoog-dimensionale Lie-groep ($dim=45$), en convergeert significant sneller dan een triviale kinetische Langevin sampler.
2. Beeldclassificatie (MNIST):
   • Setup: Een pretrained ResNet18 wordt getest op MNIST-beelden onderhevig aan onbekende affiene ($Aff(2, \mathbb{R})$) en homografie ($PGL(3, \mathbb{R})$) transformaties.
   • Baselines: Vergeleken met optimalisatie-gebaseerde methoden (LieLAC, FoCal, ITS) en sampling methoden (Kinetic Langevin).
   • Resultaten: TIED presteert significant beter dan alle baselines. Bijvoorbeeld, onder affiene transformaties met een classifier-gebaseerde energie, bereikt TIED een nauwkeurigheid van 82,64%, vergeleken met 55,48% voor de ruwe ResNet18 en 74,72% voor Kinetic Langevin. Het presteert ook beter dan gespecialiseerde equivariante netwerken (affConv) in bepaalde settings.
3. PDE Oplossen:
   • Setup: Een DeepONet neurale operator wordt getest op 1D Heat en Burgers vergelijkingen onder onbekende Lie-puntensymmetrie transformaties.
   • Resultaten: TIED bereikt de laagste relatieve $L_2$ fout over alle testgevallen, waarbij het de LieLAC en FoCal verslaat. Voor de 1D Heat vergelijking met $SL(2, \mathbb{R}) \ltimes H(1, \mathbb{R})$ symmetrieën, reduceert TIED de fout naar 0,043, vergeleken met 0,066 voor LieLAC en 0,690 voor de ongecorrigeerde DeepONet.

Betekenis en Claims

De auteurs claimen dat TIED een training-vrije pipeline biedt om elk pretrained model tijdens de testtijd equivariant te maken voor Lie-groep transformaties. Door gebruik te maken van energy-based modellen (die benaderd kunnen worden door het pretrained model zelf), vermijdt de methode de noodzaak voor gespecialiseerde equivariante architecturen of dure hertraining.

De betekenis ligt in:

• Generaliteit: Het is toepasbaar op een brede klasse van Lie-groepen, inclus\n bij niet-compacte groepen die geen bi-invariante metrieken bezitten, wat een uitdaging vormde voor eerdere diffusie of optimalisatiemethoden.
• Robuustheid: Het herstelt effectief getransformeerde inputs naar de trainingsdistributie, wat de generalisatie van pretrained netwerken naar out-of-distribution geometrische transformaties aanzienlijk verbetert.
• Efficiëntie: De diffusie-gebaseerde aanpak, gecombineerd met de trivialized score identiteit, biedt een efficiënter en stabieler sampling proces vergeleken met traditionele MCMC of iteratieve optimalisatie op ruige energielandschappen.

Het artikel concludeert dat deze aanpak een vorm van inference-time scaling voor equivariantie vertegenwoordigt, wat een praktische oplossing biedt voor robuustheid in het tijdperk van grote foundation models waar het omgaan met out-of-distribution transformaties cruciaal is.