A gradient update adjusts model parameters to reduce loss. Compute the loss and its gradients, then move weights in the opposite direction. In regression, a linear model y = wx + b uses mean squared error; repeated updates form the training loop. section 4.2
A practical update follows this pattern:
- Evaluate the loss and obtain gradients with value_and_grad(loss). section 7.3.2
- For data-parallel training, sum corresponding gradients across devices with lax.psum. section 7.3.2
- Apply w - learning_rate * gradient and repeat across batches and epochs. section 4.2 section 7.3.2
- Keep parameter copies synchronized: each device computes locally, receives aggregated gradients, and performs the same update. section 7.3.2
Updating a replicated classifier
An update function receives replicated parameters and per-device image batches. It computes categorical cross-entropy and gradients, aggregates each weight and bias gradient with lax.psum, applies a decaying learning rate, and returns updated parameters with per-device loss values. A complete loop reshapes a 256-example batch into eight groups of 32 before calling update. section 7.3.2
Listing 7.27 Loss and update functions
from functools import partial
INIT_LR = 1.0
DECAY_RATE = 0.95
DECAY_STEPS = 5
NUM_EPOCHS = 20
def loss(params, images, targets): #1
"""Categorical cross entropy loss function."""
logits = batched_predict(params, images)
log_preds = logits - logsumexp(logits)
return -jnp.mean(targets*log_preds)
@partial(jax.pmap,
axis_name='devices',
in_axes=(None, 0, 0, None),
out_axes=(None,0)) #2
def update(params, x, y, epoch_number):
loss_value, grads = value_and_grad(loss)(params, x, y)
grads = [(jax.lax.psum(dw, 'devices'), #3
jax.lax.psum(db, 'devices')) #3
for dw, db in grads] #3
lr = INIT_LR * DECAY_RATE ** (epoch_number / DECAY_STEPS)
return [(w - lr * dw, b - lr * db)
for (w, b), (dw, db) in zip(params, grads)], loss_value
#1 The same loss function as before
#2 Uses pmap to parallelize function and map over the data only; parameters are replicated and not mapped over
#3 Accumulates gradients on each of the devices
Inspect the JAX loss and update listing, including value_and_grad, pmap, lax.psum, learning-rate decay, and parameter arithmetic. section 7.3.2
