Skip to content

← JAX Autodiff step 10 of 25

Medium Primitives

Custom VJP: Stable log1pexp

Jump to the task

Softplus is the elementwise function log(1 + exp(x)). Its derivative is the sigmoid function. The direct formula can overflow in float32 because it forms exp(x) before taking the logarithm.

At x = 100, the mathematical softplus value is approximately 100, but exp(100) exceeds float32's range. A custom backward rule alone does not repair that infinite forward value. Both paths need suitable numerical formulas.

Stabilize the forward calculation

Use jnp.logaddexp(0.0, x), which calculates the log of the sum of two exponentials without directly forming the large exponential. It represents the same softplus function. An equivalent real-arithmetic identity is:

softplus(x) = max(x, 0) + log1p(exp(-abs(x)))

Prefer a library implementation for this calculation. Differentiating a hand-built expression involving max and abs at zero can introduce boundary conventions you did not intend. jax.nn.softplus is also a stable library implementation and ordinarily does not need a custom rule.

Define the reverse-mode rule

This exercise uses jax.custom_vjp to practice supplying that rule:

@jax.custom_vjp
def stable_log1pexp(x):
    return jnp.logaddexp(0.0, x)

def fwd(x):
    return jnp.logaddexp(0.0, x), x

def bwd(x, cotangent):
    return (cotangent * jax.nn.sigmoid(x),)

stable_log1pexp.defvjp(fwd, bwd)

The forward rule returns a pair: the primal output and the saved input. The backward rule receives that saved input and the output cotangent. Since this function returns a vector, its cotangent is also a vector. Multiply it elementwise by sigmoid; do not sum it or assume it is all ones. Return a one-element tuple because the primal function has one argument.

Use jax.nn.sigmoid rather than claiming that a handwritten 1 / (1 + exp(-x)) has no overflowing intermediates: its exponential can overflow for large negative x. Finite-precision implementations can still round sigmoid to zero or one in the tails.

A custom VJP defines reverse-mode differentiation. Direct jax.jvp through a custom_vjp function is not supported. Keep rule code pure and express data-dependent control flow with JAX operations.

Problem

Implement stable_log1pexp(x) with the custom VJP above:

  1. Compute elementwise softplus with a stable forward formula such as jnp.logaddexp(0.0, x) in both the primal and forward rule.
  2. Save x as the residual.
  3. Return (cotangent * jax.nn.sigmoid(x),) from the backward rule.
  4. Register both rules with defvjp.

x is a finite float32 JAX vector of shape (n,), with n >= 1. Leave it unchanged and return a JAX vector of the same shape, not a sum or gradient. Large positive entries such as 100 and 1000 must produce finite outputs near those values. At zero, the value is approximately 0.6931472 and the derivative is 0.5.

Output tests can check forward stability but cannot establish that the custom derivative rule was used. Also check a VJP with a nonuniform cotangent and compare gradients with the sigmoid formula locally.