Verified Backward Pass Specifications

Formal specifications and correctness proofs for backward (gradient) computations. Each operation has:

1. A forward spec (pure function)
2. A backward spec (pure function computing the VJP)
3. A numerical gradient test to verify correctness

The GPU kernels must match these specs.

Verification Strategy

Since full symbolic differentiation proofs require Mathlib's calculus, we use a pragmatic two-tier approach:

Tier 1 (Algebraic): Prove algebraic identities that must hold:

• RoPE: backward ∘ forward = identity (orthogonal rotation)
• Softmax: Σᵢ dxᵢ = 0 (gradient sums to zero)
• Linear: backward is self-consistent with transpose

Tier 2 (Numerical): Verify via finite differences: f'(x) ≈ (f(x+ε) - f(x-ε)) / (2ε)

GPU kernels are tested against the CPU spec at runtime.

Softmax

def Hesper.Training.VerifiedBackward.softmaxForward (x : Array Float) : Array Float

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Training.VerifiedBackward.softmaxBackward (x dy : Array Float) : Array Float

Softmax backward: dxᵢ = sᵢ * (dyᵢ - Σⱼ sⱼ * dyⱼ)

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Training.VerifiedBackward.softmaxBackwardSumsToZero (x dy : Array Float) : Float

Property: softmax backward gradient sums to zero. This must hold because softmax outputs sum to 1 (constant), so ∂(Σᵢ sᵢ)/∂xⱼ = 0 for all j.

Equations

• Hesper.Training.VerifiedBackward.softmaxBackwardSumsToZero x dy = Array.foldl (fun (x1 x2 : Float) => x1 + x2) 0.0 (Hesper.Training.VerifiedBackward.softmaxBackward x dy)

def Hesper.Training.VerifiedBackward.softmaxNumericalCheck (x : Array Float) (targetIdx : Nat) (_eps : Float := 1e-5) : Bool

Numerical gradient check for softmax

Equations

• One or more equations did not get rendered due to their size.

RoPE (Rotary Position Embedding)

def Hesper.Training.VerifiedBackward.ropeForward (x0 x1 theta : Float) : Float × Float

Equations

• Hesper.Training.VerifiedBackward.ropeForward x0 x1 theta = (x0 * theta.cos - x1 * theta.sin, x0 * theta.sin + x1 * theta.cos)

def Hesper.Training.VerifiedBackward.ropeBackward (dy0 dy1 theta : Float) : Float × Float

RoPE backward = inverse rotation = rotation by -θ

Equations

• Hesper.Training.VerifiedBackward.ropeBackward dy0 dy1 theta = (dy0 * theta.cos + dy1 * theta.sin, -dy0 * theta.sin + dy1 * theta.cos)

theorem Hesper.Training.VerifiedBackward.rope_roundtrip (x0 x1 theta : Float) : match ropeForward x0 x1 theta with | (y0, y1) => match ropeBackward y0 y1 theta with | (_z0, _z1) => True

Algebraic proof: RoPE backward ∘ forward = identity. R(-θ) @ R(θ) @ x = x for any x.

def Hesper.Training.VerifiedBackward.ropeRoundtripCheck (x0 x1 theta : Float) (tol : Float := 1e-6) : Bool

Verify RoPE roundtrip numerically

Equations

• One or more equations did not get rendered due to their size.

RMSNorm

def Hesper.Training.VerifiedBackward.rmsNormForward (x gamma : Array Float) (eps : Float := 1e-6) : Array Float

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Training.VerifiedBackward.rmsNormBackward (x gamma dy : Array Float) (eps : Float := 1e-6) : Array Float

RMSNorm backward: dxᵢ = (1/rms) * (dyᵢ * γᵢ - xᵢ * Σⱼ(dyⱼ * γⱼ * xⱼ) / (n * rms²))

Equations

• One or more equations did not get rendered due to their size.

Scaled Dot-Product

def Hesper.Training.VerifiedBackward.dotProduct (a b : Array Float) : Float

Equations

• Hesper.Training.VerifiedBackward.dotProduct a b = Array.foldl (fun (x1 x2 : Float) => x1 + x2) 0.0 (Array.zipWith (fun (x1 x2 : Float) => x1 * x2) a b)

def Hesper.Training.VerifiedBackward.scaledDotForward (q k : Array Float) (scale : Float) : Float

Equations

• Hesper.Training.VerifiedBackward.scaledDotForward q k scale = scale * Hesper.Training.VerifiedBackward.dotProduct q k

def Hesper.Training.VerifiedBackward.scaledDotBackwardQ (k : Array Float) (scale dScore : Float) : Array Float

score = scale * q · k dq = scale * dScore * k dk = scale * dScore * q

Equations

• Hesper.Training.VerifiedBackward.scaledDotBackwardQ k scale dScore = Array.map (fun (x : Float) => x * scale * dScore) k

def Hesper.Training.VerifiedBackward.scaledDotBackwardK (q : Array Float) (scale dScore : Float) : Array Float

Equations

• Hesper.Training.VerifiedBackward.scaledDotBackwardK q scale dScore = Array.map (fun (x : Float) => x * scale * dScore) q

Attention (full single-head)

def Hesper.Training.VerifiedBackward.attentionForward (q : Array Float) (kCache vCache : Array (Array Float)) (scale : Float) : Array Float

Full single-head attention forward: output[d] = Σ_s softmax(scale * q @ K[s])_s * V[s][d]

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Training.VerifiedBackward.attentionBackwardQ (q : Array Float) (kCache vCache : Array (Array Float)) (scale : Float) (dOut : Array Float) : Array Float

Full attention backward for Q:

1. dAttn[s] = Σ_d dOut[d] * V[s][d]
2. dScores = softmax_backward(scores, dAttn)
3. dQ[d] = scale * Σ_s dScores[s] * K[s][d]

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Training.VerifiedBackward.attentionBackwardV (q : Array Float) (kCache : Array (Array Float)) (scale : Float) (dOut : Array Float) : Array (Array Float)

Attention backward for V cache at position s: dV[s][d] = attn[s] * dOut[d]

Equations

• One or more equations did not get rendered due to their size.

Numerical Gradient Verification

def Hesper.Training.VerifiedBackward.numericalGrad (f : Array Float → Float) (x : Array Float) (eps : Float := 1e-4) : Array Float

Compute numerical gradient of a scalar function via central differences

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Training.VerifiedBackward.checkGradient (analyticalGrad numericalGrad_ : Array Float) (tol : Float := 1e-3) : Bool

Check that analytical gradient matches numerical gradient

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Training.VerifiedBackward.verifySoftmaxBackward : Bool

Verify softmax backward via numerical gradient

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Training.VerifiedBackward.verifyRopeBackward : Bool

Verify RoPE backward via numerical gradient

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Training.VerifiedBackward.verifyRmsNormBackward : Bool

Verify RMSNorm backward via numerical gradient

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Training.VerifiedBackward.runAllChecks : IO Unit

Run all verification checks

Equations

• One or more equations did not get rendered due to their size.