Inverting Data Transformations via Diffusion Sampling

TL;DR

This paper introduces Transformation-Inverting Energy Diffusion (TIED), a diffusion-based sampler that operates on Lie groups to invert unknown data transformations given a data-space energy prior. By using a forward diffusion on the group and a reverse-time SDE with a trivialized Lie-algebra score, TIED stays on the manifold and leverages energy gradients to estimate the required scores. The method enables test-time equivariance for pretrained models and shows strong improvements on image homographies and Lie-point symmetry PDEs compared to baselines, all without training. This approach meaningfuly broadens probabilistic inversion to general transformation groups and offers a practical, training-free path to robust predictions in real-world settings.

Abstract

We study the problem of transformation inversion on general Lie groups: a datum is transformed by an unknown group element, and the goal is to recover an inverse transformation that maps it back to the original data distribution. Such unknown transformations arise widely in machine learning and scientific modeling, where they can significantly distort observations. We take a probabilistic view and model the posterior over transformations as a Boltzmann distribution defined by an energy function on data space. To sample from this posterior, we introduce a diffusion process on Lie groups that keeps all updates on-manifold and only requires computations in the associated Lie algebra. Our method, Transformation-Inverting Energy Diffusion (TIED), relies on a new trivialized target-score identity that enables efficient score-based sampling of the transformation posterior. As a key application, we focus on test-time equivariance, where the objective is to improve the robustness of pretrained neural networks to input transformations. Experiments on image homographies and PDE symmetries demonstrate that TIED can restore transformed inputs to the training distribution at test time, showing improved performance over strong canonicalization and sampling baselines. Code is available at https://github.com/jw9730/tied.

Figures (4)

• Figure 1: Graphical model describing our problem and method (with observed variables in gray and unobserved variables in white). $\mathcal{X}$ denotes the data space and $G$ a group of transformations, here the group of image homographies ${\rm PGL}(3,\mathbb{R})$. We are provided a data sample $\tilde{\mathbf{x}}$ that is generated by transforming an unknown in-distribution sample $\mathbf{x}$ with an unknown transformation $g$. We wish to sample from the posterior over transformations. Inspired by diffusion models, we construct a fast sampler that reverts a diffusion process on the Lie group starting from a random group element. The scores of the SDE are computed in the Lie algebra $\mathfrak{g}$.
• Figure 2: Energy (top) and density (bottom) along the forward process \ref{['eq:trivialized_forward_sde']} for the group of rotations $G={\rm SO}(2)$. The energy of the prior $E_0\left(g \right)\equiv E_{\mathbf{x}}\left(g^{-1}\cdot \tilde{\mathbf{x}} \right)$ is defined using the logsumexp of classifier logits from a ResNet18 trained on MNIST. The energy at small timesteps (top left) is low for likely orientations of the MNIST digit $\tilde{\mathbf{x}}$. The posterior $p_0(g)$ (bottom left) has modes centered around the most likely transformations Since the energy is obtained from a neural network, its landscape is highly rugged, resulting in exploding and vanishing scores. Going to the right, the densities along the forward process are plotted. The landscapes are increasingly smooth and address the ruggedness.
• Figure 3: Sampling on ${\rm SO}(10)$ under energy $E:\mathbf{X}\mapsto -10 \mathbf{X}_{1,1}^2$ using a kinetic Langevin sampler kong2024convergence and TIED (Ours). The distribution of $\mathbf{X}_{1,1}$ induced by the energy has two symmetric modes around zero, and thus the mean of $\mathbf{X}_{1,1}$ approaches zero as the sampling converges.
• Figure 4: For each PDE, we show an out-of-domain case for DeepONet $f_\theta$. From left: true solution, $f_\theta$ prediction, $f_\theta$ prediction under test-time equivariance via TIED (Ours). Zoom in for a better view.

Theorems & Definitions (35)

• Proposition 3.1: Transformation inversion posterior
• Proposition 3.2: Equivariance of the posterior
• Proposition 4.1: Reverse trivialized SDE
• Proposition 4.2: Trivialized target score identity
• Proposition 4.3: Monte Carlo score estimator
• Proposition 1.1: Trivialized gradient
• proof
• Proposition 1.2: Trivialized score
• proof
• Proposition \ref{prop:transformation_posterior}: Transformation inversion posterior