Differentiating through Stochastic Differential Equations

Code accompanying the paper:

Differentiating through Stochastic Differential Equations: A Primer R. Leburu, L. Nurbekyan, and L. Ruthotto arXiv:2601.08594

Overview

This repository provides minimal, educational implementations for computing gradients through stochastic differential equations (SDEs) using the discretize-optimize approach. It supports both Ito and Stratonovich interpretations.

The main takeaways from the paper are:

1. Discretize-optimize is safe: Differentiate directly through your numerical SDE solver
2. Endpoint evaluation matters: For Ito SDEs with state-dependent diffusion, use left-point Jacobian evaluation
3. Ito-Stratonovich equivalence: Both frameworks give equivalent gradients with proper drift correction

Quick Start

# Install dependencies
pip install -r requirements.txt

# Run tests
pytest tests/ -v

# Open the tutorial notebook
jupyter notebook tutorial_figures.ipynb

Files

File	Description
adjoint_sde.py	Single-file module with all implementations (~450 lines)
tutorial_figures.ipynb	Interactive notebook reproducing all paper figures
tests/test_adjoint_sde.py	Test suite for CI

Module Contents

from adjoint_sde import (
    CEV,              # Constant Elasticity of Variance model
    EulerMaruyama,    # Ito discretization
    HeunMethod,       # Stratonovich discretization
    black_scholes_delta,  # Analytical Delta for validation
    stratonovich_drift_correction  # Ito-Stratonovich conversion
)

CEV Model

The Constant Elasticity of Variance model serves as our primary example:

dS(t) = r S(t) dt + sigma S(t)^beta dW(t)

• When beta = 1: Black-Scholes (geometric Brownian motion)
• When beta != 1: State-dependent diffusion, demonstrating the importance of proper adjoint discretization

Numerical Schemes

1. Euler-Maruyama (Ito): X_{n+1} = X_n + dt*f(X_n) + g(X_n)*dW_n
2. Heun Method (Stratonovich): Predictor-corrector scheme with trapezoidal averaging

Discrete Adjoint

The discrete adjoint computes gradients by backward recursion:

p_n = [I + dt*(df/dx)^T + dW_n^T*(dg/dx)^T] @ p_{n+1}

Key insight: Jacobians must be evaluated at the left point (t_n, X_n) for Ito SDEs.

Reproducing Paper Figures

The tutorial_figures.ipynb notebook generates all three figures from the paper:

1. Figure 1: Convergence of discrete adjoint to analytical Black-Scholes Delta
2. Figure 2: Comparison of discrete (correct) vs naive (incorrect) adjoint for CEV
3. Figure 3: Ito-Stratonovich equivalence demonstration

Citation

@article{leburu2025sde,
  title={Differentiating through Stochastic Differential Equations: A Primer},
  author={Leburu, R. and Nurbekyan, L. and Ruthotto, L.},
  journal={arXiv preprint arXiv:2601.08594},
  year={2026}
}

Acknowledgement

This work was supported in part by NSF award DMS 2038118.