vmap + grad: Per-Example Gradients

The combination of vmap and grad is JAX's flagship feature. It lets you compute individual gradients for each example in a batch — something that's naturally impossible in standard backpropagation (which computes the average gradient).

import jax
import jax.numpy as jnp

def loss_single(params, x, y):
    """Loss for a SINGLE example — keep it simple!"""
    pred = jnp.dot(x, params)
    return (pred - y) ** 2

# Regular gradient: gradient of average loss
def batch_loss(params, X, y):
    return jnp.mean(jax.vmap(loss_single, in_axes=(None, 0, 0))(params, X, y))

avg_grad = jax.grad(batch_loss)

# Per-example gradients: separate gradient for each example
per_example_grad = jax.vmap(jax.grad(loss_single), in_axes=(None, 0, 0))

# Test
params = jnp.array([1.0, 2.0, 3.0])
X = jax.random.normal(jax.random.PRNGKey(0), (32, 3))
y = jax.random.normal(jax.random.PRNGKey(1), (32,))

# Average gradient: shape (3,)
g_avg = avg_grad(params, X, y)
print(f"Average gradient shape: {g_avg.shape}")  # (3,)

# Per-example gradients: shape (32, 3)
g_per = per_example_grad(params, X, y)
print(f"Per-example gradient shape: {g_per.shape}")  # (32, 3)

# Verify: mean of per-example gradients ≈ average gradient
print(f"Match: {jnp.allclose(jnp.mean(g_per, axis=0), g_avg, atol=1e-5)}")

Application: Differential Privacy

Differential privacy requires clipping individual gradients before averaging. Without per-example gradients, this is computationally expensive. With vmap + grad, it's straightforward:

import jax
import jax.numpy as jnp

def loss_single(params, x, y):
    pred = jnp.dot(x, params)
    return (pred - y) ** 2

def dp_gradient(params, X, y, clip_norm=1.0, noise_scale=0.1, key=None):
    """Differentially private gradient computation."""
    # 1. Get per-example gradients
    per_ex_grads = jax.vmap(
        jax.grad(loss_single), in_axes=(None, 0, 0)
    )(params, X, y)

    # 2. Clip each gradient to the norm bound
    grad_norms = jnp.linalg.norm(per_ex_grads, axis=-1, keepdims=True)
    clip_factor = jnp.minimum(1.0, clip_norm / (grad_norms + 1e-8))
    clipped_grads = per_ex_grads * clip_factor

    # 3. Average and add noise
    avg_grad = jnp.mean(clipped_grads, axis=0)
    noise = noise_scale * clip_norm * jax.random.normal(key, avg_grad.shape)
    return avg_grad + noise / X.shape[0]

# Usage
key = jax.random.PRNGKey(42)
params = jnp.array([1.0, 2.0, 3.0])
X = jax.random.normal(jax.random.PRNGKey(0), (64, 3))
y = jax.random.normal(jax.random.PRNGKey(1), (64,))

dp_grad = dp_gradient(params, X, y, key=key)
print(f"DP gradient: {dp_grad}")