用扩散模型逆推未知变换,提升神经网络对输入变形的鲁棒性。
Inverting Data Transformations via Diffusion Sampling
- 基于能量模型建模变换后验,通过李群上的扩散过程采样。
- 在图像仿射变换和偏微分方程对称性上实现输入分布恢复。
- 适合需要测试时等变性的模型增强场景,如对抗攻击防御。
我们研究在一般李群上的变换逆问题:数据被未知群元素变换后,目标是恢复一个逆变换使其回归原始数据分布。这类未知变换广泛存在于机器学习与科学建模中,常显著扭曲观测结果。本文从概率视角出发,将变换后验建模为数据空间上由能量函数定义的玻尔兹曼分布。为此,我们提出一种在李群上运行的扩散过程,所有更新保持在流形上,仅需在相关李代数中计算。所提方法TIED(Transformation-Inverting Energy Diffusion)依赖于新的平凡化目标得分恒等式,实现变换后验的高效得分采样。作为关键应用,聚焦测试时等变性,旨在提升预训练神经网络对输入变换的鲁棒性。在图像仿射变换与偏微分方程对称性实验中,TIED成功在测试时将变换输入恢复至训练分布,性能优于强基线的归一化与采样方法。代码已公开于https://github.com/jw9730/tied。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。