Skip to content

← JAX Autodiff step 14 of 25

Medium Primitives

Straight-Through Estimator

Jump to the task

Keep a discrete value and choose a training derivative

jnp.sign(x) returns -1 for negative values, 0 at zero, and 1 for positive values. Away from zero, a small change in x does not change this output, so its ordinary derivative is zero. At zero, the function jumps and has no ordinary derivative. JAX uses a zero derivative for sign.

A straight-through estimator deliberately substitutes a different derivative. Here we keep sign's output but ask autodiff to treat the operation like the identity. This gives a training signal through the quantization step. It is a chosen surrogate, not the true derivative of sign or a guarantee that optimization will succeed.

Build a zero value with a nonzero derivative

stop_gradient preserves a value and blocks its derivative. For finite x, the expression x - stop_gradient(x) is zero in value, but its derivative with respect to x is one. Add the stopped sign value:

import jax
import jax.numpy as jnp

def ste_quantize(x):
    zero = x - jax.lax.stop_gradient(x)
    return zero + jax.lax.stop_gradient(jnp.sign(x))

x = jnp.array([-1.5, 0.0, 0.5, 2.0])
ste_quantize(x)  # [-1, 0, 1, 1]
jax.grad(lambda z: ste_quantize(z).sum())(x)  # [1, 1, 1, 1]

For a vector input, the Jacobian is the identity matrix, not a matrix of ones. Each output passes its incoming gradient to the matching input. For example, differentiating sum(weights * ste_quantize(x)) returns weights. A JVP similarly passes its input direction through unchanged. At zero we still choose a surrogate derivative of one.

Keep the floating-point arithmetic reliable

The alternative x + stop_gradient(sign(x) - x) looks equivalent in real arithmetic, but can lose the sign value through rounding. With float32 x = 1e20, it returns zero instead of one. Subtracting x from its stopped copy first avoids that cancellation for finite inputs.

Other straight-through estimators use different surrogate derivatives. Custom JVP or VJP rules are also valid ways to define such behavior; the stop-gradient construction is sufficient for this exercise.

Problem

Implement ste_quantize(x) for a one-dimensional real floating-point JAX array of finite values. Return an array with the same shape and dtype, containing jnp.sign(x), while giving it the identity derivative under autodiff. Zero must produce zero, so this is not strictly a two-valued binary quantizer. NaN and infinity are outside this exercise.

Use the zero-valued expression above to preserve the correct sign even for large finite inputs. The output-only tests verify the values, including this numerical edge case. They cannot distinguish a plain sign implementation from one with the required surrogate derivative; use the gradient example above to check that part of your solution.