Pith. sign in

REVIEW 2 major objections 5 minor 1 cited by

Diffusion sampling on the Lie group, driven by a new trivialized target-score identity, can invert unknown data transformations and bring distorted inputs back to the training distribution.

Reviewed by Pith at T0; open to challenge. T0 means a machine referee read the full paper against a public rubric. the ladder, T0–T4 →

A Lie-group diffusion sampler that inverts unknown data transformations at test time, using only an energy function, and improves pretrained models on affine/homography images and PDE solving.

T0 review reviewed 2026-08-03 challenge →

load-bearing objection Solid theory and strong empirical results on test-time equivariance, but the practical sampler's finite-N score estimate needs direct validation before the posterior-sampling claim is fully earned. the 2 major comments →

arxiv 2602.08267 v2 pith:ELIXE252 submitted 2026-02-09 cs.LG cs.AI

Inverting Data Transformations via Diffusion Sampling

classification cs.LG cs.AI
keywords transformation inversionLie groupsdiffusion samplingenergy-based modelstest-time equivariancecanonicalizationLie algebraLie point symmetries
verification ladder T0 review T1 audit T2 compute T3 formal T4 reserved

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

This paper tries to solve a blind inverse problem: a datum is changed by an unknown element of a Lie group, and the goal is to recover the inverse transformation. The authors model the posterior over group elements as a Boltzmann distribution defined by an energy on data space, then sample from that posterior using a diffusion process that stays on the group and works in the Lie algebra. Their key theoretical move is a new trivialized target-score identity, which lets the noisy score be estimated directly from clean energy gradients without needing clean samples or a bi-invariant metric. If the approach is correct, any pretrained neural network can be made robust to geometric transformations at test time simply by undoing the unknown transform before prediction.

Core claim

The paper claims that the transformation-inversion posterior on any connected Lie group, p(g|x̃) ∝ exp(−E(g⁻¹·x̃))|det J(g⁻¹)|, is G-equivariant, so sampling h from it and applying h⁻¹·x̃ maps the transformed input back to the training distribution. The central discovery is a 'trivialized target-score identity': for the forward SDE dg_t = d(L_{g_t})_e[γ(t) dw_t], the trivialized score ∇_g log p_t(g_t) equals an average of initial trivialized scores over the conditional path, and therefore can be written as a Monte Carlo expression involving only the energy, the noise density, and a modular-function correction. Because all gradients are expressed in the Lie algebra, this score estimate works

What carries the argument

The central machinery is the reverse trivialized SDE (Prop. 4.1) together with the trivialized target-score identity (Prop. 4.2) and its Monte Carlo form (Prop. 4.3). Trivialization means pulling gradients back to the Lie algebra via d(L_{g⁻¹})_g, so gradients evaluated at different group elements live in one common vector space and can be averaged without parallel transport. The identity expresses the noisy score as ∇_g log ∫ k_t(w) exp(−E(g w⁻¹) − log λ(w)) dµ(w), with λ the modular function accounting for non-unimodularity; this is what makes the sampler trainable from an arbitrary energy function. The reverse SDE is discretized with exponential-map updates, and the whole pipeline needs o

Load-bearing premise

The load-bearing premise is that averaging a handful of random noise draws per diffusion step gives a good enough estimate of the true noisy score; the proof only shows consistency in the limit of infinitely many draws.

What would settle it

Take images with known ground-truth affine or homography transformations, run TIED with the same energy used in the paper, and compare the sampled inverse h against the true transformation g. If, at N=2–10 Monte Carlo samples, the recovered transformations do not concentrate near the true modes of p(g|x̃), or corrected images do not have lower energy than the uncorrected ones, then the sampler is not approximating the reverse SDE and the in-distribution guarantee collapses.

Watch this falsifier. Get emailed when new claim-graph text bears on it.

