Inverting Data Transformations via Diffusion Sampling

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Kim, Jinwoo, Kaba, Sékou-Oumar, Park, Jiyun, Hong, Seunghoon, Ravanbakhsh, Siamak
Format: Preprint
Publié: 2026
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866913171687604224
author Kim, Jinwoo
Kaba, Sékou-Oumar
Park, Jiyun
Hong, Seunghoon
Ravanbakhsh, Siamak
author_facet Kim, Jinwoo
Kaba, Sékou-Oumar
Park, Jiyun
Hong, Seunghoon
Ravanbakhsh, Siamak
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2602_08267
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Inverting Data Transformations via Diffusion Sampling
Kim, Jinwoo
Kaba, Sékou-Oumar
Park, Jiyun
Hong, Seunghoon
Ravanbakhsh, Siamak
Machine Learning
Artificial Intelligence
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.
title Inverting Data Transformations via Diffusion Sampling
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2602.08267