Machine learning (ML) and deep learning models learn by adjusting their parameters (weights and biases) so that a loss function gets smaller. The algorithm that decides how to adjust those parameters at every step is the optimizer.
Picture the loss as a landscape of hills and valleys. Every set of parameter values is a point on it, and training means walking downhill to the lowest point you can find. The gradient tells you which way is downhill right where you stand. The optimizer decides how big a step to take and whether to remember where you have been.
This article walks through 10 famous optimizers in the order they build on each other. For each one you get the intuition, the exact update rule, what every symbol means, a picture, and a Keras snippet.
Notation used throughout
Every formula below uses the same symbols:
| Symbol | Meaning |
|---|---|
| model parameters at step | |
| the loss function | |
| gradient of the loss at step , computed on a mini-batch | |
| learning rate (step size) | |
| momentum coefficient | |
| decay rate of a running average (RMSProp, Adadelta) | |
| decay rates of Adam's first and second moment estimates | |
| a tiny constant that prevents division by zero |
The mini-batch gradient is the average gradient over the examples in the batch:
Squares, square roots and divisions of vectors such as or are element-wise. Every parameter gets its own value, which is what lets adaptive optimizers give each parameter its own learning rate.
All the trajectory pictures in this post use the same test surface: a long, narrow valley . It is 20 times steeper across the valley than along it, which is exactly the shape that makes plain gradient descent struggle. The ○ marks the start and ★ marks the minimum.
1 - Stochastic Gradient Descent (SGD)
Intuition: look at the slope where you are standing and take a step downhill.
- points uphill, so subtracting it moves you downhill.
- scales the step. It is the single most important hyperparameter in deep learning.
Classic gradient descent computes the gradient on the whole dataset before every step. SGD estimates it from one example or, in practice, a small random mini-batch. Each estimate is noisy, but it is so much cheaper that you can take thousands of steps in the time full-batch descent takes one. The noise can even help the optimizer escape poor regions of the landscape.
The weakness of SGD is that one learning rate must suit every direction at once. In a narrow valley, a step large enough to make progress along the valley floor is too large for the steep walls, so SGD zig-zags. A slightly larger step makes the zig-zag grow until it diverges.
Keras:
keras.optimizers.SGD(learning_rate=0.01) # Keras default: 0.01
2 - SGD with Gradient Clipping
Intuition: SGD, plus a speed limit on how big a single step can be.
Some loss surfaces, especially those of recurrent networks, contain cliffs: regions where the loss changes extremely sharply. The gradient there is huge, and a single SGD step can fling the parameters far away and undo many steps of progress. This is the exploding gradient problem. Clipping caps the gradient before it is used.
Clip by norm rescales the whole gradient vector if its length exceeds a threshold :
Clip by value caps every component separately:
The update is then ordinary SGD on the clipped gradient:
Norm clipping shrinks the vector but keeps its direction. Value clipping can change the direction, because large components get cut more than small ones. Small gradients pass through either kind unchanged.
In the right panel, the loss drops sharply at the cliff. Plain SGD reaches the edge and takes one step of about 11 units, overshooting the minimum completely. With clipping, the same optimizer simply walks down the cliff.
Keras: clipping is an argument available on every Keras optimizer.
keras.optimizers.SGD(learning_rate=0.01, global_clipnorm=1.0) # norm of ALL gradients together (the formula above)
keras.optimizers.SGD(learning_rate=0.01, clipnorm=1.0) # norm of each weight tensor separately
keras.optimizers.SGD(learning_rate=0.01, clipvalue=0.5) # each component to [-0.5, 0.5]
Clipping is standard practice for RNNs, LSTMs and Transformers, typically with a threshold around 1.0.
3 - Momentum
Intuition: a heavy ball rolling downhill. It remembers its velocity, so it speeds up in directions where the gradient stays consistent and cancels out back-and-forth jitter.
- is the velocity: a running, decaying sum of past gradient steps.
- (typically 0.9) controls how much velocity is kept from one step to the next.
Across the narrow valley the gradient flips sign every step, so those contributions cancel out in . Along the valley the gradient always points the same way, so the contributions add up. With a steady gradient, the velocity grows to , which is 10 times the plain SGD step when .
Momentum was introduced by Polyak (1964) as the "heavy ball" method. It still underpins most modern optimizers, including Adam.
Keras:
keras.optimizers.SGD(learning_rate=0.01, momentum=0.9)
4 - Nesterov Momentum (Nesterov Accelerated Gradient)
Intuition: momentum that looks before it leaps. Since you are about to move by the momentum term anyway, measure the slope at the point you are about to reach rather than where you stand now.
The only difference from classical momentum is where the gradient is evaluated. Classical momentum uses . Nesterov uses the look-ahead point . If momentum is about to carry you up the opposite wall of the valley, the look-ahead gradient notices and brakes before you overshoot.
The method comes from Yurii Nesterov (1983) and was popularised for deep learning by Sutskever et al. (2013). It costs the same as classical momentum: one gradient per step.
Keras:
keras.optimizers.SGD(learning_rate=0.01, momentum=0.9, nesterov=True)
Keras implements an equivalent rearranged form, so it never has to evaluate the gradient at a separate look-ahead point.
5 - Adagrad (Adaptive Gradient)
Intuition: give each parameter its own learning rate. Parameters that have already received big gradients get smaller steps, and rarely-updated parameters get bigger ones.
- is the running sum of squared gradients, kept separately for every parameter.
- Dividing by normalises each parameter's step by its own history.
This is excellent for sparse features, such as rare words in NLP or rare items in recommender systems. A feature that appears rarely accumulates little and keeps a large learning rate, so it can still learn when it does show up.
The catch is that only ever grows. The effective learning rate shrinks roughly like , and on long training runs it can become so small that learning effectively stops. The next two optimizers were invented to fix exactly this.
Adagrad was introduced by Duchi, Hazan and Singer (2011).
Keras:
keras.optimizers.Adagrad(learning_rate=0.01, initial_accumulator_value=0.1)
# Keras default learning_rate is 0.001; starting G at 0.1 avoids huge first steps
6 - Adadelta
Intuition: fix Adagrad's shrinking learning rate by averaging only recent gradients, and get rid of the learning rate altogether.
Adadelta keeps two exponentially decaying averages for every parameter: one of squared gradients and one of squared updates.
- The denominator is the RMS of recent gradients. Using a decaying average with around 0.95, instead of Adagrad's ever-growing sum, means the step size no longer shrinks forever.
- The numerator is the RMS of recent updates. Its author, Zeiler, argued that this makes the units consistent: the update ends up with the same units as the parameters, so no hand-tuned is needed.
- The numerator uses the update average from step , because the current update hasn't been computed yet.
There is a hidden knob, though. Because starts at zero, the first steps have size proportional to , and they grow only slowly from there. In practice quietly plays the role of the learning rate.
Adadelta was introduced by Zeiler (2012).
Keras:
keras.optimizers.Adadelta(learning_rate=1.0, rho=0.95)
# Keras multiplies the update by learning_rate (default 0.001!).
# Set learning_rate=1.0 to get the original algorithm from the paper.
7 - RMSProp (Root Mean Square Propagation)
Intuition: Adagrad with a short memory. Divide each parameter's step by the root mean square of its recent gradients.
- is an exponential moving average of squared gradients. With it effectively remembers about the last 10 steps.
- In steep directions (large gradients) the step is divided by a large number. In flat directions it is divided by a small number. That evens out the zig-zag in narrow valleys.
Because old gradients fade away, the effective learning rate levels off instead of decaying to zero (see the Adagrad figure above). The optimizer can also re-adapt when the landscape changes, which is why RMSProp works well for non-stationary problems such as reinforcement learning and RNNs.
RMSProp is unusual: it was never published in a paper. Geoffrey Hinton proposed it in Lecture 6e of his 2012 Coursera course Neural Networks for Machine Learning. RMSProp and Adadelta were developed independently at about the same time and share the same first equation. The difference is that RMSProp keeps an explicit learning rate .
Keras:
keras.optimizers.RMSprop(learning_rate=0.001, rho=0.9)
8 - Adam (Adaptive Moment Estimation)
Intuition: momentum + RMSProp. Keep a running average of the gradient (direction, like momentum) and of the squared gradient (scale, like RMSProp), then correct both for starting at zero.
- is the first moment: a moving average of gradients, playing the role of momentum.
- is the second moment: a moving average of squared gradients, playing the role of RMSProp's scaling.
- and are bias-corrected versions of those averages.
Why bias correction matters. Both averages start at 0. With , after the first step , which is ten times too small. Dividing by undoes this exactly. The correction matters most in the first few dozen steps and fades to nothing as . The effect is even stronger for , because takes thousands of steps to warm up.
Adam was introduced by Kingma and Ba (2014). Its defaults ( , , ) work reasonably well across a wide range of problems. That makes it the most common default optimizer, although the learning rate still usually needs tuning for best results.
Keras:
keras.optimizers.Adam(learning_rate=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-7)
9 - Adamax
Intuition: Adam, but scale each step by the largest recent gradient magnitude instead of the root-mean-square.
Adam's is a decaying L2 norm of past gradients. Generalise it to an Lp norm and let , and the sum collapses into a simple running maximum , the infinity norm. That gives two nice properties:
- needs no bias correction, because a max of non-negative values isn't dragged toward zero by its starting value the way an average is.
- After a large gradient spike, stays high and decays slowly, so the following steps stay cautious. A single outlier also can't push the denominator toward zero.
Adamax is proposed in the same paper as Adam (Kingma and Ba, 2014, Section 7.1), with a suggested learning rate of 0.002. It sometimes does well on problems with sparse or embedding-heavy parameters. In practice Adam is used far more often.
Keras:
keras.optimizers.Adamax(learning_rate=0.002) # Keras default: 0.001
10 - SMORMS3
Intuition: RMSProp that measures how noisy each parameter's gradient is, and slows down when the gradient is mostly noise.
SMORMS3 stands for "Squared Mean Over Root Mean Squared Cubed". It was proposed by Simon Funk in an April 2015 blog post titled "RMSprop loses to SMORMS3 – Beware the Epsilon!". For every parameter it keeps a moving average of the gradient , of the squared gradient , and a memory length , which starts at 1.
Here is how to read it:
- is a signal-to-noise ratio between 0 and 1. It is close to 1 when the gradient keeps pointing the same way (all signal) and close to 0 when it keeps flipping sign (mostly noise).
- The step is capped at . With consistent gradients you get the full learning rate. With noisy gradients the step shrinks automatically. Multiplying by and dividing by is where the name comes from.
- The memory length adapts. When gradients are noisy ( ), grows by about 1 per step, so the averages cover a longer window and smooth out more noise. When gradients are consistent ( ), falls back toward 1, so the optimizer reacts quickly to changes.
SMORMS3 never became mainstream and isn't built into Keras, PyTorch or TensorFlow. It is still a neat idea: an optimizer that tunes its own averaging window per parameter.
Keras: here is a small custom Keras 3 optimizer that implements the equations above.
import keras
from keras import ops
class SMORMS3(keras.optimizers.Optimizer):
"""SMORMS3 (Simon Funk, 2015)."""
def __init__(self, learning_rate=0.001, epsilon=1e-16, name="smorms3", **kwargs):
super().__init__(learning_rate=learning_rate, name=name, **kwargs)
self.epsilon = epsilon
def build(self, var_list):
if self.built:
return
super().build(var_list)
self._m = [self.add_variable_from_reference(v, "m") for v in var_list]
self._v = [self.add_variable_from_reference(v, "v") for v in var_list]
self._mem = [self.add_variable_from_reference(v, "mem", initializer="ones")
for v in var_list]
def update_step(self, gradient, variable, learning_rate):
lr = ops.cast(learning_rate, variable.dtype)
g = ops.cast(gradient, variable.dtype)
i = self._get_variable_index(variable)
m, v, mem = self._m[i], self._v[i], self._mem[i]
r = 1.0 / (1.0 + mem)
self.assign(m, (1.0 - r) * m + r * g)
self.assign(v, (1.0 - r) * v + r * ops.square(g))
x = ops.square(m) / (v + self.epsilon) # signal-to-noise ratio
self.assign_sub(variable, g * ops.minimum(lr, x) / (ops.sqrt(v) + self.epsilon))
self.assign(mem, 1.0 + mem * (1.0 - x))
def get_config(self):
config = super().get_config()
config.update({"epsilon": self.epsilon})
return config
model.compile(optimizer=SMORMS3(learning_rate=0.001), loss="sparse_categorical_crossentropy")
All 10 side by side
Same valley, same starting point, 60 steps each. Each optimizer has a reasonable learning rate for this surface, and the grey line is plain SGD for reference.
And here is an animated race between four of them:
A 2-D valley only shows how each optimizer moves. To see them on a real (if small) learning problem, here is a 2-layer neural network trained on the classic two-spirals dataset with mini-batches of 32. All 10 optimizers were implemented from scratch with the formulas above.
Take this as an illustration, not a benchmark. The learning rates were only lightly tuned, and the ranking changes from problem to problem. A large study by Schmidt, Schneider and Hennig (2021) found that no single optimizer wins across all tasks.
Pros and cons
The last column counts how many extra values the optimizer stores for every model parameter, which drives its memory cost.
| Optimizer | Pros | Cons | Extra state per parameter |
|---|---|---|---|
| SGD | Simple and cheap. With a good learning-rate schedule it often generalises very well. | Slow in narrow valleys. Very sensitive to the learning rate. | 0 |
| SGD + clipping | Prevents exploding gradients. Makes RNN and Transformer training stable. | The threshold is one more hyperparameter. It can hide real problems such as a learning rate that is too high or bad initialisation. | 0 |
| Momentum | Much faster in ravines. Smooths out noisy gradients. | Can overshoot. Adds the momentum hyperparameter γ. | 1 |
| Nesterov | Overshoots less than classical momentum and often converges a bit faster. Same cost. | Still needs a well-tuned learning rate and schedule. | 1 |
| Adagrad | Per-parameter learning rates. Excellent for sparse features. | The learning rate decays toward zero, so it can stop learning on long runs. | 1 |
| Adadelta | No learning rate to tune in the original form. Fixes Adagrad's decay. | Slow start, and ε effectively becomes the learning rate. Rarely used today. | 2 |
| RMSProp | Per-parameter rates that don't vanish. Good for non-stationary problems such as RL and RNNs. | Has no momentum by default. Still needs learning-rate tuning. | 1 |
| Adam | Momentum + adaptive scaling. Robust defaults. The most common starting point. | Often generalises slightly worse than well-tuned SGD + momentum on image tasks. Uses 2× extra memory. | 2 |
| Adamax | Stable with gradient spikes. Needs no bias correction on . | Rarely better than Adam in practice. | 2 |
| SMORMS3 | Adapts its own averaging window. Damps noisy parameters automatically. | Not built into major frameworks. Little published evaluation. | 3 |
Which optimizer should I use?
These are common rules of thumb, not laws:
-
Not sure? Start with Adam (or AdamW, see below) at
learning_rate=1e-3. It is the most forgiving choice. - Image models where final accuracy matters most (ResNets and other CNNs): SGD + Nesterov momentum (γ = 0.9) with a learning-rate schedule often matches or beats Adam on test accuracy.
-
RNNs, LSTMs, Transformers: use Adam or AdamW, and add gradient clipping (
global_clipnorm=1.0). - Very sparse features (large embedding tables, linear models on text): Adagrad works well here.
- Reinforcement learning or other shifting objectives: RMSProp or Adam.
- Training is unstable or the loss spikes: lower the learning rate first, then add clipping.
Whatever you choose, a learning-rate schedule (warm-up, cosine decay or step decay) often matters as much as the optimizer itself.
Using the optimizers in Keras
The snippet below trains the same small CNN on MNIST with any of the ten optimizers. The first nine are built into Keras. SMORMS3 uses the custom class from section 10.
import keras
from keras import layers
# Load MNIST and scale pixels to [0, 1]
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train = x_train[..., None].astype("float32") / 255.0
x_test = x_test[..., None].astype("float32") / 255.0
def make_model():
return keras.Sequential([
keras.Input(shape=(28, 28, 1)),
layers.Conv2D(32, 3, activation="relu"),
layers.MaxPooling2D(),
layers.Conv2D(64, 3, activation="relu"),
layers.MaxPooling2D(),
layers.Flatten(),
layers.Dropout(0.3),
layers.Dense(10, activation="softmax"),
])
optimizers = {
"SGD": keras.optimizers.SGD(learning_rate=0.01),
"SGD + clipping": keras.optimizers.SGD(learning_rate=0.01, global_clipnorm=1.0),
"Momentum": keras.optimizers.SGD(learning_rate=0.01, momentum=0.9),
"Nesterov": keras.optimizers.SGD(learning_rate=0.01, momentum=0.9, nesterov=True),
"Adagrad": keras.optimizers.Adagrad(learning_rate=0.01),
"Adadelta": keras.optimizers.Adadelta(learning_rate=1.0, rho=0.95),
"RMSprop": keras.optimizers.RMSprop(learning_rate=0.001, rho=0.9),
"Adam": keras.optimizers.Adam(learning_rate=0.001),
"Adamax": keras.optimizers.Adamax(learning_rate=0.002),
"SMORMS3": SMORMS3(learning_rate=0.001), # custom class from section 10
}
for name, optimizer in optimizers.items():
keras.utils.set_random_seed(42) # same initial weights for every run
model = make_model()
model.compile(optimizer=optimizer,
loss="sparse_categorical_crossentropy",
metrics=["accuracy"])
history = model.fit(x_train, y_train, batch_size=128, epochs=5,
validation_data=(x_test, y_test), verbose=0)
print(f"{name:15s} val accuracy: {history.history['val_accuracy'][-1]:.4f}")
To add a learning-rate schedule, pass a schedule object in place of the number:
schedule = keras.optimizers.schedules.CosineDecay(initial_learning_rate=0.1,
decay_steps=10_000)
optimizer = keras.optimizers.SGD(learning_rate=schedule, momentum=0.9, nesterov=True)
Keras 3 also includes several optimizers not covered here: AdamW, Nadam, Adafactor, Ftrl, Lion, Lamb and Muon.
Beyond the 10: worth knowing
-
AdamW (Loshchilov and Hutter, 2019): Adam with decoupled weight decay. In plain Adam, L2 regularisation gets divided by
along with the gradient, which weakens it for parameters with large gradients. AdamW subtracts the weight decay term directly from the weights instead. It is the standard optimizer for training Transformers and LLMs. In Keras:
keras.optimizers.AdamW(learning_rate=1e-3, weight_decay=0.004). -
Nadam (Dozat, 2016): Adam with Nesterov-style momentum. In Keras:
keras.optimizers.Nadam(). -
AMSGrad (Reddi et al., 2018): Adam variant that keeps the maximum of past
, which fixes a convergence problem in Adam's theory. In Keras:
keras.optimizers.Adam(amsgrad=True). -
Lion (Chen et al., 2023): uses only the sign of a momentum-like average. It needs less memory than Adam. In Keras:
keras.optimizers.Lion().
References
- Robbins & Monro (1951). A Stochastic Approximation Method.
- Polyak (1964). Some methods of speeding up the convergence of iteration methods. (momentum)
- Nesterov (1983). A method for solving the convex programming problem with convergence rate O(1/k²).
- Sutskever, Martens, Dahl & Hinton (2013). On the importance of initialization and momentum in deep learning.
- Pascanu, Mikolov & Bengio (2013). On the difficulty of training recurrent neural networks. (gradient clipping)
- Duchi, Hazan & Singer (2011). Adaptive Subgradient Methods for Online Learning and Stochastic Optimization. (Adagrad)
- Zeiler (2012). ADADELTA: An Adaptive Learning Rate Method.
- Hinton, Srivastava & Swersky (2012). Neural Networks for Machine Learning, Lecture 6e. (RMSProp)
- Kingma & Ba (2014). Adam: A Method for Stochastic Optimization. (Adam and Adamax)
- Funk (2015). RMSprop loses to SMORMS3 – Beware the Epsilon! (sifter.org blog)
- Loshchilov & Hutter (2019). Decoupled Weight Decay Regularization. (AdamW)
- Ruder (2016). An overview of gradient descent optimization algorithms.
- Schmidt, Schneider & Hennig (2021). Descending through a Crowded Valley: Benchmarking Deep Learning Optimizers.
All figures in this post were generated with a short NumPy + Matplotlib script that implements every optimizer from scratch.
import os
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch, FancyBboxPatch
import imageio.v2 as imageio
OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "images")
os.makedirs(OUT, exist_ok=True)
# ---------------------------------------------------------------- style -----
SURFACE = "#fcfcfb"
INK = "#0b0b0b"
INK2 = "#52514e"
MUTED = "#9a9893"
GRID = "#e6e5e1"
C = { # fixed categorical order (validated palette, light mode)
"blue": "#2a78d6", "orange": "#eb6834", "aqua": "#1baf7a",
"yellow": "#eda100", "magenta": "#e87ba4", "green": "#008300",
"violet": "#6250d6", "red": "#e34948",
}
plt.rcParams.update({
"figure.facecolor": SURFACE, "axes.facecolor": SURFACE,
"savefig.facecolor": SURFACE, "axes.edgecolor": MUTED,
"axes.labelcolor": INK2, "xtick.color": INK2, "ytick.color": INK2,
"text.color": INK, "font.size": 11, "axes.titlesize": 13,
"axes.titleweight": "bold", "axes.titlelocation": "left",
"axes.spines.top": False, "axes.spines.right": False,
"lines.linewidth": 2, "legend.frameon": False,
"font.family": "DejaVu Sans",
})
def save(fig, name):
fig.savefig(os.path.join(OUT, name), dpi=150, bbox_inches="tight")
plt.close(fig)
print("wrote", name)
# ------------------------------------------------------------ optimizers ----
class Optimizer:
def __init__(self, **hp):
self.hp, self.t, self.s = hp, 0, {}
def step(self, theta, grad_fn):
self.t += 1
return self.update(theta, grad_fn)
def z(self, k, like):
if k not in self.s:
self.s[k] = np.zeros_like(like)
return self.s[k]
class SGD(Optimizer):
def update(self, th, gf):
return th - self.hp["lr"] * gf(th)
class SGDClip(Optimizer):
def update(self, th, gf):
g = gf(th)
n = np.linalg.norm(g)
g = g * min(1.0, self.hp["max_norm"] / (n + 1e-12))
return th - self.hp["lr"] * g
class Momentum(Optimizer):
def update(self, th, gf):
v = self.hp["gamma"] * self.z("v", th) + self.hp["lr"] * gf(th)
self.s["v"] = v
return th - v
class Nesterov(Optimizer):
def update(self, th, gf):
v_prev = self.z("v", th)
g = gf(th - self.hp["gamma"] * v_prev) # look-ahead gradient
v = self.hp["gamma"] * v_prev + self.hp["lr"] * g
self.s["v"] = v
return th - v
class Adagrad(Optimizer):
def update(self, th, gf):
g = gf(th)
G = self.z("G", th) + g ** 2
self.s["G"] = G
return th - self.hp["lr"] * g / (np.sqrt(G) + 1e-8)
class Adadelta(Optimizer):
def update(self, th, gf):
rho, eps = self.hp["rho"], self.hp["eps"]
g = gf(th)
Eg = rho * self.z("Eg", th) + (1 - rho) * g ** 2
dx = -np.sqrt(self.z("Ed", th) + eps) / np.sqrt(Eg + eps) * g
self.s["Ed"] = rho * self.s["Ed"] + (1 - rho) * dx ** 2
self.s["Eg"] = Eg
return th + self.hp.get("lr", 1.0) * dx
class RMSProp(Optimizer):
def update(self, th, gf):
g = gf(th)
Eg = self.hp["rho"] * self.z("Eg", th) + (1 - self.hp["rho"]) * g ** 2
self.s["Eg"] = Eg
return th - self.hp["lr"] * g / (np.sqrt(Eg) + 1e-8)
class Adam(Optimizer):
def update(self, th, gf):
b1, b2, t = 0.9, 0.999, self.t
g = gf(th)
m = b1 * self.z("m", th) + (1 - b1) * g
v = b2 * self.z("v", th) + (1 - b2) * g ** 2
self.s["m"], self.s["v"] = m, v
mh, vh = m / (1 - b1 ** t), v / (1 - b2 ** t)
return th - self.hp["lr"] * mh / (np.sqrt(vh) + 1e-8)
class Adamax(Optimizer):
def update(self, th, gf):
b1, b2, t = 0.9, 0.999, self.t
g = gf(th)
m = b1 * self.z("m", th) + (1 - b1) * g
u = np.maximum(b2 * self.z("u", th), np.abs(g))
self.s["m"], self.s["u"] = m, u
return th - self.hp["lr"] / (1 - b1 ** t) * m / (u + 1e-8)
class SMORMS3(Optimizer):
def update(self, th, gf):
eps = 1e-16
g = gf(th)
mem = self.s.get("mem", np.ones_like(th))
r = 1.0 / (1.0 + mem)
m = (1 - r) * self.z("m", th) + r * g
v = (1 - r) * self.z("v", th) + r * g ** 2
x = m ** 2 / (v + eps)
step = np.minimum(self.hp["lr"], x) * g / (np.sqrt(v) + eps)
self.s["mem"] = 1 + mem * (1 - x)
self.s["m"], self.s["v"] = m, v
return th - step
def run(opt, start, grad_fn, steps):
th = np.array(start, float)
path = [th.copy()]
for _ in range(steps):
th = opt.step(th, grad_fn)
path.append(th.copy())
return np.array(path)
# ------------------------------------------------------- test surfaces ------
# A long, narrow ravine: steep in y, shallow in x. Minimum at (0, 0).
A, B = 0.05, 1.0
def ravine(x, y):
return A * x ** 2 + B * y ** 2
def ravine_grad(th):
return np.array([2 * A * th[0], 2 * B * th[1]])
START = (-9.0, 2.2)
XS, YS = np.linspace(-10.5, 3, 300), np.linspace(-3.2, 3.2, 300)
XX, YY = np.meshgrid(XS, YS)
ZZ = ravine(XX, YY)
LEVELS = np.geomspace(0.02, ravine(-10.5, 3.2), 14)
def contour(ax, title=None):
ax.contourf(XX, YY, ZZ, levels=LEVELS, cmap="Blues", alpha=0.22)
ax.contour(XX, YY, ZZ, levels=LEVELS, colors=[MUTED], linewidths=0.6,
alpha=0.6)
ax.plot(0, 0, marker="*", ms=14, color=INK, zorder=6)
ax.plot(*START, marker="o", ms=8, color=INK, mfc=SURFACE, zorder=6)
ax.set_xlim(XS[0], XS[-1]); ax.set_ylim(YS[0], YS[-1])
ax.set_xticks([]); ax.set_yticks([])
for s in ax.spines.values():
s.set_visible(False)
if title:
ax.set_title(title)
def draw_path(ax, p, color, label=None, ms=3.5, lw=2, alpha=1.0, z=4):
ax.plot(p[:, 0], p[:, 1], "-", color=color, lw=lw, alpha=alpha,
label=label, zorder=z)
ax.plot(p[:, 0], p[:, 1], "o", color=color, ms=ms, alpha=alpha,
mec=SURFACE, mew=0.8, zorder=z)
# Settings tuned so every optimizer is shown "reasonably configured".
STEPS = 60
CONFIGS = [
("SGD", lambda: SGD(lr=0.9)),
("SGD + clipping", lambda: SGDClip(lr=0.9, max_norm=1.0)),
("Momentum", lambda: Momentum(lr=0.3, gamma=0.8)),
("Nesterov", lambda: Nesterov(lr=0.3, gamma=0.8)),
("Adagrad", lambda: Adagrad(lr=2.5)),
("Adadelta", lambda: Adadelta(rho=0.9, eps=1e-2)),
("RMSProp", lambda: RMSProp(lr=0.4, rho=0.9)),
("Adam", lambda: Adam(lr=0.5)),
("Adamax", lambda: Adamax(lr=0.6)),
("SMORMS3", lambda: SMORMS3(lr=0.4)),
]
PATHS = {n: run(f(), START, ravine_grad, STEPS) for n, f in CONFIGS}
# ===================================================================== FIGS ==
def fig_hero():
fig = plt.figure(figsize=(10, 4.2))
ax = fig.add_axes([0.03, 0.12, 0.94, 0.62])
contour(ax)
sel = [("SGD", C["blue"]), ("Momentum", C["orange"]),
("RMSProp", C["aqua"]), ("Adam", C["violet"])]
for n, col in sel:
draw_path(ax, PATHS[n][:40], col, n, ms=3)
ax.legend(loc="upper center", fontsize=11, ncol=4,
bbox_to_anchor=(0.5, -0.0))
fig.text(0.03, 0.95, "10 famous Machine Learning Optimizers",
fontsize=21, weight="bold", va="top")
fig.text(0.03, 0.84, "Same start, same loss surface, different paths "
"to the minimum (★)", fontsize=12.5, color=INK2, va="top")
fig.savefig(os.path.join(OUT, "00_hero.png"), dpi=150)
plt.close(fig); print("wrote 00_hero.png")
def fig_sgd_lr():
fig, axes = plt.subplots(1, 3, figsize=(15, 4.2))
cases = [(0.15, "Too small (η = 0.15)", C["blue"]),
(0.9, "About right (η = 0.9)", C["aqua"]),
(1.02, "Too large (η = 1.02): diverges", C["red"])]
for ax, (lr, t, col) in zip(axes, cases):
contour(ax, t)
n = 9 if lr > 1 else 30
draw_path(ax, run(SGD(lr=lr), START, ravine_grad, n), col)
fig.suptitle("SGD: the learning rate η decides everything",
x=0.02, ha="left", fontsize=15, weight="bold", y=1.03)
save(fig, "01_sgd_learning_rate.png")
def fig_sgd_noise():
rng = np.random.default_rng(3)
def noisy(th):
return ravine_grad(th) + rng.normal(0, 0.35, 2)
fig, ax = plt.subplots(figsize=(10, 4.4))
contour(ax, "Full-batch gradient descent vs. stochastic (mini-batch) SGD")
draw_path(ax, run(SGD(lr=0.5), START, ravine_grad, 45), C["blue"],
"Full batch: smooth but each step costs a full pass")
draw_path(ax, run(SGD(lr=0.5), START, noisy, 45), C["orange"],
"Mini-batch SGD: noisy but cheap steps")
ax.legend(loc="lower right", bbox_to_anchor=(1, -0.16), ncol=1)
save(fig, "01b_sgd_vs_full_batch.png")
def fig_clipping():
fig, (a1, a2) = plt.subplots(1, 2, figsize=(14, 5.2),
gridspec_kw={"width_ratios": [1, 1.5]})
# left: vector picture
g = np.array([3.2, 1.2]); thr = 1.5
by_norm = g * min(1, thr / np.linalg.norm(g))
by_val = np.clip(g, -1.0, 1.0)
a1.add_patch(plt.Circle((0, 0), thr, fill=False, ls="--", color=MUTED,
lw=1.5))
a1.add_patch(plt.Rectangle((-1, -1), 2, 2, fill=False, ls=":",
color=MUTED, lw=1.5))
for vec, col, lab in [(g, C["red"], "raw gradient g"),
(by_norm, C["blue"], "clip by norm: ‖g‖ ≤ 1.5"),
(by_val, C["orange"], "clip by value: |gᵢ| ≤ 1")]:
a1.annotate("", xy=vec, xytext=(0, 0),
arrowprops=dict(arrowstyle="-|>", color=col, lw=2.5,
mutation_scale=18))
a1.plot([], [], color=col, lw=2.5, label=lab)
a1.text(-2.1, -2.05, "dashed circle = norm threshold "
"dotted square = value threshold", color=INK2, fontsize=9.5)
a1.set_xlim(-2.2, 3.6); a1.set_ylim(-2.2, 2.2); a1.set_aspect("equal")
a1.axhline(0, color=GRID, lw=1, zorder=0)
a1.axvline(0, color=GRID, lw=1, zorder=0)
a1.set_xticks([]); a1.set_yticks([])
for sp in a1.spines.values():
sp.set_visible(False)
a1.legend(loc="upper center", bbox_to_anchor=(0.5, -0.02), ncol=1,
fontsize=10.5)
a1.set_title("Norm clipping keeps the direction,\nvalue clipping can "
"bend it")
# right: a cliff (sudden drop in loss) catapults plain SGD
H, k, xc = 3.0, 25.0, -5.0
def f2(x, y):
return ravine(x, y) + H / (1 + np.exp(k * (x - xc)))
def g2(th):
sg = 1 / (1 + np.exp(-k * (th[0] - xc)))
return ravine_grad(th) + np.array([-H * k * sg * (1 - sg), 0.0])
xs = np.linspace(-10.5, 8, 500); ys = np.linspace(-3.2, 3.2, 300)
X2, Y2 = np.meshgrid(xs, ys); Z2 = f2(X2, Y2)
lv = np.linspace(0.05, Z2.max(), 22)
a2.contourf(X2, Y2, Z2, levels=lv, cmap="Blues", alpha=0.22)
a2.contour(X2, Y2, Z2, levels=lv, colors=[MUTED], linewidths=0.6,
alpha=0.6)
a2.plot(0, 0, marker="*", ms=14, color=INK, zorder=6)
st = (-9.0, 1.2)
a2.plot(*st, marker="o", ms=8, color=INK, mfc=SURFACE, zorder=7)
a2.text(xc, 2.85, "cliff: loss drops by 3 here", ha="center",
color=C["red"], fontsize=10)
p_raw = run(SGD(lr=0.7), st, g2, 14)
p_clip = run(SGDClip(lr=0.7, max_norm=1.0), st, g2, 40)
draw_path(a2, p_clip, C["blue"], "SGD + clipping (max norm 1)")
draw_path(a2, p_raw, C["red"], "plain SGD: one step of ~11 units",
alpha=0.9)
a2.annotate("one step: x = −5.0 → +6.3", xy=(6.2, 0.1),
xytext=(0.8, 1.6), color=C["red"], fontsize=10.5,
arrowprops=dict(arrowstyle="->", color=C["red"]))
a2.set_xlim(-10.5, 8); a2.set_ylim(-3.2, 3.2)
a2.set_xticks([]); a2.set_yticks([])
for sp in a2.spines.values():
sp.set_visible(False)
a2.set_title("A cliff in the loss throws plain SGD far past the minimum")
a2.legend(loc="upper center", bbox_to_anchor=(0.5, -0.02), ncol=2)
save(fig, "02_gradient_clipping.png")
def fig_momentum():
fig, ax = plt.subplots(figsize=(10, 4.4))
contour(ax, "Momentum damps zig-zags and builds speed along the valley")
draw_path(ax, run(SGD(lr=0.95), START, ravine_grad, 40), C["blue"],
"SGD (η = 0.95)")
draw_path(ax, run(Momentum(lr=0.3, gamma=0.8), START, ravine_grad, 40),
C["orange"], "Momentum (η = 0.3, γ = 0.8)")
ax.legend(loc="lower right", bbox_to_anchor=(1, -0.16), ncol=2)
save(fig, "03_momentum.png")
def arrow(ax, p, q, col, lw=2.5, ls="-", z=5):
ax.add_patch(FancyArrowPatch(p, q, arrowstyle="-|>", mutation_scale=18,
color=col, lw=lw, ls=ls, zorder=z,
shrinkA=0, shrinkB=0))
def fig_nesterov():
fig, (a1, a2) = plt.subplots(1, 2, figsize=(13, 4.8), sharey=True)
th = np.array([0.0, 0.0]); mom = np.array([3.0, 1.2])
g_here = np.array([0.3, -1.4]); g_ahead = np.array([-0.9, -1.5])
for ax, title, nest in [(a1, "Classical momentum", False),
(a2, "Nesterov momentum", True)]:
ax.plot(*th, "o", ms=10, color=INK, zorder=6)
ax.text(th[0] - 0.15, th[1] - 0.35, "θₜ", fontsize=13, ha="right")
arrow(ax, th, th + mom, C["blue"])
ax.text(*(th + mom / 2 + [-0.1, 0.35]), "−γ·vₜ₋₁ (momentum)",
color=C["blue"], fontsize=11, ha="center")
if not nest:
arrow(ax, th, th + g_here, C["orange"])
ax.text(0.45, -1.1, "−η·∇J(θₜ)\nslope where you stand",
color=C["orange"], fontsize=10.5)
end = th + mom + g_here
arrow(ax, th + mom, end, C["orange"], lw=1.5, ls="--")
else:
ahead = th + mom
ax.plot(*ahead, "o", ms=9, color=C["blue"], mfc=SURFACE, zorder=6)
ax.text(ahead[0] + 0.15, ahead[1] + 0.15, "look-ahead point\nθₜ − γvₜ₋₁",
fontsize=10.5, color=INK2)
arrow(ax, ahead, ahead + g_ahead, C["orange"])
ax.text(ahead[0] + 0.05, ahead[1] - 1.1,
"−η·∇J(θₜ − γvₜ₋₁)\nslope where\nyou'll land",
color=C["orange"], fontsize=10.5, ha="left")
end = ahead + g_ahead
arrow(ax, th, end, C["aqua"], lw=3)
ax.plot(*end, "o", ms=9, color=C["aqua"], zorder=6)
ax.text(end[0] + 0.15, end[1] - 0.3, "θₜ₊₁", fontsize=13,
color=INK)
ax.set_title(title)
ax.set_xlim(-1, 5.4); ax.set_ylim(-1.6, 2.2); ax.set_aspect("equal")
ax.set_xticks([]); ax.set_yticks([])
for s in ax.spines.values():
s.set_visible(False)
a2.plot([], [], color=C["aqua"], lw=3, label="actual update")
a2.legend(loc="lower right")
save(fig, "04_nesterov_vs_momentum.png")
def fig_adagrad_rmsprop():
# effective learning rate eta / sqrt(state) with a steady gradient of 1
T = 300; t = np.arange(1, T + 1)
g = np.ones(T)
ada = 0.1 / np.sqrt(np.cumsum(g ** 2))
Eg = np.zeros(T); e = 0
for i in range(T):
e = 0.9 * e + 0.1 * g[i] ** 2; Eg[i] = e
rms = 0.1 / np.sqrt(Eg)
fig, (a1, a2) = plt.subplots(1, 2, figsize=(13, 4.4))
a1.plot(t, ada, color=C["blue"], label="Adagrad: η / √(Σ g²)")
a1.plot(t, rms, color=C["orange"], label="RMSProp: η / √E[g²]")
a1.set_xlabel("step t"); a1.set_ylabel("effective learning rate")
a1.set_title("Effective learning rate\n(constant gradient, η = 0.1)")
a1.grid(axis="y", color=GRID); a1.legend()
a1.text(T, ada[-1] + 0.01, f"{ada[-1]:.4f}", ha="right", color=C["blue"],
fontsize=10)
contour(a2, "Same ravine, 60 steps\n ")
draw_path(a2, PATHS["Adagrad"], C["blue"], "Adagrad")
draw_path(a2, PATHS["RMSProp"], C["orange"], "RMSProp")
a2.legend(loc="lower right", bbox_to_anchor=(1, -0.14), ncol=2)
fig.suptitle("Adagrad's step keeps shrinking; RMSProp's levels off",
x=0.02, ha="left", fontsize=15, weight="bold", y=1.06)
save(fig, "05_adagrad_vs_rmsprop.png")
def fig_adadelta():
T = 400; t = np.arange(1, T + 1)
fig, ax = plt.subplots(figsize=(10, 4.2))
for eps, col in [(1e-6, C["blue"]), (1e-4, C["orange"]),
(1e-2, C["aqua"])]:
opt = Adadelta(rho=0.95, eps=eps)
th = np.array([0.0]); st = []
for _ in range(T):
new = opt.step(th, lambda x: np.array([1.0]))
st.append(abs(new[0] - th[0])); th = new
ax.plot(t, st, color=col, label=f"ε = {eps:g}")
ax.set_yscale("log"); ax.grid(axis="y", color=GRID)
ax.set_xlabel("step t"); ax.set_ylabel("|Δθₜ| (log scale)")
ax.set_title("Adadelta has no η, but its step starts proportional to √ε,\n"
"so ε quietly acts as the learning rate")
ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.17), ncol=3,
title="constant gradient g = 1")
save(fig, "06_adadelta_step_size.png")
def fig_adam_bias():
T = 40; t = np.arange(1, T + 1)
rng = np.random.default_rng(1)
g = 1.0 + rng.normal(0, 0.25, T)
b1 = 0.9; m = 0; raw, cor = [], []
for i in range(T):
m = b1 * m + (1 - b1) * g[i]
raw.append(m); cor.append(m / (1 - b1 ** (i + 1)))
fig, ax = plt.subplots(figsize=(10, 4.2))
ax.plot(t, g, "o", color=MUTED, ms=4, label="raw gradients gₜ")
ax.plot(t, raw, color=C["red"], label="mₜ (starts at 0 → biased low)")
ax.plot(t, cor, color=C["blue"], label="m̂ₜ = mₜ / (1 − β₁ᵗ) (corrected)")
ax.axhline(1, color=INK2, lw=1, ls=":")
ax.text(T + 0.5, 1, "true mean", va="center", color=INK2, fontsize=10)
ax.set_xlabel("step t"); ax.set_ylim(0, 1.7)
ax.grid(axis="y", color=GRID); ax.legend(loc="lower right")
ax.set_title("Adam's bias correction fixes the slow start (β₁ = 0.9)")
save(fig, "08_adam_bias_correction.png")
def fig_adamax():
T = 80; t = np.arange(1, T + 1)
rng = np.random.default_rng(7)
g = rng.normal(0, 0.3, T); g[25] = 6.0 # one big spike
v = 0; u = 0; V, U = [], []
for i in range(T):
v = 0.999 * v + 0.001 * g[i] ** 2
u = max(0.999 * u, abs(g[i]))
V.append(np.sqrt(v / (1 - 0.999 ** (i + 1)))); U.append(u)
fig, ax = plt.subplots(figsize=(10, 4.2))
ax.bar(t, np.abs(g), color=GRID, width=0.8, label="|gₜ|")
ax.plot(t, V, color=C["blue"], label="Adam denominator √v̂ₜ (L2)")
ax.plot(t, U, color=C["orange"], label="Adamax denominator uₜ (L∞)")
ax.set_xlabel("step t"); ax.legend(loc="center right")
ax.set_title("Adamax tracks the largest recent gradient, not the average")
ax.annotate("gradient spike", xy=(26, 6), xytext=(36, 5.2),
arrowprops=dict(arrowstyle="->", color=INK2), color=INK2)
save(fig, "09_adamax_infinity_norm.png")
def fig_smorms3():
# how the per-step learning rate min(lr, m^2/v) reacts to noise
T = 150; t = np.arange(1, T + 1)
rng = np.random.default_rng(0)
fig, ax = plt.subplots(figsize=(10, 4.2))
for noise, col, lab in [(0.1, C["blue"], "consistent gradients"),
(1.5, C["orange"], "noisy gradients")]:
mem = 1.0; m = v = 0; L = []
for i in range(T):
g = 1.0 + rng.normal(0, noise)
r = 1 / (1 + mem)
m = (1 - r) * m + r * g; v = (1 - r) * v + r * g * g
x = m * m / (v + 1e-16)
L.append(min(0.5, x)); mem = 1 + mem * (1 - x)
ax.plot(t, L, color=col, label=lab)
ax.axhline(0.5, color=INK2, ls=":", lw=1)
ax.text(T, 0.52, "η cap = 0.5", ha="right", color=INK2, fontsize=10)
ax.set_ylim(0, 0.6); ax.set_xlabel("step t")
ax.set_ylabel("effective rate min(η, m²/v)")
ax.grid(axis="y", color=GRID)
ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.18), ncol=2)
ax.set_title("SMORMS3 shrinks its step when the gradient is mostly noise")
save(fig, "10_smorms3_noise.png")
def fig_family_tree():
fig, ax = plt.subplots(figsize=(12, 6.2))
ax.set_xlim(0, 12); ax.set_ylim(0, 6.6); ax.axis("off")
nodes = {
"SGD": (6, 6.0, C["blue"]),
"SGD + clipping": (1.2, 4.4, C["blue"]),
"Momentum": (3.9, 4.4, C["orange"]),
"Adagrad": (8.6, 4.4, C["aqua"]),
"Nesterov": (1.7, 2.8, C["orange"]),
"RMSProp": (7.4, 2.8, C["aqua"]),
"Adadelta": (10.4, 2.8, C["aqua"]),
"Adam": (5.8, 1.0, C["violet"]),
"SMORMS3": (9.0, 1.0, C["aqua"]),
"Adamax": (2.6, 1.0, C["violet"]),
}
edges = [("SGD", "SGD + clipping"), ("SGD", "Momentum"),
("SGD", "Adagrad"), ("Momentum", "Nesterov"),
("Adagrad", "RMSProp"), ("Adagrad", "Adadelta"),
("Momentum", "Adam"), ("RMSProp", "Adam"),
("RMSProp", "SMORMS3"), ("Adam", "Adamax")]
w, h = 2.1, 0.62
for a, b in edges:
(x1, y1, _), (x2, y2, _) = nodes[a], nodes[b]
if abs(y1 - y2) < 0.1: # same row: side to side
sgn = np.sign(x2 - x1)
p1, p2 = (x1 + sgn * w / 2, y1), (x2 - sgn * w / 2, y2)
else:
p1, p2 = (x1, y1 - h / 2), (x2, y2 + h / 2)
ax.add_patch(FancyArrowPatch(p1, p2,
arrowstyle="-|>", mutation_scale=14, color=MUTED, lw=1.4,
connectionstyle="arc3,rad=0.0", zorder=1,
shrinkA=2, shrinkB=2))
for n, (x, y, col) in nodes.items():
ax.add_patch(FancyBboxPatch((x - w / 2, y - h / 2), w, h,
boxstyle="round,pad=0.02,rounding_size=0.15",
fc=SURFACE, ec=col, lw=2.2, zorder=3))
ax.text(x, y, n, ha="center", va="center", fontsize=12,
weight="bold", zorder=4)
for lab, col, y in [("plain / stabilised gradient", C["blue"], 6.45),
("momentum family", C["orange"], 6.1),
("adaptive learning rate family", C["aqua"], 5.75),
("momentum + adaptive (Adam family)", C["violet"],
5.4)]:
ax.plot([0.15, 0.55], [y, y], color=col, lw=4)
ax.text(0.7, y, lab, va="center", fontsize=10.5, color=INK2)
ax.text(12, 0.15, "arrows: \"builds on\"", ha="right", color=INK2,
fontsize=10)
save(fig, "family_tree.png")
def fig_grid():
fig, axes = plt.subplots(5, 2, figsize=(12, 13))
for ax, (n, _) in zip(axes.flat, CONFIGS):
contour(ax, n)
draw_path(ax, PATHS["SGD"], MUTED, ms=0, lw=1.2, alpha=0.7, z=3)
draw_path(ax, PATHS[n], C["blue"], ms=3)
d = np.linalg.norm(PATHS[n][-1])
ax.text(0.98, 0.04, f"distance to ★ after {STEPS} steps: {d:.2f}",
transform=ax.transAxes, ha="right", fontsize=9.5,
color=INK2)
fig.suptitle("All 10 optimizers on the same ravine (grey line = plain "
"SGD for reference)", x=0.02, ha="left", fontsize=15,
weight="bold", y=1.0)
fig.tight_layout()
save(fig, "comparison_grid.png")
def gif_race():
sel = [("SGD", C["blue"]), ("Momentum", C["orange"]),
("RMSProp", C["aqua"]), ("Adam", C["violet"])]
frames = []
for k in range(1, 46):
fig, ax = plt.subplots(figsize=(8, 4), dpi=90)
contour(ax, f"Optimizer race · step {k}")
for n, col in sel:
p = PATHS[n][:k + 1]
ax.plot(p[:, 0], p[:, 1], "-", color=col, lw=2, label=n)
ax.plot(p[-1, 0], p[-1, 1], "o", color=col, ms=8, mec=SURFACE,
mew=1.5, zorder=7)
ax.legend(loc="lower right", ncol=4, bbox_to_anchor=(1, -0.16),
fontsize=10)
fig.canvas.draw()
frames.append(np.asarray(fig.canvas.buffer_rgba())[..., :3].copy())
plt.close(fig)
frames += [frames[-1]] * 15
imageio.mimsave(os.path.join(OUT, "optimizer_race.gif"), frames,
duration=0.12, loop=0)
print("wrote optimizer_race.gif")
# ------------------------------------------- real training: tiny MLP --------
def spirals(n=300, seed=0):
rng = np.random.default_rng(seed)
t = np.sqrt(rng.uniform(0, 1, n)) * 3 * np.pi
X0 = np.c_[t * np.cos(t), t * np.sin(t)]
X = np.r_[X0, -X0] / 10 + rng.normal(0, 0.05, (2 * n, 2))
y = np.r_[np.zeros(n), np.ones(n)]
return X, y
def train_mlp(make_opt, epochs=150, seed=0):
X, y = spirals()
rng = np.random.default_rng(seed)
sizes = [2, 32, 32, 1]
shapes = [(sizes[i], sizes[i + 1]) for i in range(3)]
params = []
for a, b in shapes:
params += [rng.normal(0, np.sqrt(2 / a), (a, b)).ravel(), np.zeros(b)]
flat = np.concatenate(params)
cuts = np.cumsum([p.size for p in params])[:-1]
def unpack(f):
ps = np.split(f, cuts)
return [ps[i].reshape(shapes[i // 2]) if i % 2 == 0 else ps[i]
for i in range(6)]
def loss_grad(f, Xb, yb):
W1, b1, W2, b2, W3, b3 = unpack(f)
h1 = np.tanh(Xb @ W1 + b1); h2 = np.tanh(h1 @ W2 + b2)
z = (h2 @ W3 + b3).ravel(); p = 1 / (1 + np.exp(-z))
loss = -np.mean(yb * np.log(p + 1e-9) + (1 - yb) * np.log(1 - p + 1e-9))
dz = (p - yb)[:, None] / len(yb)
gW3 = h2.T @ dz; gb3 = dz.sum(0)
d2 = (dz @ W3.T) * (1 - h2 ** 2)
gW2 = h1.T @ d2; gb2 = d2.sum(0)
d1 = (d2 @ W2.T) * (1 - h1 ** 2)
gW1 = Xb.T @ d1; gb1 = d1.sum(0)
g = np.concatenate([a.ravel() for a in (gW1, gb1, gW2, gb2, gW3, gb3)])
return loss, g
opt = make_opt(); hist = []
for ep in range(epochs):
idx = rng.permutation(len(y))
for s in range(0, len(y), 32):
b = idx[s:s + 32]
flat = opt.step(flat, lambda f: loss_grad(f, X[b], y[b])[1])
hist.append(loss_grad(flat, X, y)[0])
return np.array(hist)
def fig_loss_curves():
groups = [
("Plain & momentum", [
("SGD", lambda: SGD(lr=0.1), C["blue"]),
("SGD + clipping", lambda: SGDClip(lr=0.1, max_norm=1.0),
C["orange"]),
("Momentum", lambda: Momentum(lr=0.1, gamma=0.9), C["aqua"]),
("Nesterov", lambda: Nesterov(lr=0.1, gamma=0.9), C["violet"])]),
("Adaptive learning rate", [
("Adagrad", lambda: Adagrad(lr=0.05), C["blue"]),
("Adadelta", lambda: Adadelta(rho=0.95, eps=1e-6), C["orange"]),
("RMSProp", lambda: RMSProp(lr=0.003, rho=0.9), C["aqua"])]),
("Adam family & SMORMS3", [
("Adam", lambda: Adam(lr=0.005), C["blue"]),
("Adamax", lambda: Adamax(lr=0.01), C["orange"]),
("SMORMS3", lambda: SMORMS3(lr=0.005), C["aqua"])]),
]
fig, axes = plt.subplots(1, 3, figsize=(16, 4.6), sharey=True)
finals = {}
for ax, (title, items) in zip(axes, groups):
for n, mk, col in items:
h = train_mlp(mk); finals[n] = h[-1]
ax.plot(np.arange(1, len(h) + 1), h, color=col, label=n)
ax.set_title(title); ax.set_xlabel("epoch")
ax.grid(axis="y", color=GRID)
ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.17), ncol=2)
ax.set_yscale("log")
axes[0].set_ylabel("training loss (log)")
fig.suptitle("Training a small 2-layer network on the two-spirals "
"dataset (batch 32)", x=0.02, ha="left", fontsize=15,
weight="bold", y=1.04)
save(fig, "loss_curves.png")
for n, v in finals.items():
print(f" final loss {n:15s} {v:.4f}")
if __name__ == "__main__":
fig_hero(); fig_sgd_lr(); fig_sgd_noise(); fig_clipping(); fig_momentum()
fig_nesterov(); fig_adagrad_rmsprop(); fig_adadelta(); fig_adam_bias()
fig_adamax(); fig_smorms3(); fig_family_tree(); fig_grid()
fig_loss_curves(); gif_race()
for n, p in PATHS.items():
print(f" {n:15s} dist to min after {STEPS}: {np.linalg.norm(p[-1]):.3f}")
Top comments (0)