If this is right

  • If TIED is correct, test-time equivariance no longer requires an equivariant architecture or retraining: any pretrained model can be made robust to Lie-group transformations using an energy function, including the model's own confidence.
  • Corrected inputs are guaranteed to be drawn from the data prior (Prop. 3.1), so downstream inference operates on in-distribution samples rather than on arbitrary transformed inputs.
  • The method covers noncompact, non-Abelian groups and nonlinear actions, needing only the Lie algebra and exponential map; this includes projective homographies and PDE symmetry groups such as SL(2,R)⋉H(1,R).
  • On the reported benchmarks, TIED restores accuracy and lowers FID compared with optimization, Langevin, and canonicalization baselines, and on the tested settings matches or exceeds specialized equivariant networks.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • Because only on-orbit energy differences matter, the method may tolerate fairly crude energies; a testable extension is to use an energy from a model trained on a different dataset and see whether inversion still pushes inputs in-distribution.
  • The finite-N Monte Carlo score estimate is a standing risk; a natural diagnostic is to measure the gap between TIED samples and the true Boltzmann posterior modes on a rugged energy, then adaptively increase N where the gap is large.
  • The trivialized identity suggests a route to sampling on other manifolds that admit a trivializing chart, and to faster samplers via distilling the reverse SDE into fewer steps, directions the paper hints at but does not pursue.
  • If the energy is provided by a learned generative model, TIED effectively turns blind inverse problems such as image registration or sensor-motion estimation into a single on-group sampling procedure, though that application is not demonstrated.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit.

Referee Report

2 major / 5 minor

Summary. The paper studies transformation inversion: an observed datum x̃ is generated by applying an unknown element g of a Lie group G to an in-distribution sample x. The authors model the posterior p(g|x̃) as a Boltzmann distribution from an energy E(g) = E_x(g^{-1}·x̃) − log|det J_{g^{-1}}(x̃)| and propose Transformation-Inverting Energy Diffusion (TIED), a diffusion sampler on the group. The sampler uses a new 'trivialized target-score identity' (Prop. 4.2) that expresses the score of the noised posterior as an average of clean energy gradients, and a Monte Carlo estimator (Prop. 4.3, Eq. 23) suitable for implementation. Algorithm 1 discretizes the reverse trivialized SDE using exponential-map updates. The authors prove that the posterior is equivariant (Prop. 3.2), which they use for test-time equivariance of pretrained networks. Experiments cover SO(10) sampling, MNIST with affine and perspective transformations, and PDE solving under Lie point symmetries, reporting improvements over optimization and sampling baselines.

Significance. If the practical sampler is sound, the paper makes a valuable contribution: a diffusion-based posterior sampler on general Lie groups that does not require compactness, a bi-invariant metric, or linear actions, and that enables training-free test-time equivariance for generic pretrained models. The theoretical development is largely self-contained, and the SO(10) experiment shows a substantial speedup over kinetic Langevin in a smooth, low-dimensional posterior regime. The paper includes code and compares across several energy choices and baselines. The main caveat is that the deployed algorithm uses a finite-N Monte Carlo score estimator whose bias is not controlled, so the theoretical guarantees do not automatically transfer to the experimental settings.

