神经随机微分方程的高效精确梯度计算

神经 SDE 结合了 RNN 与 SDE 的诸多优点：内存高效的训练、高容量函数逼近以及对模型空间的强先验。这使其成为建模许多类型时间动力学的自然选择。训练神经 SDE（作为 VAE 或 GAN）需要通过对 SDE 求解进行反向传播。这可以通过求解一个时间反向的 SDE 来实现，其解即为所需的参数梯度。然而，由于高计算成本与数值截断误差，该方法此前存在严重的速度与精度问题。在此，我们通过若干技术创新克服了这些问题。首先，我们引入可逆 Heun 方法。这是一种新的 SDE 求解器，具有代数可逆性：消除了数值梯度误差，且据我们所知是首个此类求解器。此外，它所需的函数评估次数仅为同类求解器的一半，带来最高 $1.98\times$ 的加速。其次，我们引入布朗区间：一种快速、内存高效且精确采样及重建布朗运动的新方法。借此，相较此前既近似又相对缓慢的技术，我们获得了最高 $10.6\times$ 的速度提升。第三，当具体将神经 SDE 作为 GAN 训练时（Kidger 等，2021），我们展示了如何通过谨慎的权重裁剪与激活函数选择来训练 SDE-GAN。这降低了计算成本（最高 $1.87\times$ 加速）并消除了与梯度惩罚相关的数值截断误差。总体而言，在训练速度以及分类、预测和 MMD 测试指标方面，我们以显著优势超越最先进水平。我们已将所有技术的实现贡献给 torchsde 库，以促进其应用。