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 →
Inverting Data Transformations via Diffusion Sampling
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
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.
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
- 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.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
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)
- [§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.
- [§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)
- [§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.
- [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.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.
- [§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.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
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
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)
- Step size Δt =
1/50 (images/PDE), 1/100 (SO(10))
- MC sample size N =
2-100 (100 for SO(10), 2-4 images, 10 PDE)
axioms (5)
- domain assumption G is a connected Lie group and the action is free and diffeomorphic
- domain assumption The noisy densities pt are smooth and strictly positive, and energy normalization is finite for noncompact groups
- ad hoc to paper k1 approximates p1 so the reverse process can be initialized from k1
- ad hoc to paper Euler discretization (line 12, Eq. 24) converges to the continuous reverse and forward SDEs
- standard math Monte Carlo estimator integrability/positivity conditions
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}
}
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
Forward citations
Cited by 1 Pith paper
-
What are the Right Symmetries for Formal Theorem Proving?
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
-
[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...
2024
-
[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;...
2019
-
[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...
1994
-
[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,
-
[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,
-
[15]
The posterior distribution ofghas densityp(g| ˜x)∝p x g−1 · ˜x |detJ g−1 (˜x)|
-
[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 ...
2012
-
[17]
A method to samplew∼k t,
-
[18]
A method to calculatelogλ(w),
-
[1993]
(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...
Pith/arXiv arXiv 2021
-
[1996]
(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...
Pith/arXiv arXiv 2024
-
[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,...
2012
-
[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...
2019
-
[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...
arXiv 2025
-
[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,...
Pith/arXiv arXiv 2016
-
[2019]
(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...
Pith/arXiv arXiv 2022
-
[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. ...
Pith/arXiv arXiv 2015
-
[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,
-
[2256]
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,...
Pith/arXiv arXiv 2022
This paper was first reviewed by deepseek-v4-flash on August 3, 2026.
discussion (0)
Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.