major comments (2)
  1. [§4.3, Eq. (23); Prop. B.12; Table 3] The practical estimator replaces the integral in Proposition 4.3 with a logsumexp over N Monte Carlo draws from k_t (Eq. 23). Proposition B.12 establishes only almost-sure consistency as N→∞; no non-asymptotic bound is given. Table 3 uses N=2–10 in all practical experiments, while the SO(10) demonstration uses N=100 with a smooth quadratic energy. Because the drift in the reverse SDE (Eq. 4) is exactly this estimated score, a finite-N bias means the sampler output is not distributed according to p(g|x̃). Consequently the in-distribution guarantee of Prop. 3.1 and the equivariance guarantee of Prop. 3.2 do not apply to the deployed algorithm. The paper needs either a non-asymptotic error bound or a direct empirical check that the sampler produces samples from the posterior in the practical N regime.
  2. [§5, Tables 1–2, Fig. 3] The evaluation reports classification accuracy and FID, but not whether TIED actually samples the posterior. Accuracy can improve even if the sampler degenerates into a deterministic optimizer, and FID measures closeness to the training image distribution, not agreement with p(g|x̃). The synthetic SO(10) experiment (Fig. 3) does validate posterior sampling in a smooth-energy, N=100 regime, but the image and PDE experiments use N=2–10 with rugged neural-network or boundary-distance energies and do not include any direct sampling-quality metric. The central claim that TIED 'samples from the posterior' and that the theoretical guarantees therefore hold at test time is thus not directly supported by the experiments. I recommend adding a direct measure of sampler quality for the practical N regime — for example, comparing TIED output histograms to ground-truth posterior densities on a small g
minor comments (5)
  1. [§4.2, Eq. (5)] The notation ∇_g log p_0(g_0) in the trivialized target-score identity is potentially misleading. The proof clarifies that the argument is g_t b with b ≡ g_t^{-1} g_0, and that the derivative is taken with respect to g_t. I recommend writing this as ∇_g log p_0(g_t b)|_{b = g_t^{-1} g_0} or using an explicit coordinate form to avoid confusion.
  2. [Tables 1 and 4] The TIED accuracy for the Aff(2,R) classifier-energy row is 85.64% in Table 1 but 82.64±0.11% in Table 4 (five repeated tests). Please clarify which number is final and, if Table 1 is a single run, present the mean and standard deviation consistently across the tables.
  3. [§3.2, Prop. 3.1] The free action assumption is stated but not discussed. For data with nontrivial stabilizers, the posterior density on the group is not defined in the same way. A brief remark on how the method behaves or fails when the action is not free would help.
  4. [§B.9, Algorithm 1, line 2] The initialization uses the approximation k_1 ≈ p_1 without empirical validation. Since the noise schedule is finite in all experiments, a sentence justifying this approximation (or a sensitivity check on γ_max) would strengthen the paper.
  5. [§5.2] FID is reported for MNIST using an inception-based feature space (Fatir, 2018). For 40×40 grayscale digits this may not be a meaningful metric; a simple distributional metric on pixels or on the classifier's feature space would be more interpretable.

Circularity Check

0 steps flagged

No significant circularity; the score identity is derived from the stated forward SDE, the energy is an input, and validation includes an independent FID check.

full rationale

The derivation chain is self-contained rather than circular. Proposition 3.1 derives the transformation-inversion posterior directly from the stated generative model and a change of variables; Proposition 3.2 verifies equivariance by direct computation. Propositions 4.1-4.3 and their proofs in Appendix B derive the reverse trivialized SDE, the trivialized target-score identity, and the Monte Carlo score estimator from the forward SDE (3), without importing the result from a citation. The energy E is an input to the algorithm, not a fitted parameter, and no parameter is fit to a subset of the reported predictions. The image experiments additionally report FID between inverse-transformed images and training images, which is an independent check of the in-distribution claim rather than a re-statement of the accuracy objective. Self-citations (Kaba et al. 2023, Kim et al. 2023, Lawrence et al. 2025, Akhound-Sadegh et al. 2024) appear in standard contexts - canonicalization, ensembling, and the Euclidean precursor of target-score identity - but the load-bearing Lie-group identity is proved in this paper and does not reduce to those citations. The main flagged weakness is statistical, not circular: Section B.8 states that the MC estimator is biased because of the log of a sample mean and proves only consistency as N goes to infinity (Prop. B.12), while the experiments use N=2-10 (Table 3). This is a finite-sample correctness/robustness concern about whether the reverse SDE actually samples the posterior, not a case where a prediction is equivalent to its input by construction.

Axiom & Free-Parameter Ledger

3 free parameters · 5 axioms · 0 invented entities

