# Inverting Data Transformations via Diffusion Sampling

arXiv: https://arxiv.org/abs/2602.08267 · HTML: https://mlsanity.com/paper/2602.08267/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 51 of its 87 references that are in the map.
- First posted on arXiv 9 Feb 2026.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 30 of 51)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Going Deeper with Convolutions](https://mlsanity.com/paper/1409.4842/) (2014)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Benchmarking Neural Network Robustness to Common Corruptions and Perturbations](https://mlsanity.com/paper/1903.12261/) (2019)
- [DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators](https://mlsanity.com/paper/1910.03193/) (2019)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [Spatial Transformer Networks](https://mlsanity.com/paper/1506.02025/) (2015)
- [Diffusion Posterior Sampling for General Noisy Inverse Problems](https://mlsanity.com/paper/2209.14687/) (2022)
- [Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges](https://mlsanity.com/paper/2104.13478/) (2021)
- [Group Equivariant Convolutional Networks](https://mlsanity.com/paper/1602.07576/) (2016)
- [Your Classifier is Secretly an Energy Based Model and You Should Treat it Like One](https://mlsanity.com/paper/1912.03263/) (2019)
- [On the Number of Linear Regions of Deep Neural Networks](https://mlsanity.com/paper/1402.1869/) (2014)
- [On the Generalization of Equivariance and Convolution in Neural Networks to the Action of Compact Groups](https://mlsanity.com/paper/1802.03690/) (2018)
- [Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks](https://mlsanity.com/paper/1807.07356/) (2018)
- [Cold Diffusion: Inverting Arbitrary Image Transforms Without Noise](https://mlsanity.com/paper/2208.09392/) (2022)
- [Riemannian Score-Based Generative Modelling](https://mlsanity.com/paper/2202.02763/) (2022)
- [A Practical Method for Constructing Equivariant Multilayer Perceptrons for Arbitrary Matrix Groups](https://mlsanity.com/paper/2104.09459/) (2021)
- [Frame Averaging for Invariant and Equivariant Network Design](https://mlsanity.com/paper/2110.03336/) (2021)
- [Denoising Diffusion Samplers](https://mlsanity.com/paper/2302.13834/) (2023)
- [Equivariance Through Parameter-Sharing](https://mlsanity.com/paper/1702.08389/) (2017)
- [Universal approximations of invariant maps by neural networks](https://mlsanity.com/paper/1804.10306/) (2018)
- [Better Aggregation in Test-Time Augmentation](https://mlsanity.com/paper/2011.11156/) (2020)
- [Riemannian Diffusion Models](https://mlsanity.com/paper/2208.07949/) (2022)
- [Iterated Denoising Energy Matching for Sampling from Boltzmann Densities](https://mlsanity.com/paper/2402.06121/) (2024)
- [Improved sampling via learned diffusions](https://mlsanity.com/paper/2307.01198/) (2023)
- [Generator Matching: Generative modeling with arbitrary Markov processes](https://mlsanity.com/paper/2410.20587/) (2024)
- [Polar Transformer Networks](https://mlsanity.com/paper/1709.01889/) (2017)

## Built on by (most cited first, 2 of 2)

- [What are the Right Symmetries for Formal Theorem Proving?](https://mlsanity.com/paper/2605.22257/) (2026)
- [Zero-Shot Test-Time Canonicalization using Out-of-Distribution Scoring](https://mlsanity.com/paper/2606.24178/) (2026)

## Abstract (arXiv)

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 the 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.