Adam with a learning-rate schedule
The problem
Smoothing: 100,000 noisy points dᵢ, at tᵢ evenly spread over [0, 1], become a curve xᵢ that follows them without their noise, by penalized least squares:
f(x) = Σᵢ (xᵢ − dᵢ)² + λ Σᵢ (xᵢ₊₁ − xᵢ)², λ = 50
The first sum keeps the curve near the data, the second keeps it smooth. Each of the 100,000 values xᵢ is a parameter.
The data are made so that the answer is known exactly. The curve
x*ᵢ = sin(2π tᵢ) + 0.3 sin(10π tᵢ) + εᵢ, εᵢ uniform in [−0.01, 0.01]
is smooth but for a little noise of its own, and the data are dᵢ = xᵢ + λ (L x)ᵢ, with L the Laplacian of the chain of points ((L x)ᵢ = 2xᵢ − xᵢ₋₁ − xᵢ₊₁ inside, one neighbor at the ends). The gradient of f is 2(x − d) + 2λ L x, so it's 0 exactly where (I + λ L) x = d: at x*, the only minimum, since f is a convex quadratic. The data's noise is the curve's, amplified up to 1 + 4λ = 201 times by L: about ±2 around a curve of amplitude 1.3.
What makes it hard
Many parameters, and coupled ones: each value is pulled toward its data point and toward its two neighbors. The Hessian, 2(I + λ L), has curvatures from 2 to 2(1 + 4λ) = 402, so the curve's rough components settle in a few steps and its smooth ones in hundreds. A method that stores a matrix of the parameters (10¹⁰ entries) is out of the question, and a line search would cost evaluations every step.
Representation
A Real genome of 100,000 genes in [−10, 10], starting from a flat curve, all zeros. The fitness
is f, to minimize, with its gradient: Differentiable in Rust, and gradient=True in Python, where
the function returns the value and the gradient. The gradient is computed in the same order of
operations in both, so both versions take the same steps.
Algorithm
FirstOrder with Step::Adam: Adam (Kingma and Ba 2015, Algorithm 1), with β₁ = 0.9, β₂ = 0.999
and ε = 1e-8. Each step keeps averages of the gradient and of its square per value, corrects them
for their start at 0, and moves each value by about the learning rate α in the direction of the
averaged gradient, scaled by its typical size. One gradient per step, and the memory of two vectors.
The learning rate starts at 0.05 and is halved every 500 steps, by Engine::control
(set_learning_rate in Rust, running.learning_rate in Python). The run stops by itself
(StopReason::Converged) when no component of the gradient exceeds 1e-9.
Then the same number of steps again, with the learning rate kept at 0.05, as a contrast.
Output
A row every 100 steps and the last: the learning rate, the loss f, the largest component of the gradient, and the distance to the answer, the largest |xᵢ − x*ᵢ|.
For 500 steps at 0.05, Adam gets within about 1e-2 of the answer and stays there: its steps are too long for the stiffest components of the curve, which it keeps overshooting, so the gradient doesn't shrink. Halved to 0.025, the steps fit, and every 100 steps gain more than two digits, until the gradient is below 1e-9 after 925 steps: every value is then within 2.5e-12 of the answer. With the learning rate kept at 0.05, the same 925 steps end within 1.3e-2.
The project page plays both runs back: the distance to the answer and the learning rate at every step.
Good results
The minimum is f(x) = 50,609.073…, at distance 0 from x. The run converges in 925 steps, one gradient each, within 2.5e-12 of x*; a step takes about 1.5 ms (the method's own share, 0.7 ms, is linear in the parameters).
Adam's strength is a step of about α per parameter whatever the scale of its gradient, and an average that smooths noisy gradients (mini-batches): on this deterministic, well-scaled quadratic the others do as well or better. With the same convergence test, from the same start:
| Step rule | Steps to converge |
|---|---|
| Adam, α = 0.05 halved every 500 steps (this example) | 925 |
| Adam, α constant at 0.025 or 0.04 | none in 200,000 (stays within 4e-3 to 6e-3) |
| Gradient descent, α = 0.0049 (just below the stability limit 2 / 402) | 2,202 |
| Polyak's momentum, α = 0.005, μ = 0.9 | 433 |
| Nesterov's accelerated gradient, α = 0.003, μ = 0.9 | 367 |
Lbfgsb, the quasi-Newton method with a line search, learns the curvatures from its last
gradients: from the same start, with the gradient tolerance at 1e-9, it stops after 93 iterations
and 218 evaluations, within 4.0e-7 of the answer, when its line search can no longer lower the
loss: near the minimum, the loss of 50,609 changes by less than its rounding, while the gradient
still points the way. It's the better choice for a smooth deterministic function, where a few
digits of the answer suffice or the loss is small. Adam and momentum follow the gradient alone,
which brings them closer here, and suit the cases a line search can't serve: gradients that are
noisy estimates, such as a model trained on mini-batches, or evaluations too costly to spend on a
line search.
Reference: Kingma, D. P. and Ba, J. (2015). Adam: a method for stochastic optimization. ICLR 2015.
Known optimum: the curve x* the data are made from, at distance 0
Source: examples/adam
Interactive run: tachsin.gr/projects/genoxide/examples/adam
cargo run --release --example adam
//! Adam with a learning-rate schedule: smooth 100,000 noisy points into a curve, the curve's
//! 100,000 values the parameters, by penalized least squares with its gradient.
//!
//! The data are made so that the exact answer is known, and the run ends at it: Adam's learning
//! rate is halved every 500 steps by `control`, and the run stops when the gradient vanishes.
//! Then the same number of steps with the learning rate kept constant, which hovers around the
//! answer instead.
//!
//! With `GENOXIDE_TRACE=<file>`, it also writes a trace of its runs for the plot on the example's
//! page, with `trace.rs`.
//!
//! ```text
//! cargo run --release --example adam
//! ```
mod trace;
use genoxide::algorithm::first_order::Step;
use genoxide::math;
use genoxide::prelude::*;
use std::f64::consts::TAU;
// the number of points, and of the curve's values
const POINTS: usize = 100_000;
// the weight of the roughness against the misfit
const LAMBDA: f64 = 50.0;
// Adam's first learning rate, halved every `HALVING` steps
const RATE: f64 = 0.05;
const HALVING: u64 = 500;
// a row of the table every this many steps
const EVERY: u64 = 100;
fn main() -> Result<()> {
let (exact, data) = problem()?;
// the misfit to the data and the roughness, with its gradient
let smoothing = Differentiable(|x: &Reals, gradient: &mut [f64]| loss(x, &data, gradient));
let adam = || {
FirstOrder::builder(Real::uniform(POINTS, -10.0..=10.0)?)
.step(Step::adam(RATE))
.initial_genome(Reals::from(vec![0.0; POINTS]))
.gradient_tolerance(1e-9)
.minimize()
.build()
};
// with GENOXIDE_TRACE=<file>, a trace of the runs for the plot on the example's page
let mut trace = trace::Trace::from_env();
let mut rows = Vec::new();
let mut engine = Engine::new(adam()?, smoothing)
.stop_when(Stop::generations(10_000))
.control(|adam: &mut FirstOrder, progress| {
let point = &adam.population()[0];
let distance = distance(point.genome(), &exact);
let rate = adam.step().learning_rate();
trace.scheduled(progress.generation(), distance, rate);
rows.push(Row {
step: progress.generation(),
rate,
loss: point.fitness().and_then(Fitness::score).expect("valid"),
gradient: adam.gradient_norm(),
distance,
});
// the learning rate of the next step: halved every `HALVING` steps
let halvings = (progress.generation() / HALVING) as i32;
adam.set_learning_rate(RATE * 0.5f64.powi(halvings))
});
let outcome = engine.run()?;
let steps = engine.algorithm().iterations();
drop(engine);
println!("Smoothing {POINTS} noisy points: the misfit plus {LAMBDA} times the roughness");
println!("Adam from a flat curve, its learning rate {RATE} halved every {HALVING} steps");
println!(" step learning rate loss largest gradient distance to the answer");
let last = rows.len() - 1;
for (index, row) in rows.iter().enumerate() {
if row.step.is_multiple_of(EVERY) || index == last {
println!(
"{:>6} {:>13} {:>11} {:>16} {:>22}",
row.step,
scientific(row.rate, 1),
scientific(row.loss, 7),
scientific(row.gradient, 1),
scientific(row.distance, 1)
);
}
}
assert_eq!(outcome.stop_reason(), StopReason::Converged);
let end = &rows[last];
println!(
"converged after {steps} steps and {} evaluations: every value within {} of the answer",
outcome.evaluations(),
scientific(end.distance, 1)
);
// the same number of steps, the learning rate constant
let mut constant = Vec::new();
Engine::new(adam()?, smoothing)
.stop_when(Stop::generations(steps))
.control(|adam: &mut FirstOrder, progress| {
let distance = distance(adam.population()[0].genome(), &exact);
trace.constant(progress.generation(), distance, adam.step().learning_rate());
constant.push(distance);
Ok(())
})
.run()?;
println!(
"the learning rate kept at {RATE}, after the same {steps} steps: within {} of the answer",
scientific(constant[constant.len() - 1], 1)
);
trace.write();
Ok(())
}
// a row of the table
struct Row {
step: u64,
rate: f64,
loss: f64,
gradient: f64,
distance: f64,
}
// the exact answer x*, a smooth curve with a little noise of its own, and the data that make it
// the answer: d = x* + λ L x*, where L is the Laplacian of the chain of points, so that
// (I + λ L) x* = d, the normal equations of the loss
fn problem() -> Result<(Vec<f64>, Vec<f64>)> {
let noise =
Real::uniform(POINTS, -0.01..=0.01)?.random_genome(&mut StreamRng::seed_from_u64(1));
let exact: Vec<f64> = (0..POINTS)
.map(|i| {
let t = i as f64 / (POINTS - 1) as f64;
math::sin(TAU * t) + 0.3 * math::sin(5.0 * TAU * t) + noise[i]
})
.collect();
let mut data = exact.clone();
for i in 0..POINTS - 1 {
let step = exact[i + 1] - exact[i];
data[i + 1] += LAMBDA * step;
data[i] -= LAMBDA * step;
}
Ok((exact, data))
}
// Σ (xᵢ − dᵢ)² + λ Σ (xᵢ₊₁ − xᵢ)², and its gradient into `gradient`
fn loss(x: &[f64], data: &[f64], gradient: &mut [f64]) -> f64 {
let mut value = 0.0;
for ((g, &xi), &di) in gradient.iter_mut().zip(x).zip(data) {
let misfit = xi - di;
*g = 2.0 * misfit;
value += misfit * misfit;
}
let weight = 2.0 * LAMBDA;
for i in 0..x.len() - 1 {
let step = x[i + 1] - x[i];
gradient[i + 1] += weight * step;
gradient[i] -= weight * step;
value += LAMBDA * step * step;
}
value
}
// the largest difference between the curve and the answer
fn distance(x: &[f64], exact: &[f64]) -> f64 {
x.iter()
.zip(exact)
.map(|(a, b)| (a - b).abs())
.fold(0.0, f64::max)
}
// `digits` digits after the point, e.g. 1.2e-7
fn scientific(value: f64, digits: usize) -> String {
format!("{value:.digits$e}")
}
python examples/adam/main.py
"""Adam with a learning-rate schedule: smooth 100,000 noisy points into a curve, the curve's
100,000 values the parameters, by penalized least squares with its gradient.
The data are made so that the exact answer is known, and the run ends at it: Adam's learning rate
is halved every 500 steps by ``control``, and the run stops when the gradient vanishes. Then the
same number of steps with the learning rate kept constant, which hovers around the answer
instead.
With ``GENOXIDE_TRACE=<file>``, it also writes a trace of its runs for the plot on the example's
page, with trace.py.
python examples/adam/main.py
"""
import math
import numpy as np
import genoxide as gx
from trace import Trace
# the number of points, and of the curve's values
POINTS = 100_000
# the weight of the roughness against the misfit
LAMBDA = 50.0
# Adam's first learning rate, halved every HALVING steps
RATE = 0.05
HALVING = 500
# a row of the table every this many steps
EVERY = 100
def problem():
"""The exact answer x*, a smooth curve with a little noise of its own, and the data that make
it the answer: d = x* + lambda L x*, where L is the Laplacian of the chain of points, so that
(I + lambda L) x* = d, the normal equations of the loss."""
noise = gx.Real((-0.01, 0.01), length=POINTS).random_genome(1)
t = np.arange(POINTS) / (POINTS - 1)
exact = gx.math.sin(math.tau * t) + 0.3 * gx.math.sin(5.0 * math.tau * t) + noise
data = exact.copy()
step = LAMBDA * np.diff(exact)
data[1:] += step
data[:-1] -= step
return exact, data
EXACT, DATA = problem()
def loss(x):
"""The misfit to the data plus lambda times the roughness, and its gradient, in the order of
the Rust example's operations, so that both take the same steps."""
misfit = x - DATA
gradient = 2.0 * misfit
step = np.diff(x)
weighted = (2.0 * LAMBDA) * step
gradient[1:] += weighted
gradient[:-1] -= weighted
value = float(np.dot(misfit, misfit) + LAMBDA * np.dot(step, step))
return value, gradient
def distance(x):
"""The largest difference between the curve and the answer."""
return float(np.max(np.abs(x - EXACT)))
def scientific(value, digits):
"""``digits`` digits after the point, e.g. 1.2e-7."""
mantissa, exponent = f"{value:.{digits}e}".split("e")
return f"{mantissa}e{int(exponent)}"
def adam():
return gx.FirstOrder(
gx.Real((-10.0, 10.0), length=POINTS),
step="adam",
learning_rate=RATE,
initial_genome=np.zeros(POINTS),
gradient_tolerance=1e-9,
objective="minimize",
)
# with GENOXIDE_TRACE=<file>, a trace of the runs for the plot on the example's page
trace = Trace()
rows = []
def schedule(running, progress):
point = progress.population[0]
rate = running.learning_rate
trace.scheduled(progress.generation, distance(point), rate)
rows.append(
(progress.generation, rate, progress.scores[0], running.gradient_norm, distance(point))
)
# the learning rate of the next step: halved every HALVING steps
running.learning_rate = RATE * 0.5 ** (progress.generation // HALVING)
result = adam().run(loss, gradient=True, generations=10_000, control=schedule)
steps = result.generations
print(f"Smoothing {POINTS} noisy points: the misfit plus {LAMBDA:g} times the roughness")
print(f"Adam from a flat curve, its learning rate {RATE} halved every {HALVING} steps")
print(" step learning rate loss largest gradient distance to the answer")
for index, (step, rate, value, gradient, gap) in enumerate(rows):
if step % EVERY == 0 or index == len(rows) - 1:
print(
f"{step:>6} {scientific(rate, 1):>13} {scientific(value, 7):>11} "
f"{scientific(gradient, 1):>16} {scientific(gap, 1):>22}"
)
assert result.stop_reason == "converged"
print(
f"converged after {steps} steps and {result.evaluations} evaluations: every value within "
f"{scientific(rows[-1][4], 1)} of the answer"
)
# the same number of steps, the learning rate constant
constant = []
def keep(running, progress):
gap = distance(progress.population[0])
trace.constant(progress.generation, gap, running.learning_rate)
constant.append(gap)
adam().run(loss, gradient=True, generations=steps, control=keep)
print(
f"the learning rate kept at {RATE}, after the same {steps} steps: within "
f"{scientific(constant[-1], 1)} of the answer"
)
trace.write()
What it prints, from a seeded run:
Smoothing 100000 noisy points: the misfit plus 50 times the roughness
Adam from a flat curve, its learning rate 0.05 halved every 500 steps
step learning rate loss largest gradient distance to the answer
0 5.0e-2 1.0544890e5 6.4e0 1.3e0
100 5.0e-2 5.0613069e4 2.5e0 2.4e-2
200 5.0e-2 5.0612759e4 3.3e0 9.5e-3
300 5.0e-2 5.0615543e4 3.9e0 1.2e-2
400 5.0e-2 5.0619309e4 4.8e0 1.3e-2
500 5.0e-2 5.0621841e4 4.6e0 1.4e-2
600 2.5e-2 5.0609073e4 2.0e-2 5.8e-5
700 2.5e-2 5.0609073e4 1.2e-4 3.3e-7
800 2.5e-2 5.0609073e4 6.3e-7 1.8e-9
900 2.5e-2 5.0609073e4 3.1e-9 1.0e-11
925 2.5e-2 5.0609073e4 8.5e-10 2.5e-12
converged after 925 steps and 926 evaluations: every value within 2.5e-12 of the answer
the learning rate kept at 0.05, after the same 925 steps: within 1.3e-2 of the answer