可微分随机微分方程求解器的 PyTorch 实现

本库提供了支持 GPU 且具有高效反向传播功能的随机微分方程 (SDE) 求解器。

---

安装

pip install torchsde

要求： Python >=3.8 和 PyTorch >=1.6.0。

文档

可在 这里 查阅。

示例

快速示例

import torch
import torchsde

batch_size, state_size, brownian_size = 32, 3, 2
t_size = 20

class SDE(torch.nn.Module):
    noise_type = 'general'
    sde_type = 'ito'

    def __init__(self):
        super().__init__()
        self.mu = torch.nn.Linear(state_size, 
                                  state_size)
        self.sigma = torch.nn.Linear(state_size, 
                                     state_size * brownian_size)

    # 驱动项
    def f(self, t, y):
        return self.mu(y)  # 形状为 (batch_size, state_size)

    # 扩散项
    def g(self, t, y):
        return self.sigma(y).view(batch_size, 
                                  state_size, 
                                  brownian_size)

sde = SDE()
y0 = torch.full((batch_size, state_size), 0.1)
ts = torch.linspace(0, 1, t_size)
# 初始状态 y0，SDE 在区间 [ts[0], ts[-1]] 上求解。
# ys 的形状为 (t_size, batch_size, state_size)
ys = torchsde.sdeint(sde, y0, ts)

笔记本

examples/demo.ipynb 提供了一个关于如何求解 SDE 的简短指南，其中包括一些细微之处，例如固定求解器中的随机性以及 噪声类型 的选择。

隐式 SDE

examples/latent_sde.py 学习一个 隐式随机微分方程，如 [1] 第 5 节所述。该示例将 SDE 拟合到数据上，同时通过正则化使其类似于 Ornstein-Uhlenbeck 先验过程。该模型可以粗略地看作是一个 变分自编码器，其先验分布和近似后验分布均为 SDE。可以通过以下命令运行此示例：

python -m examples.latent_sde --train-dir <TRAIN_DIR>

程序会将图像输出到 <TRAIN_DIR> 指定的路径。在默认超参数下，训练应在 500 次迭代后趋于稳定。

作为 GAN 的神经 SDE

examples/sde_gan.py 学习一个作为 GAN 的 SDE，如 [2]、[3] 所述。该示例训练一个 SDE 作为 GAN 的生成器，同时使用一个 神经 CDE [4] 作为判别器。可以通过以下命令运行此示例：

python -m examples.sde_gan

引用

如果您在研究中发现本代码库很有用，请考虑引用以下一篇或两篇论文：

@article{li2020scalable,
  title={Scalable gradients for stochastic differential equations},
  author={Li, Xuechen and Wong, Ting-Kam Leonard and Chen, Ricky T. Q. and Duvenaud, David},
  journal={International Conference on Artificial Intelligence and Statistics},
  year={2020}
}

@article{kidger2021neuralsde,
  title={Neural {SDE}s as {I}nfinite-{D}imensional {GAN}s},
  author={Kidger, Patrick and Foster, James and Li, Xuechen and Oberhauser, Harald and Lyons, Terry},
  journal={International Conference on Machine Learning},
  year={2021}
}

参考文献

[1] Xuechen Li, Ting-Kam Leonard Wong, Ricky T. Q. Chen, David Duvenaud. “随机微分方程的可扩展梯度”。 国际人工智能与统计会议，2020 年。[arXiv]

[2] Patrick Kidger, James Foster, Xuechen Li, Harald Oberhauser, Terry Lyons. “神经 SDE 作为无限维 GAN”。 国际机器学习会议，2021 年。[arXiv]

[3] Patrick Kidger, James Foster, Xuechen Li, Terry Lyons. “神经 SDE 的高效准确梯度”。2021 年。[arXiv]

[4] Patrick Kidger, James Morrill, James Foster, Terry Lyons，“用于不规则时间序列的神经控制微分方程”。 神经信息处理系统，2020 年。[arXiv]

---

这是一个研究项目，而非 Google 的官方产品。

torchsde 快速上手指南

torchsde 是一个基于 PyTorch 的随机微分方程（SDE）求解器库，支持 GPU 加速和高效反向传播，适用于构建潜变量 SDE、神经 SDE-GAN 等前沿模型。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• Python: 版本 >= 3.8
• PyTorch: 版本 >= 1.6.0
• 系统: 支持 Linux, macOS, Windows (需配置好 CUDA 以启用 GPU 加速)

国内加速建议：如果下载 PyTorch 或依赖包速度较慢，推荐使用清华或阿里镜像源。

pip install torch torchvision torchaudio -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

使用 pip 直接安装官方发布版本：

pip install torchsde

若需使用国内镜像加速安装：

pip install torchsde -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

以下是最小化的使用示例，演示如何定义一个 SDE 模型并求解。

代码示例

import torch
import torchsde

# 设置维度参数
batch_size, state_size, brownian_size = 32, 3, 2
t_size = 20

class SDE(torch.nn.Module):
    # 指定噪声类型和 SDE 类型 (Ito 或 Stratonovich)
    noise_type = 'general'
    sde_type = 'ito'

    def __init__(self):
        super().__init__()
        # 定义漂移项 (Drift) 网络
        self.mu = torch.nn.Linear(state_size, state_size)
        # 定义扩散项 (Diffusion) 网络
        self.sigma = torch.nn.Linear(state_size, state_size * brownian_size)

    # 漂移函数 f(t, y)
    def f(self, t, y):
        return self.mu(y)  # 形状: (batch_size, state_size)

    # 扩散函数 g(t, y)
    def g(self, t, y):
        return self.sigma(y).view(batch_size, state_size, brownian_size)

# 实例化模型
sde = SDE()

# 初始状态 y0
y0 = torch.full((batch_size, state_size), 0.1)

# 时间网格 ts，求解区间为 [ts[0], ts[-1]]
ts = torch.linspace(0, 1, t_size)

# 执行求解
# ys 的输出形状为 (t_size, batch_size, state_size)
ys = torchsde.sdeint(sde, y0, ts)

说明

1. 定义模型：继承 torch.nn.Module，必须设置 noise_type 和 sde_type 类属性，并实现 f (漂移) 和 g (扩散) 方法。
2. 准备数据：创建初始状态张量 y0 和时间点张量 ts。
3. 求解：调用 torchsde.sdeint 即可得到求解轨迹，该操作完全可微，可直接用于后续的损失计算和反向传播。