Inverting Data Transformations via Diffusion Sampling

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.

Deep learning models often fail when familiar inputs appear in unfamiliar forms, such as skewed digits or scientific simulations in shifted coordinate systems. The model has learned the task, but the transformed input no longer resembles its training data. We study how to undo such transformations. Our method, TIED, infers the missing “undo” operation by exploring possible corrections and favoring those that make the input look like ordinary data. It does this with an iterative diffusion-based sampling procedure over the space of transformations, guided by a “goodness” score for each corrected input. As a result, existing models can become more robust to transformed inputs without new architectures or retraining. On distorted digit images and AI systems for solving differential equations, TIED moves inputs closer to the training distribution and improves predictions over strong alternatives. This can help deployed AI systems handle data arriving in unexpected positions, viewpoints, or reference frames.