Study Plan

Per-Example Gradient Clipping

JAX

Medium

Compute per-example gradients of a linear model, clip each one to a maximum norm and average the clipped gradients, as one step of DP-SGD does before it adds noise.

The model has a weight vector w and a scalar bias b, stored in the dict params under the keys w and b. Example $n$ has a feature row $x_n$ and a target $y_n$ , and its loss is

\ell_n(w, b) = \tfrac{1}{2}\big(x_n \cdot w + b - y_n\big)^2

Let $g_n$ be the gradient of $\ell_n$ with respect to both parameters, and let $\lVert g_n \rVert$ be its L2 norm taken over all D + 1 entries together, the D weight entries and the bias entry. With $C$ equal to max_norm, each gradient is clipped to

\bar g_n = \frac{g_n}{\max\big(1, \lVert g_n \rVert / C\big)}

and the result is the mean of $\bar g_n$ over the N examples. Here x has shape (N, D) with rows $x_n$ and y has shape (N,).

Return a dict with two keys. The value under grads is a dict with keys w and b holding the mean clipped gradient, w as a float32 JAX array of shape (D,) and b as a float32 JAX array of shape (). The value under norms is a float32 JAX array of shape (N,) holding $\lVert g_n \rVert$ before clipping.

Examples

Input: params = {"w":[1.0,-1.0],"b":0.0}, x = [[1.0,0.0],[2.0,2.0],[0.0,3.0]], y = [0.0,1.0,1.0], max_norm = 1.0

Output: {"grads":{"w":[0.01348,-0.53845],"b":0.019182},"norms":[1.414214,3.0,12.649111]}

Explanation: The first example has error 1 and gradient ([1, 0], 1) with norm 1.414214, so it is scaled down to length 1. The other two examples have norms 3 and 12.649111 and are scaled down too, and the three clipped gradients are then averaged.

Input: params = {"w":[0.5],"b":0.5}, x = [[1.0],[-1.0],[2.0],[0.0]], y = [1.5,0.5,1.0,0.0], max_norm = 10.0

Output: {"grads":{"w":[0.25],"b":0.0},"norms":[0.707107,0.707107,1.118034,0.5]}

Input: params = {"w":[2.0,1.0],"b":-1.0}, x = [[1.0,1.0],[0.5,-1.0]], y = [2.0,-4.0], max_norm = 2.0

Output: {"grads":{"w":[0.333333,-0.666667],"b":0.666667},"norms":[0.0,4.5]}

Hint 1

jax.vmap(jax.grad(loss), in_axes=(None, 0, 0)) gives one gradient per example.

Hint 2

jnp.linalg.norm(g, axis=-1) or jnp.sqrt(jnp.sum(g ** 2, axis=...)) gives each example's norm.

Hint 3

factor[:, None] broadcasts an (N,) scale against (N, D) gradients.

Requirements

Write the loss for a single example and get the per-example gradients by applying jax.vmap to jax.grad of it, mapping over the rows of x and y while sharing params. Solutions that do not call both transforms are rejected.
An example whose gradient is exactly zero must contribute zero without producing NaN.
Leave the params dict unchanged.

Constraints

params holds w, a float32 JAX array of shape (D,), and b, a float32 JAX array of shape ().
x is a float32 JAX array of shape (N, D) with N between 1 and 8 and D between 1 and 4.
y is a float32 JAX array of shape (N,).
max_norm is a positive Python float between 0.1 and 100.

Try Similar Problems

Hessian-Vector ProductMedium Straight-Through Estimator with Custom VJPMedium Value and Grad of a LossEasy Pad and Mask to Bucket LengthMedium Traced Learning Rate ScheduleMedium Exponential Moving Average with ScanMedium Spectral Norm by Power IterationMedium Batched Attention with vmapMedium Pairwise Distances with Nested vmapMedium Global Norm Gradient ClippingMedium AdamW from ScratchMedium Mixed Precision with Loss ScalingMedium