技术摘要：通过扩散采样实现数据变换的反转

问题陈述

本文研究了通用李群（Lie groups）上的**变换反转（transformation inversion）**问题。在许多机器学习和科学建模场景中，观测到的数据 $\tilde{x}$ 是通过对分布内样本 $x$ 应用一个未知的李群 $g$ 变换生成的（即 $\tilde{x} = g \cdot x$）。目标是恢复一个逆变换 $g^{-1}$，该变换能将分布外（out-of-distribution）的观测值映射回原始数据分布。

由于这种扰动变换是未知的，且解通常不唯一，这被形式化为一个盲逆问题（blind inverse problem）。现有方法通常在数据空间内操作，或专门针对特定的变换族（如仿射群或旋转群），而本研究的目标是通用的李群作用，包括非紧致、非阿贝尔且不具备双不变度量（bi-invariant metrics）的群（例如图像单应性 $PGL(3, \mathbb{R})$ 和偏微分方程的李点对称性）。

本文强调的一个关键应用是测试时等变性（test-time equivariance）。预训练神经网络往往缺乏对几何变换的鲁棒性。作者提出在测试时反转这些变换，将输入带回模型的训练分布，从而在无需专门的等变架构或重新训练的情况下提高鲁棒性。

方法论：变换反转能量扩散 (TIED)

所提出的方法 TIED 将变换的后验概率 $p(g | \tilde{x})$ 建模为定义在数据空间上的能量函数 $E$ 的玻尔兹曼分布。该能量函数源自数据先验 $p_x(x)$，其中 $E(x) = -\log p_x(x)$。后验密度由下式给出：
$$p(g | \tilde{x}) \propto p_x(g^{-1} \cdot \tilde{x}) |\det J_{g^{-1}}(\tilde{x})|$$

为了从该后验分布中采样，TIED 直接在李群上引入了一个扩散过程。该方法克服了两个主要挑战：能量景观（通常源自神经网络）的崎岖多峰性质，以及流形上采样的几何约束。

1. 李群上的平凡化扩散 (Trivialized Diffusion on Lie Groups)

TIED 并非在任意群元素的切空间中进行更新，而是利用**平凡化（trivialization）**技术。更新在李代数 $\mathfrak{g}$（单位元处的切空间）中进行，并通过指数映射传输到当前的群元素。

• 前向过程： 在群上定义一个方差爆炸的随机微分方程（SDE），通过在李代数中抽取无穷小扰动，并通过右乘将其传输：
  $$dg_t = d(L_{g_t})_e [\gamma(t) dw^{\mathfrak{g}}_t]$$
  其中 $w^{\mathfrak{g}}_t$ 是李代数中的布朗运动。
• 反向过程： 推导反向时间 SDE 以恢复后验分布的样本。至关重要的是，漂移项涉及平凡化得分（trivialized score） $\nabla_g \log p_t(g_t)$，它是一个李代数元素。

2. 平凡化目标得分恒等式 (Trivialized Target Score Identity)

一个核心理论贡献是一个新的平凡化目标得分恒等式。标准的得分匹配（score matching）需要访问干净的样本，但在这种反转场景中这些样本是不可用的。相反，TIED 将噪声分布的得分表示为干净能量函数的梯度。
作者证明，对于一般的李群（不要求具有双不变度量），可以通过对噪声核 $k_t$ 的蒙特卡洛期望来估计平凡化得分：
$$\nabla_g \log p_t(g_t) = \nabla_g \log \int_G k_t(w) \exp\left(-E(g_t w^{-1}) - \log \lambda(w)\right) d\mu(w)$$
其中 $\lambda(w)$ 解释了反转下 Haar 测度的变化。这使得仅通过能量函数 $E$ 及其梯度（通过自动微分计算）即可估计得分。

3. 实际实现

采样算法（算法 1）对反向 SDE 进行离散化。在每一步中，它执行以下操作：

1. 从噪声核 $k_t$ 中采样噪声 $w$。
2. 计算变换候选者的能量 $E$。
3. 使用蒙特卡洛恒等式估计平凡化得分（采用 log-sum-exp 技巧以保证稳定性）。
4. 使用指数映射更新群元素。

该方法支持多种能量定义，包括源自变分自编码器（VAE）或直接源自预训练分类器预测置信度的能量。

核心贡献

1. TIED 采样器： 一种用于在一般李群上反转数据变换的新型扩散采样器。它保证更新始终保持在流形上，并能处理非紧致、非线性的作用，且仅需了解李代数的信息。
2. 理论框架： 证明了变换反转可用于预训练模型的测试时等变性。作者证明了变换反转后验具有 $G$-等变性，确保了该方法的合理性。
3. 新的得分恒等式： 将目标得分恒等式推广到不具备双不变度量的李群，从而能够利用能量梯度而非干净数据样本进行高效的基于得分的采样。
4. 实证验证： 在图像单应性（$PGL(3, \mathbb{R})$）和 PDE 对称性（$SL(2, \mathbb{R}) \ltimes H(1, \mathbb{R})$）上的演示，展示了优于强基准方法的持续性能提升。

实验结果

论文在三个主要任务上评估了 TIED：

1. $SO(10)$ 上的合成采样： TIED 成功地在高维李群（维度 $dim=45$）的多峰能量分布上进行采样，其收敛速度明显快于平凡化动力学 Langevin 采样器（trivialized kinetic Langevin sampler）。
2. 图像分类 (MNIST)：
   • 设置： 测试了一个遭受未知仿射（$Aff(2, \mathbb{R})$）和单应性（$PGL(3, \mathbb{R})$）变换的 MNIST 图像上的预训练 ResNet18。
   • 基准模型： 与基于优化的方法（LieLAC, FoCal, ITS）以及采样方法（Kin kinetic Langevin）进行对比。
   • 结果： TIED 显著优于所有基准模型。例如，在采用分类器能量的仿射变换下，TIED 达到了 82.64% 的准确率，而原始 ResNet18 为 55.48%，Kinetic Langevin 为 74.72%。在某些设置下，它甚至优于专门的等变网络（affConv）。
3. PDE 求解：
   • 设置： 测试了一个 DeepONet 神经算子在未知李点对称性变换下的 1D 热传导方程和 Burgers 方程。
   • 结果： TIED 在所有测试用例中均实现了最低的相对 $L_2$ 误差，优于 LieLAC 和 FoCal。对于具有 $SL(2, \mathbb{R}) \ltimes H(1, \mathbb{R})$ 对称性的 1D 热传导方程，TIED 将误差降低至 0.043，相比之下，LieLAC 为 0.066，未经修正的 DeepONet 为 0.690。

意义与主张

作者声称 TIED 提供了一个无需训练的流水线，使任何预训练模型在测试时都具备对李群变换的等变性。通过利用能量模型（可以通过预训练模型本身来近似），该方法避免了对专门等变架构的需求或昂贵的重新训练。

其重要性在于：

• 通用性： 它适用于广泛的李群类别，包括那些非紧致且不具备双不变度量的群，而这些群对以往的扩散或优化方法来说具有挑战性。
• 鲁棒性： 它有效地将变换后的输入恢复到训练分布，显著提高了预训练网络对分布外几何变换的泛化能力。
• 效率： 与传统的 MCMC 或在崎岖能量景观上的迭代优化相比，基于扩散的方法结合平凡化得分恒等式，提供了一种更高效、更稳定的采样过程。

论文总结道，这种方法代表了一种等变性的推理时扩展（inference-time scaling for equivariance），为在大模型时代处理分布外变换提供了实用的解决方案，而在该时代，处理此类问题至关重要。