The method introduces no new physical entities. Its central load-bearing choices are the hand-tuned diffusion hyperparameters and several analytic assumptions about Lie-group structure, smoothness, and discretization accuracy that are plausible but not fully verified.

free parameters (3)
  • Noise schedule γmin, γmax = per setting: SO(10) (0.01,10); Aff(2,R) (0.1,1); PGL(3,R) (0.01/0.05,0.5); PDE (0.01,0.5)
    Tab.3; chosen by hand following VE diffusion (Song et al. 2020), not derived.
  • Step size Δt = 1/50 (images/PDE), 1/100 (SO(10))
    Tab.3; chosen by hand.
  • MC sample size N = 2-100 (100 for SO(10), 2-4 images, 10 PDE)
    Tab.3; trades cost/quality.
axioms (5)
  • domain assumption G is a connected Lie group and the action is free and diffeomorphic
    Sec. 3.1-3.2, Prop. 3.1; experiments use Aff(2,R) and PGL(3,R), which are disconnected, and no freedom check for actual images.
  • domain assumption The noisy densities pt are smooth and strictly positive, and energy normalization is finite for noncompact groups
    Sec. 4.1 Prop 4.1, "Under mild conditions on the energy"; no explicit conditions verified for neural-network energies.
  • ad hoc to paper k1 approximates p1 so the reverse process can be initialized from k1
    Sec. B.9, "we approximate k1≈p1 as usually done for variance-exploding diffusion"; no bound on the approximation error.
  • ad hoc to paper Euler discretization (line 12, Eq. 24) converges to the continuous reverse and forward SDEs
    Sec. B.9; no discretization error analysis for noncompact groups.
  • standard math Monte Carlo estimator integrability/positivity conditions
    Prop B.12 assumes integrability and positivity; these are not verified for the energies used.

reviewed 2026-08-03 · how reviews work

0 comments
Cite this review

Pith. "Pith review of Inverting Data Transformations via Diffusion Sampling." pith.science (2026). https://pith.science/paper/ELIXE252

@misc{pith2026260208267,
  author       = {Pith},
  title        = {Pith review of: Inverting Data Transformations via Diffusion Sampling},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/ELIXE252}},
  note         = {Machine review of arXiv:2602.08267}
}
Share X Bluesky LinkedIn Reddit HN
read the original 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 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.

Figures

Figures reproduced from arXiv: 2602.08267 by Jinwoo Kim, Jiyun Park, S\'ekou-Oumar Kaba, Seunghoon Hong, Siamak Ravanbakhsh.

Figure 1
Figure 1. Figure 1: Graphical model describing our problem and method (with observed variables in gray and unobserved variables in white). X denotes the data space and G a group of transformations, here the group of image homographies PGL(3, R). We are provided a data sample x˜ that is generated by transforming an unknown in-distribution sample x with an unknown transformation g. We wish to sample from the posterior over tran… view at source ↗
Figure 2
Figure 2. Figure 2: Energy (top) and density (bottom) along the forward process (3) for the group of rotations G = SO(2). The energy of the prior E0 (g) ≡ Ex [PITH_FULL_IMAGE:figures/full_fig_p005_2.png] view at source ↗
Figure 3
Figure 3. Figure 3: Sampling on SO(10) under energy E : X 7→ −10X2 1,1 using a kinetic Langevin sampler (Kong & Tao, 2024) and TIED (Ours). The distribution of X1,1 induced by the energy has two sym￾metric modes around zero, and thus the mean of X1,1 approaches zero as the sampling converges. problems using a trained convolutional neural network un￾der unknown affine and homography (perspective) transfor￾mations; and (iii) tw… view at source ↗
Figure 4
Figure 4. Figure 4: For each PDE, we show an out-of-domain case for DeepONet fθ. From left: true solution, fθ prediction, fθ prediction under test-time equivariance via TIED (Ours). Zoom in for a better view. challenging transformation groups. The first is the group of affine transformations Aff(2, R), and the second is the group of homography (perspective) transformations, isomor￾phic to the projective general linear group P… view at source ↗

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Forward citations

