Automatic Stochastic Differentiation in Lean 4

Stochastic calculus is a fundamental tool in mathematical finance, physics, and many other fields. In this post, I’ll implement automatic stochastic differentiation in Lean 4, inspired by Ji-Ha Kim’s excellent article.

The Key Idea

In standard calculus, when we compute derivatives, higher-order infinitesimals vanish. If we denote an infinitesimal by \(dt\), then any term like \((dt)^k\) with \(k > 1\) is negligible in the limit.

However, in stochastic calculus, a key difference arises: the quadratic variation of a Brownian motion \(B_t\) is nonzero. In fact, we have \(dB_t^2 = dt\). This means that when differentiating functions of a Brownian path, second-order terms cannot be dropped.

Algebraic Representation

To capture this behavior algebraically, we use the ring \(\mathbb{R}[\epsilon]/\epsilon^3\), with the following interpretation:

• \(\epsilon \sim dB_t\) (the Brownian increment)
• \(\epsilon^2 \sim dt\) (the time increment)
• \(\epsilon^3 = 0\) (all higher-order terms vanish)

For a stochastic differential equation such as \(dX_t = \mu\, dt + \sigma\, dB_t\), we write \(dX_t = \sigma\,\epsilon + \mu\,\epsilon^2\).

Implementation in Lean 4

Let’s implement this in Lean 4 by defining a DualNumber structure:

structure DualNumber where
  real : Float -- Standard real part
  e : Float    -- Coefficient for ε (Brownian increment)
  e2 : Float   -- Coefficient for ε² (time increment)
  deriving Repr

We then define operations on these dual numbers, ensuring that \(\epsilon^3 = 0\):

/-- Multiplication of dual numbers, with ε³ = 0 -/
def mul (a b : DualNumber) : DualNumber :=
  { real := a.real * b.real
  , e := a.real * b.e + a.e * b.real
  , e2 := a.real * b.e2 + a.e * b.e + a.e2 * b.real
    -- Note: a.e * b.e contributes to e2 because ε * ε = ε²
    -- Terms with ε³ are dropped (ε³ = 0)
  }

Deriving Itô’s Lemma

Using our dual number system, we can derive Itô’s lemma. For a function \(f(X_t)\), Itô’s lemma states:

Let’s verify this for \(f(X) = X^2\):

def itoLemmaExample (x mu sigma : Float) : String := 
  let X : DualNumber := DualNumber.ofFloat x
  let dX := DualNumber.itoIncrement sigma mu
  let X_dual := X + dX
  
  -- Compute f(X + dX) = (X + dX)²
  let f_X_sq := DualNumber.square X_dual
  
  -- Subtract f(X) to get df
  let df_dual := f_X_sq - DualNumber.square X
  
  -- Extract the Itô differential components
  let (dt_coef, dB_coef) := DualNumber.itoDifferential df_dual
  
  -- Expected values from Itô's lemma
  let expected_dt := 2 * x * mu + sigma * sigma
  let expected_dB := 2 * x * sigma
  
  s!"Differential df:\ndf = {dt_coef} dt + {dB_coef} dB\n\nExpected from Itô's lemma:\ndf = {expected_dt} dt + {expected_dB} dB"

Stratonovich Calculus

In Stratonovich calculus, the chain rule holds in its usual form without correction terms. We can implement this by modifying the increment:

/-- Create a Stratonovich increment: dX = σ*(ε + ε²/2) + μ*ε² -/
def stratonovichIncrement (sigma mu : Float) : DualNumber :=
  { real := 0
  , e := sigma
  , e2 := sigma / 2 + mu }

Running the Examples

When we run our examples with \(X = 5.0\), \(\mu = 2.0\), and \(\sigma = 3.0\), we get:

Itô Calculus Example:
Differential df:
df = 29.0 dt + 30.0 dB

Expected from Itô's lemma:
df = 29.0 dt + 30.0 dB

Stratonovich Calculus Example:
Stratonovich Differential df:
df = 27.5 dt + 30.0 dB

Expected from Stratonovich calculus:
df = 27.5 dt + 30.0 dB

This confirms that our implementation correctly reproduces both Itô and Stratonovich calculus.

Why This Matters

This algebraic approach to stochastic calculus has several advantages:

1. Clarity: It makes the relationship between Itô and Stratonovich calculus explicit
2. Computability: It’s easy to implement in computer algebra systems
3. Generality: The approach extends to more complex stochastic processes

The full implementation is available in the Blog/StochasticDiff.lean file in the repository.

Conclusion

Automatic stochastic differentiation provides an elegant algebraic framework for stochastic calculus. By working in the ring \(\mathbb{R}[\epsilon]/\epsilon^3\) and interpreting \(\epsilon\) as \(dB_t\) and \(\epsilon^2\) as \(dt\), we can derive Itô’s lemma and understand the difference between Itô and Stratonovich calculus.

This implementation in Lean 4 demonstrates how formal verification systems can be used to implement and verify stochastic calculus concepts.