Inverting Data Transformations via Diffusion Sampling

2026-02-09

AI Summary

This work proposes TIED, a method for solving inverse problems involving unknown data transformations on general Lie groups. By modeling the transformation posterior as a Boltzmann distribution defined via an energy function and constructing a diffusion process in the Lie algebra to preserve manifold structure, TIED enables efficient posterior sampling. Its key innovation is the introduction of a trivialized objective score identity, which— for the first time—enables efficient score-based posterior inference on Lie groups. The approach supports test-time equivariance, substantially enhancing model robustness. In tasks involving image homography and PDE symmetries, TIED successfully recovers transformed inputs back to the original training distribution, outperforming existing normalization and sampling baselines.

Technology Category

• Machine Learning: Learning with Manifolds
• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• transformation inversion
• Lie groups
• data distribution
• test-time equivariance
• unknown transformations

Innovation

• diffusion sampling
• Lie groups
• transformation inversion
• score-based generative modeling
• test-time equivariance