loss_surrogate

MSE Surrogate技巧

在深度学习的常规实践中，我们定义一个标量损失函数（Loss），然后依赖自动微分（Autograd）框架来计算模型参数的梯度，并以此更新网络。但设想一个特殊的场景：我们不关心损失值本身，而是通过某种理论分析或外部模块，已经精确地知道了某个中间张量 x 应有的梯度 g_x。我们该如何将这个“钦定”的梯度 g_x “注入”到反向传播的链条中呢？

这就是 MSE Surrogate 的使用情景。它为我们提供了一种优雅且高效的方式，在不破坏自动微分框架工作流的前提下，精确地指定任意张量的梯度

目标场景

我们有一个由参数 $\theta$ 控制的生成网络 $G$ ，其输出为 $x = G(\theta, z)$ 。在某些高级应用中（如某些强化学习策略、生成模型的特定优化技巧），我们可能已经得到了一个理想的梯度向量：

g_x = \text{desired gradient} \in \mathbb{R}^{\text{shape}(x)}

我们希望模型参数 $\theta$ 的最终更新方向恰好是这个理想梯度 $g_x$ 通过雅可比-向量积（JVP）反向传播的结果：

\nabla_\theta L = \left(\frac{\partial x}{\partial \theta}\right)^\top g_x

直接计算并存储巨大的雅可比矩阵 $(\frac{\partial x}{\partial \theta})$ 是不现实的。我们需要一个标量损失 $L_{\text{sur}}$ ，它必须满足一个核心条件：其对于中间张量 $x$ 的梯度恰好等于我们想要的 $g_x$ 。

\nabla_x L_{\text{sur}} = g_x

只要满足这个条件，自动微分框架便会为我们无缝地完成后续的链式法则计算，得到正确的 $\nabla_\theta L$

MSE Surrogate 的核心思想与构造

MSE Surrogate 的巧妙之处在于，它构造了一个看似普通的均方误差损失，但其梯度却精确地等于我们指定的 g_x。

核心公式:

L_{\text{sur}} = \tfrac12 \|x - \text{stopgrad}(x - g_x)\|_2^2

让我们来解析这个构造：

我们定义一个“目标” target，即 target = x - g_x。
关键一步：我们使用 stop_gradient（在 PyTorch 中是 .detach()）来处理这个 target，使其在反向传播的计算图中被视为一个常数
然后，我们计算 x 与这个“伪造”的常数 target 之间的 MSE 损失

在反向传播时，由于 target 被视为常数，损失 $L_{\text{sur}}$ 对 $x$ 的梯度为：

\nabla_x L_{\text{sur}} = \nabla_x \left( \tfrac12 \|x - \text{target}\|^2 \right) = x - \text{target}

将 target 的定义代入：

\nabla_x L_{\text{sur}} = x - (x - g_x) = g_x

就这样我们精确地注入了梯度 g_x。

# 假设 x 是需要注入梯度的网络输出
# g_x 是我们计算出的、形状与 x 相同的目标梯度
with torch.no_grad():
    g_x = compute_desired_gradient(x)

# 核心步骤：构造 target 并 detach
target = (x - g_x).detach()

# 构造 surrogate loss
loss_sur = 0.5 * (x - target).pow(2).sum() # 使用 .sum() 来避免缩放问题

# 现在可以像普通 loss 一样使用
loss_sur.backward()

替代方案：线性Surrogate及其对比

除了 MSE Surrogate，我们还有其他方法可以实现相同的目标吗？答案是肯定的。最直接的替代方案是线性 Surrogate：

L_{\text{linear}} = (x \cdot \operatorname{stopgrad}(g_x)).sum()

其中 $g_x$ 同样需要被 detach 以免引入不必要的二阶项。对这个损失求导，结果直接就是 $g_x$ 。

\nabla_x L_{\text{linear}} = g_x

为什么实践中 MSE Surrogate 更受欢迎呢？

特性	MSE Surrogate	线性 Surrogate `(x * g_x).sum()`
对 x 的梯度	$g_x$	$g_x$
Loss 数值特性	始终非负，易于监控。当梯度注入成功时，loss值接近 `0.5 *
日志可读性	像一个标准的优化目标，易于理解和调试	数值波动大，难以判断优化状态
框架兼容性	许多训练框架（如Lightning）期望损失是非负的	可能不满足某些框架的假设
直观性	“将 x 拉向一个目标”的比喻，易于团队协作理解	“最大化点积”的解释相对抽象

虽然线性 Surrogate 在数学上等价且更简洁，但 MSE Surrogate 因其良好的数值特性、可读性和框架兼容性，在实践中通常是更稳健和首选的方案

等价的 `x.backward(g_x)` 写法

对于简单场景，PyTorch 提供了一个更直接的快捷方式：

1 2	`# g_x 必须与 x 无计算图依赖，或已被 detach x.backward(gradient=g_x)`

这在功能上与 MSE/线性 Surrogate 等价，都是为了实现 $J^\top g_x$ 的计算。但它的缺点是不产生一个标量损失值，这使得它难以与其他损失组合、不方便日志记录，也无法用于需要标量损失的训练框架中。因此，Surrogate 方法在复杂项目中更具优势

代码验证

验证 $x=Wz, \nabla_W L = g_x z^\top$ 。

import torch

d_x, d_z = 4, 3
W = torch.randn(d_x, d_z, requires_grad=True)
z = torch.randn(d_z)
x = W @ z

g_x = torch.randn_like(x)

# Surrogate 方法
target = (x - g_x).detach()
loss_sur = 0.5 * (x - target).pow(2).sum()
loss_sur.backward()

# 理论梯度
grad_theory = g_x.unsqueeze(1) @ z.unsqueeze(0)

print("Autograd 和理论梯度的最大差值:", (W.grad - grad_theory).abs().max().item())
# 差值应接近于0

总结

MSE Surrogate 是一种强大而灵活的技巧，它将一个看似困难的“梯度注入”问题，转化为一个自动微分框架可以轻松处理的标准 MSE 损失问题。它的核心是通过 detach() 构造一个伪常数目标，从而将你想要的梯度向量伪装成 MSE 的残差。

E2E_auto_driving 下一篇