When Automatic Differentiation Goes Wrong#
Automatic differentiation (autodiff) computes exact derivatives of the program as written. However, this doesn’t mean the derivatives are always accurate in the numerical sense. This notebook explores common pitfalls where autodiff produces unreliable gradients.
Key insight: Autodiff differentiates the computation, not the mathematical function. If the computation is numerically unstable, so are the gradients.
Topics Covered#
Catastrophic cancellation
Overflow and underflow
Ill-conditioned functions
Numerical comparison with finite differences
Strategies for improving accuracy
import jax.numpy as jnp
import jax
jax.config.update("jax_enable_x64", True)
import numpy as np
import matplotlib.pyplot as plt
1. Catastrophic Cancellation#
When subtracting two nearly equal numbers, significant digits are lost. This affects both the function value AND its gradient.
Example: Derivative of \((1 - \cos(x))/x^2\) near \(x=0\)#
Mathematically, \(\lim_{x \to 0} \frac{1 - \cos(x)}{x^2} = \frac{1}{2}\)
But numerically, \(1 - \cos(x) \approx 0\) when \(x\) is small, causing cancellation.
# Naive implementation - suffers from cancellation
def f_naive(x):
return (1 - jnp.cos(x)) / x**2
# Stable implementation using Taylor series: 1 - cos(x) = x^2/2 - x^4/24 + ...
def f_stable(x):
# Use identity: 1 - cos(x) = 2*sin^2(x/2)
return 2 * (jnp.sin(x/2) / x)**2
# Test at various x values approaching 0
x_values = jnp.array([1e-1, 1e-2, 1e-4, 1e-6, 1e-8, 1e-10, 1e-12])
print("Function values (should approach 0.5):")
print(f"{'x':<12} {'Naive':<20} {'Stable':<20}")
print("-" * 52)
for x in x_values:
naive_val = float(f_naive(x))
stable_val = float(f_stable(x))
print(f"{float(x):<12.0e} {naive_val:<20.10f} {stable_val:<20.10f}")
WARNING:2026-01-10 20:27:07,228:jax._src.xla_bridge:852: An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
Function values (should approach 0.5):
x Naive Stable
----------------------------------------------------
1e-01 0.4995834722 0.4995834722
1e-02 0.4999958333 0.4999958333
1e-04 0.4999999970 0.4999999996
1e-06 0.5000444503 0.5000000000
1e-08 0.0000000000 0.5000000000
1e-10 0.0000000000 0.5000000000
1e-12 0.0000000000 0.5000000000
# Now check the GRADIENTS
grad_naive = jax.grad(f_naive)
grad_stable = jax.grad(f_stable)
# True derivative at x=0 is 0 (by symmetry and L'Hopital)
print("\nGradient values (should approach 0):")
print(f"{'x':<12} {'Naive grad':<20} {'Stable grad':<20}")
print("-" * 52)
for x in x_values:
naive_grad = float(grad_naive(x))
stable_grad = float(grad_stable(x))
print(f"{float(x):<12.0e} {naive_grad:<20.6e} {stable_grad:<20.6e}")
print("\nNotice how the naive gradient becomes garbage for small x!")
Gradient values (should approach 0):
x Naive grad Stable grad
----------------------------------------------------
1e-01 -8.327779e-03 -8.327779e-03
1e-02 -8.333278e-04 -8.333278e-04
1e-04 4.410805e-05 -8.333336e-06
1e-06 -8.890058e+01 -8.323696e-08
1e-08 1.000000e+08 1.490116e-08
1e-10 1.000000e+10 1.907349e-06
1e-12 1.000000e+12 -1.220703e-04
Notice how the naive gradient becomes garbage for small x!
2. Overflow and Underflow#
When intermediate values exceed floating-point range, we get inf or 0, and gradients become nan.
Example: Softmax with large inputs#
# Naive softmax - overflows for large inputs
def softmax_naive(x):
exp_x = jnp.exp(x)
return exp_x / jnp.sum(exp_x)
# Stable softmax - subtract max first
def softmax_stable(x):
x_shifted = x - jnp.max(x)
exp_x = jnp.exp(x_shifted)
return exp_x / jnp.sum(exp_x)
# Test with large values
x_large = jnp.array([1000.0, 1001.0, 1002.0])
print("Softmax with large inputs [1000, 1001, 1002]:")
print(f"Naive: {softmax_naive(x_large)}")
print(f"Stable: {softmax_stable(x_large)}")
Softmax with large inputs [1000, 1001, 1002]:
Naive: [nan nan nan]
Stable: [0.09003057 0.24472847 0.66524096]
# Check gradients
def loss_naive(x):
return jnp.sum(softmax_naive(x) * jnp.array([1.0, 0.0, 0.0])) # Cross-entropy with [1,0,0]
def loss_stable(x):
return jnp.sum(softmax_stable(x) * jnp.array([1.0, 0.0, 0.0]))
print("\nGradients:")
print(f"Naive gradient: {jax.grad(loss_naive)(x_large)}")
print(f"Stable gradient: {jax.grad(loss_stable)(x_large)}")
print("\nNaive gives NaN gradients due to 0/0 or inf/inf!")
Gradients:
Naive gradient: [nan nan nan]
Stable gradient: [ 0.08192507 -0.02203304 -0.05989202]
Naive gives NaN gradients due to 0/0 or inf/inf!
Example: Log-sum-exp#
# Naive log-sum-exp
def logsumexp_naive(x):
return jnp.log(jnp.sum(jnp.exp(x)))
# Stable log-sum-exp
def logsumexp_stable(x):
x_max = jnp.max(x)
return x_max + jnp.log(jnp.sum(jnp.exp(x - x_max)))
# Test
x_large = jnp.array([1000.0, 1001.0, 999.0])
print("Log-sum-exp with large inputs:")
print(f"Naive: {logsumexp_naive(x_large)}")
print(f"Stable: {logsumexp_stable(x_large)}")
print(f"JAX built-in: {jax.scipy.special.logsumexp(x_large)}")
print("\nGradients:")
print(f"Naive: {jax.grad(logsumexp_naive)(x_large)}")
print(f"Stable: {jax.grad(logsumexp_stable)(x_large)}")
Log-sum-exp with large inputs:
Naive: inf
Stable: 1001.4076059644444
JAX built-in: 1001.4076059644444
Gradients:
Naive: [nan nan nan]
Stable: [0.24472847 0.66524096 0.09003057]
3. Ill-Conditioned Functions#
Some functions have gradients that are extremely sensitive to small perturbations, even when the function itself looks fine.
Example: Condition number and gradient amplification#
Ill-conditioned problems amplify errors. The condition number κ tells us how much input errors get magnified in the output.
# Wilkinson's polynomial: (x-1)(x-2)...(x-20)
# Famous for being ill-conditioned with respect to coefficient perturbations
# The gradient at a root can be computed analytically: p'(k) = prod_{i≠k}(k-i)
import math
def wilkinson_poly(x):
result = jnp.ones_like(x)
for i in range(1, 21):
result = result * (x - i)
return result
# At x = 10 (exactly), the analytical gradient is:
# p'(10) = 9! * (-1)^10 * 10! = 362880 * 3628800 = 1,316,818,944,000
analytical_grad_at_10 = float(math.factorial(9) * math.factorial(10))
print(f"Analytical gradient at x=10: {analytical_grad_at_10:.6e}")
# Test autodiff at x = 10 exactly
x_exact = jnp.array(10.0)
print(f"\nAt x = 10 exactly:")
print(f" Value: {wilkinson_poly(x_exact):.6e}")
print(f" Autodiff gradient: {jax.grad(wilkinson_poly)(x_exact):.6e}")
print(f" Analytical gradient: {analytical_grad_at_10:.6e}")
# The issue with ill-conditioning appears when evaluating NEAR a root
# Tiny perturbations in x cause huge changes in the function value
x_near = 10.0 + 1e-10
print(f"\nAt x = 10 + 1e-10:")
print(f" Value: {wilkinson_poly(jnp.array(x_near)):.6e}")
print(f" The value changed dramatically from 0 due to massive gradient!")
# The gradient magnitude itself is a sign of ill-conditioning
print(f"\nThe gradient of ~10^12 means tiny errors in x get amplified 10^12 times!")
Analytical gradient at x=10: 1.316819e+12
At x = 10 exactly:
Value: 0.000000e+00
Autodiff gradient: 1.316819e+12
Analytical gradient: 1.316819e+12
At x = 10 + 1e-10:
Value: 1.316819e+02
The value changed dramatically from 0 due to massive gradient!
The gradient of ~10^12 means tiny errors in x get amplified 10^12 times!
Example: Matrix inverse gradient#
# Gradient of trace(A^{-1}) with respect to A
# For ill-conditioned A, the gradient is very sensitive
def trace_inv(A):
return jnp.trace(jnp.linalg.inv(A))
# Well-conditioned matrix
A_good = jnp.array([[2.0, 0.1], [0.1, 2.0]])
print(f"Well-conditioned matrix (cond = {jnp.linalg.cond(A_good):.1f}):")
print(f"trace(A^-1) = {trace_inv(A_good):.6f}")
print(f"Gradient norm = {jnp.linalg.norm(jax.grad(trace_inv)(A_good)):.6f}")
# Ill-conditioned matrix
A_bad = jnp.array([[1.0, 0.9999], [0.9999, 1.0]])
print(f"\nIll-conditioned matrix (cond = {jnp.linalg.cond(A_bad):.1f}):")
print(f"trace(A^-1) = {trace_inv(A_bad):.6f}")
print(f"Gradient norm = {jnp.linalg.norm(jax.grad(trace_inv)(A_bad)):.6f}")
print("\nThe gradient is huge even though the function value looks reasonable!")
Well-conditioned matrix (cond = 1.1):
trace(A^-1) = 1.002506
Gradient norm = 0.357984
Ill-conditioned matrix (cond = 19999.0):
trace(A^-1) = 10000.500025
Gradient norm = 100000000.000072
The gradient is huge even though the function value looks reasonable!
4. Comparing with Finite Differences#
A useful sanity check is to compare autodiff gradients with finite difference approximations.
def finite_diff_grad(f, x, eps=1e-6):
"""Central finite difference gradient."""
grad = jnp.zeros_like(x)
for i in range(len(x)):
x_plus = x.at[i].add(eps)
x_minus = x.at[i].add(-eps)
grad = grad.at[i].set((f(x_plus) - f(x_minus)) / (2 * eps))
return grad
def check_gradient(f, x, name="function"):
"""Compare autodiff and finite difference gradients."""
ad_grad = jax.grad(f)(x)
fd_grad = finite_diff_grad(f, x)
rel_error = jnp.linalg.norm(ad_grad - fd_grad) / (jnp.linalg.norm(ad_grad) + 1e-10)
print(f"\n{name}:")
print(f" Autodiff gradient: {ad_grad}")
print(f" Finite diff gradient: {fd_grad}")
print(f" Relative error: {rel_error:.2e}")
if rel_error > 1e-4:
print(" WARNING: Large discrepancy!")
else:
print(" OK: Gradients match.")
return rel_error
# Test various functions
x = jnp.array([1.0, 2.0, 3.0])
# Simple function - should work fine
check_gradient(lambda x: jnp.sum(x**2), x, "sum of squares")
# Stable softmax - should work fine
check_gradient(lambda x: jnp.sum(softmax_stable(x)), x, "stable softmax")
# Function with potential issues at certain points
check_gradient(lambda x: jnp.sum(jnp.sqrt(jnp.abs(x) + 1e-10)), x, "sqrt(|x|)")
sum of squares:
Autodiff gradient: [2. 4. 6.]
Finite diff gradient: [2. 4. 6.]
Relative error: 1.26e-10
OK: Gradients match.
stable softmax:
Autodiff gradient: [-1.50252347e-17 -4.08428226e-17 5.58680573e-17]
Finite diff gradient: [0.00000000e+00 5.55111512e-11 0.00000000e+00]
Relative error: 5.55e-01
WARNING: Large discrepancy!
sqrt(|x|):
Autodiff gradient: [0.5 0.35355339 0.28867513]
Finite diff gradient: [0.5 0.35355339 0.28867513]
Relative error: 4.90e-10
OK: Gradients match.
Array(4.90223736e-10, dtype=float64)
5. Strategies for Improving Accuracy#
Strategy 1: Use mathematically equivalent but numerically stable forms#
# Example: log(1 + exp(x)) - the "softplus" function
def softplus_naive(x):
return jnp.log(1 + jnp.exp(x))
def softplus_stable(x):
# For large x: log(1 + exp(x)) ≈ x
# For small x: log(1 + exp(x)) ≈ log(1 + x) when x << 1
return jnp.where(x > 20, x, jnp.log1p(jnp.exp(jnp.minimum(x, 20))))
# Test
x_test = jnp.array([0.0, 10.0, 100.0, 1000.0])
print("Softplus function:")
print(f"x = {x_test}")
print(f"Naive: {softplus_naive(x_test)}")
print(f"Stable: {softplus_stable(x_test)}")
print("\nGradients at x=1000:")
print(f"Naive: {jax.grad(lambda x: softplus_naive(x))(1000.0)}")
print(f"Stable: {jax.grad(lambda x: softplus_stable(x))(1000.0)}")
print("(Should be ~1.0 for large x)")
Softplus function:
x = [ 0. 10. 100. 1000.]
Naive: [ 0.69314718 10.0000454 100. inf]
Stable: [6.93147181e-01 1.00000454e+01 1.00000000e+02 1.00000000e+03]
Gradients at x=1000:
Naive: nan
Stable: 1.0
(Should be ~1.0 for large x)
Strategy 2: Use higher precision when needed#
# We already enabled float64 with jax.config.update("jax_enable_x64", True)
# Float64 has ~16 significant digits vs ~7 for float32
# Example: sum of many small numbers (accumulation error)
def accumulated_sum(n):
"""Sum 1/n, n times - should equal 1.0"""
return jnp.sum(jnp.ones(n) / n)
n_values = [100, 10000, 1000000]
print("Accumulated sum error (should be 1.0):")
print(f"{'n':<12} {'Float32':<20} {'Float64':<20}")
print("-" * 52)
for n in n_values:
val32 = float(accumulated_sum(jnp.array(n, dtype=jnp.int32)).astype(jnp.float32))
ones_32 = jnp.ones(n, dtype=jnp.float32)
val32 = float(jnp.sum(ones_32 / n))
ones_64 = jnp.ones(n, dtype=jnp.float64)
val64 = float(jnp.sum(ones_64 / n))
print(f"{n:<12} {val32:<20.15f} {val64:<20.15f}")
# Example: computing variance with cancellation
print("\nVariance calculation (catastrophic cancellation):")
print("Data: 1e8 + [0, 1, 2, 3, 4] - variance should be 2.0")
def naive_variance(x):
"""Textbook formula: E[X^2] - E[X]^2 - numerically unstable"""
n = len(x)
mean_sq = jnp.sum(x**2) / n
sq_mean = (jnp.sum(x) / n)**2
return mean_sq - sq_mean
# Data with large offset but small variance
data_offset = jnp.array([0.0, 1.0, 2.0, 3.0, 4.0]) # true variance = 2.0
offset = 1e8
data32 = (data_offset + offset).astype(jnp.float32)
data64 = (data_offset + offset).astype(jnp.float64)
print(f"Float32 variance: {naive_variance(data32):.6f}")
print(f"Float64 variance: {naive_variance(data64):.6f}")
print(f"True variance: 2.0")
print("\nFloat32 loses precision due to squaring large numbers then subtracting!")
Accumulated sum error (should be 1.0):
n Float32 Float64
----------------------------------------------------
100 0.999999940395355 1.000000000000000
10000 1.000000476837158 0.999999999999999
1000000 1.000000238418579 1.000000000000000
Variance calculation (catastrophic cancellation):
Data: 1e8 + [0, 1, 2, 3, 4] - variance should be 2.0
Float32 variance: 0.000000
Float64 variance: 2.000000
True variance: 2.0
Float32 loses precision due to squaring large numbers then subtracting!
Strategy 3: Use JAX’s built-in numerically stable functions#
# JAX provides many stable implementations
print("JAX stable functions:")
print(f"jax.nn.softmax: {jax.nn.softmax(jnp.array([1000., 1001., 1002.]))}")
print(f"jax.nn.log_softmax: {jax.nn.log_softmax(jnp.array([1000., 1001., 1002.]))}")
print(f"jax.scipy.special.logsumexp: {jax.scipy.special.logsumexp(jnp.array([1000., 1001., 1002.]))}")
print(f"jax.nn.softplus: {jax.nn.softplus(jnp.array([0., 100., 1000.]))}")
# All of these have well-behaved gradients
print("\nTheir gradients work correctly even for extreme inputs.")
JAX stable functions:
jax.nn.softmax: [0.09003057 0.24472847 0.66524096]
jax.nn.log_softmax: [-2.40760596 -1.40760596 -0.40760596]
jax.scipy.special.logsumexp: 1002.4076059644444
jax.nn.softplus: [6.93147181e-01 1.00000000e+02 1.00000000e+03]
Their gradients work correctly even for extreme inputs.
Summary#
Issue |
Example |
Solution |
|---|---|---|
Catastrophic cancellation |
|
Use equivalent stable form: |
Overflow |
|
Subtract max before exp |
Underflow |
|
Work in log space |
Ill-conditioning |
Matrix inverse of nearly singular A |
Regularization, pseudoinverse |
Loss of precision |
Float32 for sensitive calculations |
Use Float64 |
Golden rules:
Always verify gradients with finite differences for new functions
Use JAX’s built-in stable implementations when available
Work in log-space for products of probabilities
Be suspicious of gradients that are
nan,inf, or unexpectedly large