A gradient tells you how a function changes with respect to its inputs. In neural-network training, those inputs are usually model parameters, and the gradient points toward the steepest increase in the loss. Subtracting a scaled gradient moves the parameters toward lower loss. section 2.6 section 4.1
Use gradients in JAX as explicit function transformations:
- Define a loss function that receives parameters and data, then obtain its gradients with grad(loss)(params, x, y). section 2.2 section 2.6.3
- Update each parameter by subtracting the learning-rate-scaled gradient; an exponentially decaying learning rate makes later steps smaller. section 2.6.3
- Choose value_and_grad when training needs both the loss value and its gradient. section 2.6.3
- Use argnums to select differentiated arguments, has_aux for auxiliary results, and repeated grad calls for higher-order derivatives. section 4.4
A minimal parameter update
from jax import grad def update(params, x, y, learning_rate): grads = grad(loss)(params, x, y) return [(w - learning_rate * dw, b - learning_rate * db) for (w, b), (dw, db) in zip(params, grads)] The transformed function returns gradients for each parameter. The update applies them in the opposite direction, while leaving the loss definition separate from the optimization step. section 2.6.3
Visualizing gradient descent steps along the loss landscape
Visualizes repeated steps across a loss landscape, making the opposite-gradient direction and convergence toward a minimum concrete. section 2.6
Listing 2.11 Implementing the gradient update step
from jax import grad #1
INIT_LR = 1.0 #2
DECAY_RATE = 0.95 #3
DECAY_STEPS = 5 #4
def update(params, x, y, epoch_number):
grads = grad(loss)(params, x, y) #5
lr = INIT_LR * DECAY_RATE ** (epoch_number / DECAY_STEPS) #6
return [(w - lr * dw, b - lr * db) #7
for (w, b), (dw, db) in zip(params, grads)] #7
Provides a runnable update-step listing that pairs grad with parameter-wise learning-rate updates. section 2.6.3
