PyTorch中如何对X维度而非Y维度求和雅可比矩阵？

你需要的雅可比沿X维度求和的需求，可以通过前向模式自动微分实现，仅需一次类前向计算，开销和普通推理接近，无需构造全雅可比，也无需循环遍历Y的元素，完全适配大输入输出规模的神经网络场景。

实现原理 #

我们的目标是对Y中每个元素y_k，计算所有X元素对y_k的偏导之和，等价于计算雅可比矩阵J(Y,X)与全1的X形向量的乘积J·v。
前向模式自动微分原生支持雅可比-向量乘积运算，输入方向向量形状与X一致，输出结果形状与Y完全一致，正好匹配需求。

代码示例 #

import torch

# 示例输入
X = torch.eye(2)
X.requires_grad_()

# 你的前向计算逻辑，替换为你自己的模型前向即可
def forward(x):
    return torch.sum(x*x, dim=0)

# 前向模式梯度计算
v = torch.ones_like(X)
with torch.autograd.forward_ad.dual_level():
    # 构造对偶张量：输入+方向向量
    X_dual = torch.autograd.forward_ad.make_dual(X, v)
    Y_dual = forward(X_dual)
    # 提取切向量，即为所需结果
    res = torch.autograd.forward_ad.unpack_dual(Y_dual).tangent

print(res)
# 输出：tensor([2., 2.])，与预期完全一致

兼容老版本PyTorch的备选方案 #

如果你的PyTorch版本低于1.11（不支持前向自动微分），可以使用标量扰动+批量化反向梯度的方案，仅当Y的元素规模较小时适用：

import torch

X = torch.eye(2)
X.requires_grad_()
Y = torch.sum(X*X, dim=0)

t = torch.tensor(0., requires_grad=True)
X_t = X + t
Y_t = torch.sum(X_t * X_t, dim=0)

# 批量化计算每个Y元素对t的梯度
grad_outputs = torch.eye(Y_t.numel(), device=Y_t.device)
batch_grads = torch.autograd.grad(Y_t, t, grad_outputs=grad_outputs, is_grads_batched=True)[0]
res = batch_grads.reshape(Y_t.shape)
print(res)
# 输出：tensor([2., 2.])

内容的提问来源于stack exchange，提问作者Jonas De Schouwer