Skip to main content

Closed-form derivatives

The central object in omnibias is the derivative tower of an activation:

σ(z),  σ'(z),  σ''(z),  …,  σⁿ(z)

For a special but very useful class of activations, every entry in this tower is a polynomial in the single value σ(z) you already computed. That is the whole trick.

The Riccati identity​

An activation is in the Riccati class if its derivative is a polynomial in its own value:

σ'(z) = P(σ(z))

The two canonical examples:

ActivationValueDerivativeRiccati P
sigmoids = 1/(1 + e^-z)s(1 - s)P(s) = s - s²
tanht = tanh(z)1 - t²P(t) = 1 - t²

Differentiate the identity once more and the chain rule keeps you inside polynomials in σ(z):

σ''(z) = P'(σ)·σ' = P'(σ)·P(σ)

Iterating gives a recurrence for a sequence of polynomials Tₙ with σⁿ(z) = Tₙ(σ(z)). No new transcendental evaluation is ever needed.

Three recurrences​

omnibias implements one fast-path recurrence per activation family:

  • sigmoid uses the Eulerian polynomial recursion.
  • tanh uses a Legendre-style recurrence T(n+1) = (1 - t²)·T'(n).
  • softplus reuses sigmoid's tower because softplus⁽ⁿ≥¹⁾(z) = σ⁽ⁿ⁻¹⁾(z).
  • gaussian uses the probabilist's Hermite identity g⁽ⁿ⁾(z) = (-1)ⁿ·Heₙ(z)·g(z).

Using the fast path​

Every backend activation spec exposes fastpath(z, n):

from omnibias.jax import get_activation
import jax.numpy as jnp

spec = get_activation("tanh")
z = jnp.array([0.3, 0.7])
print(spec.fastpath(z, 5))   # 5th derivative of tanh, closed form

The contract:

  • n < 0 raises ValueError.
  • An order that is genuinely unimplemented for that activation raises NotImplementedError.
  • Otherwise you get σⁿ(z) at machine precision.

Why this is fast and stable​

  • Flat in order. The cost of σⁿ is one activation evaluation plus an O(n) polynomial recurrence — it does not grow like a nested autodiff graph.
  • No catastrophic cancellation. Unlike finite differences, there is no subtraction of nearly-equal large numbers, so you do not lose ~1 digit per order.
  • Bit-identical across backends. The polynomial coefficients live once in omnibias-core (pure Python); every backend imports them, so the numbers agree to float64 ULP. See Cross-backend parity.

From the tower to operators​

A scalar field f(x) = Σ_h c_h·σ(W_h·x + β_h) has a Laplacian that is a direct contraction of the second entry of the tower:

∇²f(x) = Σ_h c_h·σ''(z_h)·‖W_h‖²,    where  z_h = W_h·x + β_h

The only dimension-dependent term, ‖W_h‖², is computed once and reused — which is why the Laplacian overhead is O(1) in the input dimension D. Iterating this gives the poly-Laplacian Δᵏ for high-order PDEs.

Continue
Last updated on