Title: Inverting Data Transformations via Diffusion Sampling
Primary area: learning on graphs and other geometries & topologies
Keywords: transformation, symmetry, diffusion, sampling, equivariance, test-time, pde, homography, deep learning

Abstract: We study transformation inversion in general Lie groups, where an input has been transformed by an unknown element and the goal is to recover an inverse that maps it back to the data distribution. This problem is fundamental across machine learning and scientific modeling, where unknown transformations often distort observations. As an application, we focus on test-time equivariance, which improves the robustness of neural networks to input transformations. We adopt a probabilistic approach and model the posterior over transformations as a Boltzmann distribution parametrized by a data-space energy.  To sample from this distribution, we introduce a scheme for diffusion on Lie groups that keeps updates on-manifold and requires only Lie algebra computations.  Our method, Transformation-Inverting Energy Diffusion (TIED), uses a new trivialized target-score identity and can efficiently sample from the Boltzmann posterior on the group. TIED addresses curved group geometry and rough, multimodal energy landscapes, which are limitations of existing methods. It also applies to a wide class of groups and in contrast to prior works does not assume compactness or a bi-invariant metric. Experiments on image homographies and PDE symmetries show test-time equivariance of pretrained networks with improved performance over strong canonicalization baselines.

TL;DR: We propose method to invert transformations of data by sampling from a Boltzmann distribution on a Lie group.