Inverting Data Transformations via Diffusion Sampling
This addresses the issue of distorted observations in machine learning and scientific modeling, offering a method to enhance robustness to input transformations, though it appears incremental as it builds on existing diffusion and equivariance techniques.
The paper tackles the problem of inverting unknown data transformations on Lie groups to recover the original data distribution, introducing TIED, a diffusion-based method that improves test-time equivariance for pretrained neural networks, showing better performance than baselines on image homographies and PDE symmetries.
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.