Apply one step of stochastic gradient descent with momentum to every array in a parameter pytree. Here the pytrees are nestings of Python dicts and lists whose innermost entries, called leaves, are arrays. The three inputs params, grads and velocity share one structure, and for every leaf you compute
vnew=μv+g,pnew=p−ηvnewHere p is a leaf of params, g is the matching leaf of grads, v is the matching leaf of velocity, μ is momentum and η is learning_rate. The structure can be anything from a single bare array to dicts that hold lists of dicts.
Return a tuple (new_params, new_velocity). Both are pytrees with exactly the structure of params, and every leaf is a JAX array with the shape and dtype of the matching leaf of params.
Input: params = {"w": [0.5, -1.0], "b": 0.2}, grads = {"w": [0.1, 0.4], "b": -0.5}, velocity = {"w": [0.0, 0.0], "b": 0.0}, learning_rate = 0.1, momentum = 0.9
Output: (new_params, new_velocity) = ({"b": 0.25, "w": [0.49, -1.04]}, {"b": -0.5, "w": [0.1, 0.4]})
Explanation: The velocity starts at zero, so the new velocity equals the gradient. Each parameter then moves against its gradient by a tenth of it, so w[0] becomes 0.5 minus 0.1 times 0.1, which is 0.49, and b becomes 0.2 plus 0.1 times 0.5, which is 0.25.
Input: params = {"dense": {"w": [[1.0, 2.0], [3.0, 4.0]], "b": [0.0, 0.5]}, "scale": 1.0}, grads = {"dense": {"w": [[0.5, -0.5], [1.0, 0.0]], "b": [0.2, -0.2]}, "scale": 0.3}, velocity = {"dense": {"w": [[1.0, 1.0], [-1.0, 0.0]], "b": [0.0, 2.0]}, "scale": -0.5}, learning_rate = 0.05, momentum = 0.5
Output: (new_params, new_velocity) = ({"dense": {"b": [-0.01, 0.46], "w": [[0.95, 2.0], [2.975, 4.0]]}, "scale": 0.9975}, {"dense": {"b": [0.2, 0.8], "w": [[1.0, 0.0], [0.5, 0.0]]}, "scale": 0.05})
Input: params = {"layers": [{"w": [1.0, -1.0], "b": 0.0}, {"w": [2.0], "b": 1.0}]}, grads = {"layers": [{"w": [0.2, 0.2], "b": 1.0}, {"w": [-1.0], "b": 0.5}]}, velocity = {"layers": [{"w": [0.1, -0.1], "b": 0.0}, {"w": [0.4], "b": -0.2}]}, learning_rate = 0.2, momentum = 0.8
Output: (new_params, new_velocity) = ({"layers": [{"b": -0.2, "w": [0.944, -1.024]}, {"b": 0.932, "w": [2.136]}]}, {"layers": [{"b": 1.0, "w": [0.28, 0.12]}, {"b": 0.34, "w": [-0.68]}]})
jax.tree.map(lambda p, v: ..., params, velocity) works leaf by leaf.
Call jax.tree.map twice, once for the new velocity and once for the new params.
A Python float times a bfloat16 array stays bfloat16.
Accepts: object, array, number
Accepts: object, array, number
Accepts: object, array, number
Accepts: number
Accepts: number
Apply one step of stochastic gradient descent with momentum to every array in a parameter pytree. Here the pytrees are nestings of Python dicts and lists whose innermost entries, called leaves, are arrays. The three inputs params, grads and velocity share one structure, and for every leaf you compute
vnew=μv+g,pnew=p−ηvnewHere p is a leaf of params, g is the matching leaf of grads, v is the matching leaf of velocity, μ is momentum and η is learning_rate. The structure can be anything from a single bare array to dicts that hold lists of dicts.
Return a tuple (new_params, new_velocity). Both are pytrees with exactly the structure of params, and every leaf is a JAX array with the shape and dtype of the matching leaf of params.
Input: params = {"w": [0.5, -1.0], "b": 0.2}, grads = {"w": [0.1, 0.4], "b": -0.5}, velocity = {"w": [0.0, 0.0], "b": 0.0}, learning_rate = 0.1, momentum = 0.9
Output: (new_params, new_velocity) = ({"b": 0.25, "w": [0.49, -1.04]}, {"b": -0.5, "w": [0.1, 0.4]})
Explanation: The velocity starts at zero, so the new velocity equals the gradient. Each parameter then moves against its gradient by a tenth of it, so w[0] becomes 0.5 minus 0.1 times 0.1, which is 0.49, and b becomes 0.2 plus 0.1 times 0.5, which is 0.25.
Input: params = {"dense": {"w": [[1.0, 2.0], [3.0, 4.0]], "b": [0.0, 0.5]}, "scale": 1.0}, grads = {"dense": {"w": [[0.5, -0.5], [1.0, 0.0]], "b": [0.2, -0.2]}, "scale": 0.3}, velocity = {"dense": {"w": [[1.0, 1.0], [-1.0, 0.0]], "b": [0.0, 2.0]}, "scale": -0.5}, learning_rate = 0.05, momentum = 0.5
Output: (new_params, new_velocity) = ({"dense": {"b": [-0.01, 0.46], "w": [[0.95, 2.0], [2.975, 4.0]]}, "scale": 0.9975}, {"dense": {"b": [0.2, 0.8], "w": [[1.0, 0.0], [0.5, 0.0]]}, "scale": 0.05})
Input: params = {"layers": [{"w": [1.0, -1.0], "b": 0.0}, {"w": [2.0], "b": 1.0}]}, grads = {"layers": [{"w": [0.2, 0.2], "b": 1.0}, {"w": [-1.0], "b": 0.5}]}, velocity = {"layers": [{"w": [0.1, -0.1], "b": 0.0}, {"w": [0.4], "b": -0.2}]}, learning_rate = 0.2, momentum = 0.8
Output: (new_params, new_velocity) = ({"layers": [{"b": -0.2, "w": [0.944, -1.024]}, {"b": 0.932, "w": [2.136]}]}, {"layers": [{"b": 1.0, "w": [0.28, 0.12]}, {"b": 0.34, "w": [-0.68]}]})
jax.tree.map(lambda p, v: ..., params, velocity) works leaf by leaf.
Call jax.tree.map twice, once for the new velocity and once for the new params.
A Python float times a bfloat16 array stays bfloat16.
Accepts: object, array, number
Accepts: object, array, number
Accepts: object, array, number
Accepts: number
Accepts: number
