mlsanityData to 7 Oct 2026

Inverting Data Transformations via Diffusion Sampling

arXiv 2602.08267 ↗First posted 9 Feb 2026cs.LG cs.AIcited by 2

On the map: Diffusion models → Diffusion foundations and sampling

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.

Key facts. Cited by 2 papers in the map (data through 7 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.

Builds on

51 of 87 references.

Closest ancestors · 7

Same direction, newest first

  1. 2024Target Score Matching13 Feb
  2. 2023Consistency Models2 Mar
  3. 2023Denoising Diffusion Samplers27 Feb
  4. 2022Cold Diffusion: Inverting Arbitrary Image Transforms Without Noise19 Aug
  5. 2022Riemannian Diffusion Models16 Aug
  6. 2022Riemannian Score-Based Generative Modelling6 Feb
  7. 2020Score-Based Generative Modeling through Stochastic Differential Equations26 Nov

Other references · 44

  1. 2015Deep Residual Learning for Image Recognitioncited by 51,809
  2. 2015Deep Unsupervised Learning using Nonequilibrium Thermodynamicscited by 7,005
  3. 2014Going Deeper with Convolutionscited by 4,231
  4. 2019Generative Modeling by Estimating Gradients of the Data Distributioncited by 4,067
  5. 2019Benchmarking Neural Network Robustness to Common Corruptions and Perturbationscited by 2,857
  6. 2020Fourier Neural Operator for Parametric Partial Differential Equationscited by 2,806
  7. 2019DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operatorscited by 1,965
  8. 2015Spatial Transformer Networkscited by 1,423
  9. 2022Diffusion Posterior Sampling for General Noisy Inverse Problemscited by 1,375
  10. 2016Group Equivariant Convolutional Networkscited by 1,154
  11. 2021Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gaugescited by 1,131
  12. 2019Your Classifier is Secretly an Energy Based Model and You Should Treat it Like Onecited by 402
32 more
  1. 2014On the Number of Linear Regions of Deep Neural Networkscited by 378
  2. 2018On the Generalization of Equivariance and Convolution in Neural Networks to the Action of Compact Groupscited by 296
  3. 2018Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networkscited by 278
  4. 2021A Practical Method for Constructing Equivariant Multilayer Perceptrons for Arbitrary Matrix Groupscited by 185
  5. 2021Frame Averaging for Invariant and Equivariant Network Designcited by 166
  6. 2017Equivariance Through Parameter-Sharingcited by 139
  7. 2018Universal approximations of invariant maps by neural networkscited by 127
  8. 2020Better Aggregation in Test-Time Augmentationcited by 121
  9. 2024Iterated Denoising Energy Matching for Sampling from Boltzmann Densitiescited by 113
  10. 2023Improved sampling via learned diffusionscited by 103
  11. 2024Generator Matching: Generative modeling with arbitrary Markov processescited by 82
  12. 2017Polar Transformer Networkscited by 81
  13. 2019Trivializations for Gradient-Based Optimization on Manifoldscited by 78
  14. 2020Solving Linear Inverse Problems Using the Prior Implicit in a Denoisercited by 78
  15. 2022Lie Point Symmetry Data Augmentation for Neural PDE Solverscited by 76
  16. 2020Learning Loss for Test-Time Augmentationcited by 56
  17. 2024Equivariant Frames and the Impossibility of Continuous Canonicalizationcited by 41
  18. 2023Equivariant Adaptation of Large Pretrained Modelscited by 40
  19. 2021Equivariant neural networks for inverse problemscited by 29
  20. 2023Learning Probabilistic Symmetrization for Architecture Agnostic Equivariancecited by 25
  21. 2023Lie Point Symmetry and Physics Informed Networkscited by 24
  22. 2024Lie Algebra Canonicalization: Equivariant Neural Operators under arbitrary Lie Groupscited by 23
  23. 2021Enabling equivariance for arbitrary Lie groupscited by 22
  24. 2024Trivialized Momentum Facilitates Diffusion Generative Modeling on Lie Groupscited by 19
  25. 2025Improving Equivariant Networks with Probabilistic Symmetry Breakingcited by 15
  26. 2020Variational Optimization on Lie Groups, with Examples of Leading (Generalized) Eigenvalue Problemscited by 12
  27. 2024A Generative Model of Symmetry Transformationscited by 10
  28. 2024A comparison between humans and AI at recognizing objects in unusual posescited by 7
  29. 2024Tilt your Head: Activating the Hidden Spatial-Invariance of Classifierscited by 7
  30. 2024Convergence of Kinetic Langevin Monte Carlo on Lie groupscited by 6
  31. 2024Stochastic Neural Network Symmetrisation in Markov Categoriescited by 5
  32. 2025Test-Time Canonicalization by Foundation Models for Robust Perceptioncited by 5

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2026: 2.

2 papers.

  1. 2026What are the Right Symmetries for Formal Theorem Proving?cited by 0
  2. 2026Zero-Shot Test-Time Canonicalization using Out-of-Distribution Scoringcited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2025Diffusion Generative Modeling on Lie Group Representations
  2. 2025Generative diffusion model with inverse renormalization group flows
  3. 2026Diffusion Model's Generalization Can Be Characterized by Inductive Biases toward a Data-Dependent Ridge Manifold
  4. 2024DEFT: Efficient Fine-Tuning of Diffusion Models by Learning the Generalised h-transform
  5. 2026IMS3: Breaking Distributional Aggregation in Diffusion-Based Dataset Distillation
  6. 2023U-Turn Diffusion
  7. 2023Neural Diffusion Models
  8. 2023Reflected Diffusion Models