TorchSDE 项目解析:基于 PyTorch 的随机微分方程求解指南
随机微分方程基础
随机微分方程(Stochastic Differential Equations, SDEs)是描述包含随机扰动系统的微分方程,广泛应用于金融、物理、生物和机器学习等领域。TorchSDE 项目提供了高效求解以下形式 SDE 的功能:
dy(t) = f(t, y(t)) dt + g(t, y(t)) dW(t)
其中:
t
是标量时间变量y(t)
是 n 维状态向量f(t, y(t))
是 n 维漂移项g(t, y(t))
是 n×m 维扩散项矩阵W(t)
是 m 维布朗运动
核心功能:sdeint 函数
TorchSDE 的核心是 sdeint
函数,其基本调用方式如下:
from torchsde import sdeint
ys = sdeint(sde, y0, ts)
参数详解
-
sde 对象:
- 必须是
torch.nn.Module
子类 - 需要定义
sde_type
("ito" 或 "stratonovich") - 需要定义
noise_type
("scalar"、"additive"、"diagonal" 或 "general") - 必须实现
f(t, y)
方法计算漂移项 - 必须实现
g(t, y)
方法计算扩散项 - 可选实现
h(t, y)
方法计算先验漂移项(用于KL散度计算)
- 必须是
-
初始条件 y0:
- 形状为
(batch_size, state_size)
的张量 - 表示在时间
ts[0]
时的初始状态
- 形状为
-
时间序列 ts:
- 形状为
(t_size,)
的张量 - 指定需要输出解的时间点
- 形状为
噪声类型详解
-
标量噪声 (scalar):
- 扩散项输出形状为
(batch_size, state_size, 1)
- 使用1维布朗运动
- 扩散项输出形状为
-
加性噪声 (additive):
- 扩散项不依赖于状态 y
- 输出形状为
(batch_size, state_size, brownian_size)
- 使用多维布朗运动
-
对角噪声 (diagonal):
- 扩散项是逐元素的
- 输出形状为
(batch_size, state_size)
- 使用与状态维度相同的布朗运动
-
一般噪声 (general):
- 最通用的形式
- 输出形状为
(batch_size, state_size, brownian_size)
- 使用多维布朗运动
求解器选择指南
Ito SDE 求解器
-
Euler-Maruyama 方法:
- 最简单的求解器
- 计算成本最低
- 精度较低
-
Milstein 方法:
- 中等计算成本
- 比 Euler 方法精度更高
- 不支持一般噪声
-
随机 Runge-Kutta 方法 (SRK):
- 计算成本最高
- 精度最高
- 同样不支持一般噪声
Stratonovich SDE 求解器
-
Euler-Heun 方法:
- 计算成本最低
- 适用于不需要高精度的场景
-
Heun 方法:
- 中等计算成本
- 比 Euler-Heun 精度更高
-
可逆 Heun 方法 (reversible_heun):
- 特别适合伴随方法
- 计算效率高
- 数值误差小
求解器选择建议
-
训练神经网络 SDE(不使用伴随方法):
- Ito SDE:选择 "euler"
- Stratonovich SDE:选择 "reversible_heun"
-
使用伴随方法训练:
- 强烈推荐 Stratonovich SDE + "reversible_heun"
伴随方法详解
TorchSDE 提供两种反向传播方式:
-
直接反向传播:
- 通过求解器内部操作反向传播
- 使用
sdeint
函数 - 内存消耗较大
-
伴随方法:
- 通过求解伴随 SDE 计算梯度
- 使用
sdeint_adjoint
函数 - 内存效率高
- 计算时间较长
伴随方法使用技巧
- 优先使用 Stratonovich SDE,其伴随 SDE 计算成本更低
- 使用
method="reversible_heun"
和adjoint_method="adjoint_reversible_heun"
组合 - 对于精度要求高的场景,考虑使用自适应步长或减小步长
布朗运动控制
TorchSDE 提供了灵活的布朗运动控制方式:
from torchsde import BrownianInterval
bm = BrownianInterval(t0=0., t1=1., size=(4, 1))
关键参数
-
确定性随机数:
BrownianInterval(..., entropy=42, tol=1e-5, halfway_tree=True)
-
速度优化:
BrownianInterval(..., cache_size=None) # 使用更多内存提高速度
-
Levy 区域近似:
- 用于高阶求解器
- 可选 "none"、"space-time"、"davie" 或 "foster"
实际应用建议
-
神经网络 SDE 训练:
- 对于常规训练,使用 Euler 或 reversible_heun 求解器
- 对于需要高精度梯度的场景,使用伴随方法
-
数值模拟:
- 根据精度需求选择 Milstein 或 SRK 方法
- 对于长时间模拟,考虑使用自适应步长
-
可重复实验:
- 使用固定的随机种子 (entropy)
- 设置适当的容差 (tol)
TorchSDE 项目为 PyTorch 生态提供了强大的 SDE 求解能力,特别适合机器学习研究和应用场景。通过合理选择求解器和参数配置,可以在计算效率和数值精度之间取得良好平衡。
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考