Cited by 1 Pith paper

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score.

  1. What are the Right Symmetries for Formal Theorem Proving?

    cs.LG 2026-05 unverdicted novelty 7.0

    Introduces rewriting categories to formalize proof equivariance and success invariance, shows LLM provers violate both, and demonstrates test-time aggregation recovers invariance and boosts performance.

Reference graph

Works this paper leans on

19 extracted references · 6 linked inside Pith · cited by 1 Pith paper

  1. [1]

    However, on a Lie group, gradients of a function at different points live in different tangent spaces and hence are not directly compatible, e.g., cannot be added or averaged

    In energy-based diffusion, one usually employs an expectation form of the score function that averages energy gradients evaluated across the state space (De Bortoli et al., 2024; Akhound-Sadegh et al., 2024). However, on a Lie group, gradients of a function at different points live in different tangent spaces and hence are not directly compatible, e.g., c...

  2. [2]

    We can sidestep both difficulties by always working in the Lie algebra instead of arbitrary tangent spaces

    In general, during sampling on a Lie group, the update directions live in tangent spaces that change over time simultaneously as a sample is updated, requiring a careful handling. We can sidestep both difficulties by always working in the Lie algebra instead of arbitrary tangent spaces. This technique is called (left-)trivialization (Lezcano Casado, 2019;...

  3. [3]

    For the first item, we recall the relationship gt =g 0wt with g0 ∼p 0 and wt ∼k t independent (proof of Theorem 4.2), and recall the forward process (3)

    A method to compute trivialized gradient∇ g of generalf:G→R. For the first item, we recall the relationship gt =g 0wt with g0 ∼p 0 and wt ∼k t independent (proof of Theorem 4.2), and recall the forward process (3). Together, these imply that wt ∼k t is described by the following SDE which is identical to the forward SDE but starts at the identity: dwt = d...

  4. [6]

    (pages 3, 7, 8, 13) Krizhevsky, A., Sutskever, I., and Hinton, G. E. Imagenet classification with deep convolutional neural networks. Advances in neural information processing systems, 25,

  5. [8]

    J.Applications of Lie groups to differential equa- tions, volume 107

    (page 1) Olver, P. J.Applications of Lie groups to differential equa- tions, volume 107. Springer Science & Business Media,

  6. [15]

    The posterior distribution ofghas densityp(g| ˜x)∝p x g−1 · ˜x |detJ g−1 (˜x)|

  7. [16]

    The random variablex ′ =h −1 · ˜x, withh∼p(g| ˜x)has densityp(x ′ | ˜x)∝p x (x′)1 G·x (˜x)with respect toµ G·x. Proof.For the first statement, using independence ofxandgand the relation ˜x=g·x, we have p(x, ˜x|g) =p x (x)p( ˜x|x, g) p(x, ˜x|g) =p x (x)δ( ˜x−g·x) 14 Inverting Data Transformations via Diffusion Sampling Marginalizing with respect toxyields ...

  8. [17]

    A method to samplew∼k t,

  9. [18]

    A method to calculatelogλ(w),

  10. [1993]

    J., and Lipman, Y

    (pages 1, 9, 13) 11 Inverting Data Transformations via Diffusion Sampling Puny, O., Atzmon, M., Ben-Hamu, H., Misra, I., Grover, A., Smith, E. J., and Lipman, Y . Frame averaging for invariant and equivariant network design.arXiv preprint arXiv:2110.03336, 2021. (page 3) Ravanbakhsh, S., Schneider, J., and Poczos, B. Equivariance through parameter-sharing...

  11. [1996]

    and Stober, S

    (page 4) Schmidt, J. and Stober, S. Tilt your head: Activating the hidden spatial-invariance of classifiers.arXiv preprint arXiv:2405.03730, 2024. (pages 2, 3, 4, 8) Shanmugam, D., Blalock, D., Balakrishnan, G., and Guttag, J. When and why test-time augmentation works.arXiv preprint arXiv:2011.11156, 2020. (pages 3, 4) Shao, J.Mathematical Statistics. Spr...

  12. [2003]

    Background We provide an overview of the mathematical background, and refer the readers to Lee (2012); Tu (2010) for more details

    (pages 1, 2) 12 Inverting Data Transformations via Diffusion Sampling A. Background We provide an overview of the mathematical background, and refer the readers to Lee (2012); Tu (2010) for more details. Lie groupA Lie group G is a group that is also a smooth manifold, such that multiplications of elements and taking inverses are smooth. In deep learning,...

  13. [2011]

    Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks.Neurocomputing, 338: 34–45, 2019

    (page 3) Wang, G., Li, W., Aertsen, M., Deprest, J., Ourselin, S., and Vercauteren, T. Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks.Neurocomputing, 338: 34–45, 2019. (page 3) Welling, M. and Teh, Y . W. Bayesian learning via stochastic gradient langevin dynamics. InProceedin...

  14. [2012]

    Im- proving equivariant networks with probabilistic symmetry breaking

    (page 3) Lawrence, H., Portilheiro, V ., Zhang, Y ., and Kaba, S.-O. Im- proving equivariant networks with probabilistic symmetry breaking. InThe Thirteenth International Conference on Learning Representations, 2025. (pages 3, 4) Lee, J.Introduction to Smooth Manifolds. Graduate Texts in Mathematics. Springer New York, 2012. (pages 13, 14, 16) Lee, T. and...

  15. [2013]

    Deep residual learn- ing for image recognition

    (pages 1, 2) He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learn- ing for image recognition. InProceedings of the IEEE conference on computer vision and pattern recognition, pp. 770–778, 2016. (page 7) Hendrycks, D. and Dietterich, T. Benchmarking neural net- work robustness to common corruptions and perturbations. arXiv preprint arXiv:1903.12261,...

  16. [2019]

    W., and Doucet, A

    (page 2) De Bortoli, V ., Mathieu, E., Hutchinson, M., Thornton, J., Teh, Y . W., and Doucet, A. Riemannian score-based gen- erative modelling.Advances in Neural Information Pro- cessing Systems, 35:2406–2422, 2022. (page 6) De Bortoli, V ., Hutchinson, M., Wirnsberger, P., and Doucet, A. Target score matching.arXiv preprint arXiv:2402.08667, 2024. (pages...

  17. [2022]

    Spatial transformer networks.Advances in neural information processing systems, 28, 2015

    (page 6) Jaderberg, M., Simonyan, K., Zisserman, A., et al. Spatial transformer networks.Advances in neural information processing systems, 28, 2015. (pages 1, 2) Kaba, S.-O., Mondal, A. K., Zhang, Y ., Bengio, Y ., and Ravanbakhsh, S. Equivariance with learned canonicaliza- tion functions. InInternational Conference on Machine Learning, pp. 15546–15566. ...

  18. [2024]

    D., Zoph, B., Mane, D., Vasudevan, V ., and Le, Q

    (page 3) Cubuk, E. D., Zoph, B., Mane, D., Vasudevan, V ., and Le, Q. V . Autoaugment: Learning augmentation strategies from data. InProceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 113–123,

  19. [2256]

    (page 9) Bronstein, M

    PMLR, 2022. (page 9) Bronstein, M. M., Bruna, J., Cohen, T., and Veli ckovi´c, P. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges.arXiv preprint arXiv:2104.13478, 2021. (page 2) Celledoni, E., Ehrhardt, M. J., Etmann, C., Owren, B., Sch¨onlieb, C.-B., and Sherry, F. Equivariant neural net- works for inverse problems.Inverse Problems,...

This paper was first reviewed by deepseek-v4-flash on August 3, 2026.