TopolearnJAX's best-known feature computes gradients; this lesson explains what they are and why training needs them.
Think of a function as a machine: numbers go in, a number comes out. Its derivative answers a simple question: if I nudge the input up a tiny bit, how much does the output change, and in which direction? For f(x) = x ** 2 at x = 3, a nudge of 0.001 raises the output by about 0.006, so the derivative there is 6 (the formula is 2 * x).
Most interesting functions take many inputs. The gradient is the list of derivatives, one per input: for each input, how strongly and in which direction it pushes the output. If the output is height on a hillside and the inputs are your position, the gradient points straight uphill.
The biggest user of gradients is machine learning: instead of writing rules by hand, you let a program learn them from examples.
The program is a model, a formula with adjustable numbers called weights. A loss function scores how wrong the model's answers are on the examples: one number, zero for perfect answers. The question "how should I change each weight to lower the loss?" is answered by the gradient of the loss with respect to the weights.
Training repeats one move: compute that gradient, then step every weight a little in the opposite direction, downhill. This is gradient descent. The size of each step is the learning rate: too big and you overshoot the valley, too small and you crawl.
With millions of weights, deriving the gradient formula by hand is out of the question, and estimating it by nudging each weight in turn would be far too slow. JAX's jax.grad (next lesson) computes the exact gradient automatically, at about the cost of running the function a couple of times. The same tool works for any numeric code, from physics simulations to finance, not only machine learning.
Example
# Estimate a derivative by nudging the input, then use it to train one weight.
def f(x):
return x ** 2
h = 0.001
print((f(3 + h) - f(3)) / h) # about 6.0, the derivative of x**2 at 3
examples = [(1, 3), (2, 6), (3, 9)] # (input, right answer): the truth is y = 3x
def loss(w): # how wrong the model guess = w * x is
return sum((w * x - y) ** 2 for x, y in examples)
w, learning_rate = 0.0, 0.01
for step in range(20):
grad = (loss(w + h) - loss(w)) / h # derivative of the loss w.r.t. w
w -= learning_rate * grad # step downhill
print(round(w, 2)) # 3.0: the model learned y = 3xRemember
- A derivative says how fast the output changes when an input changes; a gradient lists one per input
- Training = gradient descent: step each weight against the gradient of the loss
- The learning rate is the step size; jax.grad computes gradients exactly and automatically
Watch out for
- Choosing a learning rate so large that the loss jumps around or grows instead of shrinking.