Inverting Data Transformations via Diffusion Sampling
Jinwoo Kim, S\'ekou-Oumar Kaba, Jiyun Park, Seunghoon Hong, Siamak Ravanbakhsh
TL;DR
This paper introduces TIED, a diffusion-based method for inverting unknown transformations on data modeled on Lie groups, improving robustness of neural networks to input distortions.
Contribution
We develop a diffusion process on Lie groups with a new score identity, enabling efficient probabilistic inversion of transformations for improved model robustness.
Findings
TIED effectively restores transformed inputs to original data distribution.
It outperforms canonicalization and sampling baselines in experiments.
Demonstrates robustness improvements on image homographies and PDE symmetries.
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…
Peer Reviews
Decision·Submitted to ICLR 2026
This paper presents a very interesting and working algorithm. The problem of group-equivariant learning has been studied for a long time. As an important property of physics (symmetry), the learning/prediction outcomes should be equivariant to group transformations. However, most of the prior methods (i) data augmentation (ii) learning a standard position or (iii) group convolutions can be (i) not robust (ii) hard to define (iii) expensive. This framework nicely combines the strength of the di
The main weakness of this paper is the presentation. Some figures and experimental results are hard to understand easily (even after a couple of re-reads from the reviewer's side - who is familiar with the field). Please see the questions below.
- Overall, I found the setup both novel and conceptually interesting. The method enables density estimation on Lie groups using purely data-driven constraints, bridging geometry and probabilistic modeling elegantly. - The experiments show that applying TIED before inference improves neural network accuracy and consistency under unseen transformations, providing a convincing proof of concept for the proposed idea.
- Despite its technical novelty, the paper suffers from weak presentation and clarity. The writing is often dense, and key visual aids (e.g., Figure 1) are underexplained—either the caption needs more detail or the figure itself should be redesigned to convey the main idea clearly. - Conceptually, the work frames “canonicalization” as a way to achieve equivariance through energy-based sampling on Lie groups. However, the paper lacks a clear discussion of how this approach fits within the broader
1. The work proposes a novel framework for inverting unknown data transformations from general Lie groups, which breaks through the limitations of traditional diffusion models and provides a new research perspective for the field of transformation inversion. 2. Compared with traditional diffusion models, TIED infers the posterior distribution of transformations based on an energy prior, which significantly accelerates the reconstruction process. Experimental results further confirm that TIED ac
1. The experimental setup lacks sufficient clarity. In Section 5.2, the specific parameters of the affine matrix (a critical factor directly affecting experimental results) are not provided. This omission makes it difficult to fully verify the reliability of the experimental conclusions and hinders the reproducibility of the study. 2. In existing works on affine reconstruction, experimental results are typically presented as quantitative data distributed across the real number axis to reflect r
Videos
No videos yet. Explain this paper in a talk, walkthrough, or lecture? Add one.
Taxonomy
TopicsGenerative Adversarial Networks and Image Synthesis · Domain Adaptation and Few-Shot Learning · Morphological variations and asymmetry