Technical Summary: Inverting Data Transformations via Diffusion Sampling

Problem Statement

The paper addresses the problem of transformation inversion on general Lie groups. In many machine learning and scientific modeling scenarios, observed data $\tilde{x}$ is generated by applying an unknown transformation $g$ from a Lie group $G$ to an in-distribution sample $x$ (i.e., $\tilde{x} = g \cdot x$). The goal is to recover an inverse transformation $g^{-1}$ that maps the out-of-distribution observation back to the original data distribution.

This is formulated as a blind inverse problem because the nuisance transformation is unknown and the solution is generally non-unique. While existing methods often operate in the data space or are specialized to specific transformation families (e.g., affine or rotation groups), this work targets general Lie group actions, including non-compact, non-Abelian groups without bi-invariant metrics (e.g., image homographies $PGL(3, \mathbb{R})$ and Lie point symmetries for PDEs).

A key application highlighted is test-time equivariance. Pretrained neural networks often lack robustness to geometric transformations. The authors propose inverting these transformations at test time to bring inputs back into the model's training distribution, thereby improving robustness without requiring specialized equivariant architectures or retraining.

Methodology: Transformation-Inverting Energy Diffusion (TIED)

The proposed method, TIED, models the posterior over transformations $p(g | \tilde{x})$ as a Boltzmann distribution defined by an energy function $E$ on the data space. The energy is derived from a data prior $p_x(x)$, where $E(x) = -\log p_x(x)$. The posterior density is given by:
$$p(g | \tilde{x}) \propto p_x(g^{-1} \cdot \tilde{x}) |\det J_{g^{-1}}(\tilde{x})|$$

To sample from this posterior, TIED introduces a diffusion process directly on the Lie group. The methodology overcomes two main challenges: the rugged, multimodal nature of energy landscapes (often derived from neural networks) and the geometric constraints of sampling on manifolds.

1. Trivialized Diffusion on Lie Groups

Instead of performing updates in the tangent spaces of arbitrary group elements, TIED utilizes trivialization. Updates are performed in the Lie algebra $\mathfrak{g}$ (the tangent space at the identity) and transported to the current group element via the exponential map.

• Forward Process: A variance-exploding Stochastic Differential Equation (SDE) is defined on the group by drawing infinitesimal perturbations in the Lie algebra and transporting them via right multiplication:
  $$dg_t = d(L_{g_t})_e [\gamma(t) dw^{\mathfrak{g}}_t]$$
  where $w^{\mathfrak{g}}_t$ is Brownian motion in the Lie algebra.
• Reverse Process: The reverse-time SDE is derived to recover samples from the posterior. Crucially, the drift term involves the trivialized score $\nabla_g \log p_t(g_t)$, which is an element of the Lie algebra.

2. Trivialized Target Score Identity

A core theoretical contribution is a new trivialized target score identity. Standard score matching requires access to clean samples, which are unavailable in this inversion setting. Instead, TIED expresses the score of the noisy distribution in terms of the gradient of the clean energy function.
The authors prove that for general Lie groups (without requiring a bi-invariant metric), the trivialized score can be estimated via a Monte Carlo expectation over the noise kernel $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)$$
where $\lambda(w)$ accounts for the change of the Haar measure under inversion. This allows the score to be estimated using only the energy function $E$ and its gradients, computed via automatic differentiation.

3. Practical Implementation

The sampling algorithm (Algorithm 1) discretizes the reverse SDE. At each step, it:

1. Samples noise $w$ from the noise kernel $k_t$.
2. Computes the energy $E$ for transformed candidates.
3. Estimates the trivialized score using the Monte Carlo identity (employing the log-sum-exp trick for stability).
4. Updates the group element using the exponential map.

The method supports various energy definitions, including those derived from Variational Autoencoders (VAEs) or directly from the prediction confidence of a pretrained classifier.

Key Contributions

1. TIED Sampler: A novel diffusion sampler for inverting data transformations on general Lie groups. It guarantees updates remain on the manifold and handles non-compact, non-linear actions using only knowledge of the Lie algebra.
2. Theoretical Framework: A proof that transformation inversion can be used for test-time equivariance of pretrained models. The authors demonstrate that the transformation inversion posterior is $G$-equivariant, ensuring the soundness of the approach.
3. New Score Identity: A generalization of the target score identity to Lie groups that works without bi-invariant metrics, enabling efficient score-based sampling using energy gradients rather than clean data samples.
4. Empirical Validation: Demonstrations on image homographies ($PGL(3, \mathbb{R})$) and PDE symmetries ($SL(2, \mathbb{R}) \ltimes H(1, \mathbb{R})$), showing consistent performance improvements over strong baselines.

Experimental Results

The paper evaluates TIED on three main tasks:

1. Synthetic Sampling on $SO(10)$: TIED successfully samples from a multimodal energy distribution on a high-dimensional Lie group ($dim=45$), converging significantly faster than a trivialized kinetic Langevin sampler.
2. Image Classification (MNIST):
   • Setup: A pretrained ResNet18 is tested on MNIST images subjected to unknown affine ($Aff(2, \mathbb{R})$) and homography ($PGL(3, \mathbb{R})$) transformations.
   • Baselines: Compared against optimization-based methods (LieLAC, FoCal, ITS) and sampling methods (Kinetic Langevin).
   • Results: TIED significantly outperforms all baselines. For example, under affine transformations with classifier-based energy, TIED achieves 82.64% accuracy compared to 55.48% for the raw ResNet18 and 74.72% for Kinetic Langevin. It also outperforms specialized equivariant networks (affConv) in some settings.
3. PDE Solving:
   • Setup: A DeepONet neural operator is tested on 1D Heat and Burgers equations under unknown Lie point symmetry transformations.
   • Results: TIED achieves the lowest relative $L_2$ error across all test cases, outperforming LieLAC and FoCal. For the 1D Heat equation with $SL(2, \mathbb{R}) \ltimes H(1, \mathbb{R})$ symmetries, TIED reduces error to 0.043, compared to 0.066 for LieLAC and 0.690 for the uncorrected DeepONet.

Significance and Claims

The authors claim that TIED provides a training-free pipeline to make any pretrained model equivariant to Lie group transformations at test time. By leveraging energy-based models (which can be approximated by the pretrained model itself), the method avoids the need for specialized equivariant architectures or expensive retraining.

The significance lies in:

• Generality: It applies to a broad class of Lie groups, including those that are non-compact and lack bi-invariant metrics, which have been challenging for previous diffusion or optimization methods.
• Robustness: It effectively restores transformed inputs to the training distribution, significantly improving the generalization of pretrained networks to out-of-distribution geometric transformations.
• Efficiency: The diffusion-based approach, combined with the trivialized score identity, offers a more efficient and stable sampling process compared to traditional MCMC or iterative optimization on rugged energy landscapes.

The paper concludes that this approach represents a form of inference-time scaling for equivariance, offering a practical solution for robustness in the era of large foundation models where handling out-of-distribution transformations is critical.