如何用torch.autograd.grad无循环计算随机过程路径对初值的梯度

问题描述 #

我有一条从初始点出发的随机过程采样路径，代码实现如下：

import torch
import torch.nn as nn
import torchsde

class SDE_ou_1d(nn.Module):
    def __init__(self):
        super().__init__()
        self.sde_type = "ito"
        self.noise_type = "diagonal"

    def f(self, t, y): # 漂移项
        return -y

    def g(self, t, y): # 扩散项
        return torch.ones_like(y)

t_vec = torch.linspace(0, 1, 100)  # 时间数组
mySDE = SDE_ou_1d()

x0 = torch.zeros(1, 1, requires_grad=True).to(t_vec)
X_t = torchsde.sdeint(mySDE, x0, t_vec, method='euler')

我希望用torch.autograd.grad()计算该路径对初始条件的梯度，得到与X_t形状相同（即100x1）的输出，反映路径在每个时间点的变化。

尝试以下代码时，梯度会对所有t值求和，无法得到每个时间点的单独梯度：

X_grad = torch.autograd.grad(outputs=X_t, inputs=x0,
                           grad_outputs=torch.ones_like(X_t),
                           create_graph=False, retain_graph=True, only_inputs=True, allow_unused=True)[0]

用循环逐个计算虽然可行，但速度极慢：

X_grad_loop = torch.zeros_like(X_t)

for i in range(X_t.shape[0]):  # 遍历X_t的时间维度
    grad_i = torch.autograd.grad(outputs=X_t[i,...], inputs=x0,
                                    grad_outputs=torch.ones_like(X_t[i,...]),
                                    create_graph=False, retain_graph=True, only_inputs=True, allow_unused=True)[0]
    X_grad_loop[i,...] = grad_i

请问是否存在无需循环，直接用torch.autograd.grad()计算该梯度的方法？

解决方案 #

可以利用Jacobian矩阵的批量计算特性，通过调整grad_outputs为单位矩阵的形式，让torch.autograd.grad一次性输出每个时间点梯度组成的张量，彻底避免循环。

核心思路 #

当outputs是形状为(T, 1)的张量，inputs是形状为(1,1)的张量时，我们需要计算的是Jacobian矩阵的每一行（对应每个时间点的输出对初始条件的梯度）。通过构造grad_outputs为单位矩阵扩展维度后的形式，可以让自动梯度系统针对每个时间点的输出单独计算梯度，最终拼接成与X_t形状一致的结果。

代码实现 #

# 获取时间步长数量
T = X_t.shape[0]

# 构造grad_outputs：形状为(T, T, 1)，每个位置(T,i,1)对应第i个时间点的单位梯度信号
grad_outputs = torch.eye(T, device=X_t.device).unsqueeze(-1)

# 计算批量梯度
X_grad = torch.autograd.grad(
    outputs=X_t,
    inputs=x0,
    grad_outputs=grad_outputs,
    create_graph=False,
    retain_graph=True,
    only_inputs=True,
    allow_unused=True
)[0]

# 调整维度至(T, 1)，与X_t形状完全匹配
X_grad = X_grad.squeeze(0).transpose(0, 1)

关键细节解释 #

• torch.eye(T)生成T×T的单位矩阵，unsqueeze(-1)将其扩展为T×T×1的张量。这样每个grad_outputs[i]是仅第i个位置为1、其余为0的张量，对应只触发第i个时间点输出对初始条件的梯度计算。
• 初始计算得到的梯度形状为(1, T, 1)，通过squeeze(0)去掉多余的维度，再用transpose(0,1)交换维度，最终得到(T,1)的张量，与X_t的形状完全一致。

这种方法利用PyTorch的批量梯度计算能力，效率远高于循环版本，且代码简洁易维护。

内容的提问来源于stack exchange，提问作者GigaByte123