The guide

Training

The spec is immutable and the state is yours: parameters, steps, optimizers, what-ifs by cloning, checkpoints, and the curves.

Spec and state#

The architecture is an immutable value. The weights are yours. Training does not version the graph.

When live weights live inside the graph, a training step has to produce a new graph, and then every name has to say which generation it binds to. Topos splits the two: the Network is structure, shapes, record-site initials, and input defaults — runnable standalone, with no live weights and no lock; the Parameters are a caller-owned table of payloads, born from those initials or from a checkpoint, passed into every run and every plan.

A training step is then a pure data transform of the table. No new spec, no generation. Cloning the state is honest and costs what the weights cost; that is the whole price of a what-if. Optimizer moments are more tables of the same shape.

network.describe() unchanged after 100 steps: true
network.parameters().of(w)   0   (the recorded initial, materialized fresh)
parameters.of(w)             1.9999999993640727   (the table the loop stepped)
parameters.clone()           a what-if, costing 1 payloads

Plate 1. Training does not version the graph. After a hundred steps the network prints the same lines and still answers the initial parameters; the loop's own table is what moved.

The code that produced it

use topos::{Detach, Tape, Tensor};

let (network, [w, x, y, loss]) = Tape::record(|tape| {
    let w = tape.parameter(0.0_f64);
    let x = tape.input(0.0);
    let y = tape.input(0.0);
    let error = w * x - y;
    [w, x, y, error * error].detach()
});
let before = network.describe();
let mut parameters = network.parameters();
let samples = [(1.0, 2.0), (2.0, 4.0), (3.0, 6.0)];
for step in 0..100 {
    let (sample_x, sample_y) = samples[step % samples.len()];
    let run = network.forward(&parameters, [(x, sample_x.into()), (y, sample_y.into())]);
    let gradients = run.backward(loss).parameters(&parameters);
    parameters = parameters.step(&gradients, |w, g| {
        w.clone() - g.clone() * Tensor::from(0.02)
    });
}
assert_eq!(network.describe(), before);
show::lines([
    format!("network.describe() unchanged after 100 steps: {}", network.describe() == before),
    format!("network.parameters().of(w)   {}   (the recorded initial, materialized fresh)", network.parameters().of(w)),
    format!("parameters.of(w)             {}   (the table the loop stepped)", parameters.of(w)),
    format!("parameters.clone()           a what-if, costing {} payloads", parameters.len()),
])

The loop#

let mut parameters = network.parameters();
for step in 0..steps {
    let run = network.forward(&parameters, feeds(step));
    let gradients = run.backward(loss).parameters(&parameters);
    parameters = parameters.step(&gradients, |parameter, gradient| {
        parameter.clone() - gradient.clone() * learning_rate.broadcast_like(gradient)
    });
}

Field::parameters projects a backward pass onto the parameter slots, so the update is O(parameters) and never touches a graph-sized buffer. step calls the rule once per slot, in recording order, and answers a new table; step_each also hands the rule each parameter's symbol, for per-parameter policy. The learning rate is a per-step argument spread over each gradient explicitly, because nothing broadcasts by accident.

For compiled training, record the derivative and lower a forward-only plan over the adjoints' roots; then each step is one plan.forward and a recorded_gradients read, and no backward pass executes at all. The walkthrough does this on its graph.

Learning rates are clones#

One spec, several states. Each learning rate below descends from a clone of the same initial table; nothing about the network is copied.

Plate 2. Three learning rates on one spec, each descending from a clone of the same initial state. On a log scale a constant convergence rate is a straight line.

The code that produced it

use malevich::{Line, Plot};
use topos::{Detach, Tape, Tensor};

let (network, (w, b, loss)) = Tape::record(|tape| {
    let w = tape.parameter(0.0_f64);
    let b = tape.parameter(0.0);
    // Three samples of the line y = 2x + 1, as leaves; the total
    // loss is the sum of the squared errors.
    let mut loss = None;
    for (x, y) in [(1.0, 3.0), (2.0, 5.0), (3.0, 7.0)] {
        let error = w * tape.leaf(x) + b - tape.leaf(y);
        let squared = error * error;
        loss = Some(match loss {
            Some(total) => total + squared,
            None => squared,
        });
    }
    (w, b, loss.expect("at least one sample")).detach()
});
let initial = network.parameters();

let mut plot = Plot::new()
    .title("one spec, three learning rates")
    .x_label("step")
    .y_label("loss")
    .log_y();
for learning_rate in [0.005, 0.02, 0.05] {
    // Cloning the state is all a what-if costs.
    let mut parameters = initial.clone();
    let mut losses = Vec::new();
    for _ in 0..300 {
        let run = network.forward(&parameters, []);
        losses.push(run.of(loss).scalar());
        let gradients = run.backward(loss).parameters(&parameters);
        parameters = parameters.step(&gradients, |parameter, gradient| {
            parameter.clone() - gradient.clone() * Tensor::from(learning_rate)
        });
    }
    let label = format!(
        "lr {learning_rate}: w {:.2}, b {:.2}",
        parameters.of(w).scalar(),
        parameters.of(b).scalar()
    );
    plot = plot.layer(Line::y(losses).label(label));
}
show::plot(plot)

Parallel what-ifs work the same way: clone the table per thread, share the network by reference. The sealed spec has no lock to fight over, and feeds are run state, so gradient_descent.rs trains three states on three threads over one network, and makemore_mlp_parallel.rs shards a minibatch across threads and sums the shard gradients as tables.

Optimizers are instruments#

An Optimizer is an open trait: gradients and parameters in, next parameters out, with whatever state it needs held as more parameter-aligned tables. Sgd is stateless; Adam and AdamW carry moments. A custom optimizer is an ordinary implementation with the same standing as the built-ins, and a comparison loop can iterate them side by side.

Plate 3. Plain SGD against Adam on one small regression, from the same initial state and the same recorded gradients. The optimizer is a caller-owned instrument the loop hands the gradients to.

The code that produced it

use malevich::{Line, Plot};
use topos::{Activation, Adam, Mlp, Module, Optimizer, Sgd, Shape, Tape, Tensor, init};

// A tanh MLP fit to sin x over 48 samples, recorded once.
let features: Tensor<f32> = init::uniform(3, 3.0)(&Shape::new([48, 1]));
let targets: Vec<f32> = features.to_vec().iter().map(|x| x.sin()).collect();
let tape: Tape<f32> = Tape::new();
let mlp = Mlp::new(&tape, &[1, 16, 1], Activation::Tanh, init::xavier(7));
let error = mlp.express(tape.input(features)) - tape.input(Tensor::new([48, 1], targets));
let loss = (error * error).sum();
let adjoints = tape.differentiate(loss, mlp.parameters());
let loss = loss.symbol();
let network = tape.into_network();
let plan = network.entry(adjoints.roots()).lower();

// The loop owns the state and the learning rate; the optimizer is
// the strategy it hands the gradients to.
let mut curves = Vec::new();
let mut sgd = Sgd;
let mut adam = Adam::new(Tensor::from(0.9_f32), Tensor::from(0.999), Tensor::from(1e-8));
let runs: [(&str, &mut dyn Optimizer<f32>, f32); 2] =
    [("sgd, lr 0.003", &mut sgd, 0.003), ("adam, lr 0.005", &mut adam, 0.005)];
for (label, optimizer, learning_rate) in runs {
    let learning_rate = Tensor::from(learning_rate);
    let mut parameters = network.parameters();
    let mut losses = Vec::new();
    for _ in 0..400 {
        let run = plan.forward(&parameters, []);
        losses.push(run.of(loss).scalar());
        let gradients = run.recorded_gradients(&adjoints);
        parameters = optimizer.step(&parameters, &gradients, &learning_rate);
    }
    curves.push((label, losses));
}
let mut plot = Plot::new()
    .title("two optimizers, one recorded gradient")
    .x_label("step")
    .y_label("sum of squared errors")
    .log_y();
for (label, losses) in curves {
    plot = plot.layer(Line::y(losses).label(label));
}
show::plot(plot)

Hyperparameters are caller-owned and visible at the call site: Adam::new takes its betas and epsilon as payloads you write, and the learning rate stays an argument of every step. A facade never chooses a float constant.

Feeds are the data#

Input shapes are fixed at record time, and the batch arrives as a feed. A graph recorded over a [2, 2] input trains on any number of two-sample minibatches without ever growing the tape — and rasterizes its learned surface through the very same two-row expression afterwards, two grid cells per run.

Plate 4. XOR learned by a [2, 4, 1] perceptron fed two samples per run, and the surface it learned, rasterized through the very same two-row expression training used.

The code that produced it

use malevich::{Cells, Plot};
use topos::{Activation, Mlp, Module, Tape, Tensor, init};

let tape: Tape<f32> = Tape::new();
let mlp = Mlp::new(&tape, &[2, 4, 1], Activation::Tanh, init::uniform(7, 0.5));
// The defaults only fix the shapes: two samples of two features.
let x = tape.input(Tensor::filled([2, 2], 0.0_f32));
let y = tape.input(Tensor::filled([2, 1], 0.0));
let predicted = mlp.express(x);
let error = predicted - y;
let loss = (error * error).sum();
let (x, y, loss, predicted) = (x.symbol(), y.symbol(), loss.symbol(), predicted.symbol());
let network = tape.into_network();
let mut parameters = network.parameters();

let minibatches = [
    (Tensor::new([2, 2], [0.0, 0.0, 0.0, 1.0]), Tensor::new([2, 1], [-1.0, 1.0])),
    (Tensor::new([2, 2], [1.0, 0.0, 1.0, 1.0]), Tensor::new([2, 1], [1.0, -1.0])),
];
let learning_rate = Tensor::from(0.05_f32);
for step in 0..3000 {
    let (batch_x, batch_y) = &minibatches[step % 2];
    let run = network.forward(&parameters, [(x, batch_x.clone()), (y, batch_y.clone())]);
    let gradients = run.backward(loss).parameters(&parameters);
    parameters = parameters.step(&gradients, |parameter, gradient| {
        parameter.clone() - gradient.clone() * learning_rate.broadcast_like(gradient)
    });
}

// The surface, two grid cells per run through the same expression.
let (columns, rows) = (24, 12);
let centers: Vec<(f32, f32)> = (0..rows)
    .flat_map(|row| (0..columns).map(move |column| ((column as f32 + 0.5) / columns as f32, (row as f32 + 0.5) / rows as f32)))
    .collect();
let mut surface = Vec::new();
for pair in centers.chunks(2) {
    let feed = Tensor::new([2, 2], [pair[0].0, pair[0].1, pair[1].0, pair[1].1]);
    surface.extend(network.forward(&parameters, [(x, feed)]).of(predicted).to_vec());
}
show::plot(
    Plot::new()
        .layer(Cells::matrix(columns, surface).extents((0.0, 1.0), (0.0, 1.0)))
        .colorbar()
        .title("the learned xor surface"),
)

Twins#

When the shape you train on is not the shape you want to read, record both expressions of the same parameters on one tape: a batch-shaped one for training and a grid-shaped or single-row twin for reading. The twin shares every parameter symbol, so one forward with the trained table answers it, and an entry that declares only the training root never schedules the twin.

Plate 5. A tanh MLP fit to a noisy sine. The fit line is the trained expression itself, fed an even grid in place of the training samples.

The code that produced it

use malevich::{Line, Plot, Points};
use topos::{Activation, Mlp, Module, Shape, Tape, Tensor, init};

let features: Tensor<f32> = init::uniform(3, 3.0)(&Shape::new([96, 1]));
let noise: Tensor<f32> = init::normal(5, 0.1)(&Shape::new([96, 1]));
let xs = features.to_vec();
let ys: Vec<f32> = xs.iter().zip(noise.to_vec()).map(|(x, n)| x.sin() + n).collect();

let tape: Tape<f32> = Tape::new();
let mlp = Mlp::new(&tape, &[1, 16, 1], Activation::Tanh, init::xavier(7));
let input = tape.input(features);
let predicted = mlp.express(input);
let error = predicted - tape.input(Tensor::new([96, 1], ys.clone()));
let loss = (error * error).sum();
let adjoints = tape.differentiate(loss, mlp.parameters());
let (input, predicted) = (input.symbol(), predicted.symbol());
let network = tape.into_network();

let plan = network.entry(adjoints.roots()).lower();
let learning_rate = Tensor::from(0.002_f32);
let mut parameters = network.parameters();
for _ in 0..1500 {
    let gradients = plan.forward(&parameters, []).recorded_gradients(&adjoints);
    parameters = parameters.step(&gradients, |parameter, gradient| {
        parameter.clone() - gradient.clone() * learning_rate.broadcast_like(gradient)
    });
}

// The fit: the trained expression fed an even grid instead of the
// samples. Same nodes, different feed.
let grid: Vec<f32> = (0..96).map(|i| -3.0 + 6.0 * i as f32 / 95.0).collect();
let run = network.forward(&parameters, [(input, Tensor::new([96, 1], grid.clone()))]);
let fit = run.of(predicted).to_vec();
show::plot(
    Plot::new()
        .layer(Points::xy(xs, ys).label("samples"))
        .layer(Line::function(-3.0..3.0, f64::sin).label("sin x"))
        .layer(Line::xy(grid, fit).label("fit"))
        .title("a tanh mlp fit to noisy sin x")
        .x_label("x"),
)

Plate 6. The two moons, separated: the decision surface of a [2, 16, 16, 1] MLP after two thousand full-batch steps, drawn by a grid-shaped twin of the training expression.

The code that produced it

use std::f32::consts::PI;

use malevich::{Cells, Plot, Points};
use topos::{Activation, Mlp, Module, Shape, Tape, Tensor, init};

// Two interleaved half-moons of a hundred noisy points each.
let noise: Tensor<f32> = init::normal(5, 0.1)(&Shape::new([200, 2]));
let noise = noise.to_vec();
let mut xs = Vec::new();
let mut ys = Vec::new();
let mut labels = Vec::new();
for moon in 0..2 {
    for index in 0..100 {
        let angle = PI * index as f32 / 99.0;
        let (x, y) = if moon == 0 {
            (angle.cos(), angle.sin())
        } else {
            (1.0 - angle.cos(), 0.5 - angle.sin())
        };
        let at = 2 * (moon * 100 + index);
        xs.push(x + noise[at]);
        ys.push(y + noise[at + 1]);
        labels.push(if moon == 0 { 1.0_f32 } else { -1.0 });
    }
}
let features: Vec<f32> = xs.iter().zip(&ys).flat_map(|(x, y)| [*x, *y]).collect();

let tape: Tape<f32> = Tape::new();
let mlp = Mlp::new(&tape, &[2, 16, 16, 1], Activation::Tanh, init::xavier(7));
let predicted = mlp.express(tape.input(Tensor::new([200, 2], features)));
let error = predicted - tape.input(Tensor::new([200, 1], labels));
let loss = (error * error).sum();
let adjoints = tape.differentiate(loss, mlp.parameters());

// A grid-shaped twin of the same parameters rasterizes the surface.
let (columns, rows) = (48, 16);
let (x_span, y_span) = ((-1.5_f32, 2.5_f32), (-1.0_f32, 1.5_f32));
let centers: Vec<f32> = (0..rows)
    .flat_map(|row| {
        (0..columns).flat_map(move |column| {
            [
                x_span.0 + (column as f32 + 0.5) / columns as f32 * (x_span.1 - x_span.0),
                y_span.0 + (row as f32 + 0.5) / rows as f32 * (y_span.1 - y_span.0),
            ]
        })
    })
    .collect();
let surface = mlp.express(tape.input(Tensor::new([columns * rows, 2], centers))).symbol();
let network = tape.into_network();

let plan = network.entry(adjoints.roots()).lower();
let learning_rate = Tensor::from(0.0003_f32);
let mut parameters = network.parameters();
for _ in 0..2000 {
    let gradients = plan.forward(&parameters, []).recorded_gradients(&adjoints);
    parameters = parameters.step(&gradients, |parameter, gradient| {
        parameter.clone() - gradient.clone() * learning_rate.broadcast_like(gradient)
    });
}
let values = network.forward(&parameters, []).of(surface).to_vec();
let (upper_x, lower_x) = xs.split_at(100);
let (upper_y, lower_y) = ys.split_at(100);
show::plot(
    Plot::new()
        .layer(Cells::matrix(columns, values).extents(
            (f64::from(x_span.0), f64::from(x_span.1)),
            (f64::from(y_span.0), f64::from(y_span.1)),
        ))
        .layer(Points::xy(upper_x.to_vec(), upper_y.to_vec()).label("class +1"))
        .layer(Points::xy(lower_x.to_vec(), lower_y.to_vec()).label("class -1"))
        .colorbar()
        .title("the learned decision surface"),
)

The language models in the gallery use the same idiom at scale: a full-context expression and a one-token decode step, on one tape, sharing 124M parameters.

Checkpoints#

A checkpoint is pure state: snapshot and restore copy a module's parameter payloads in visit order, and named_snapshot and named_restore do it under structured paths — the form that survives code evolution and maps onto a foreign layout. No graph is touched either way. Loading GPT-2's released weights is one named_restore over the paths the model tree announces itself.

0.weights [3, 4]
0.bias [4]
1.weights [4, 2]
1.bias [2]

restored == original, payload for payload: true

Plate 7. A checkpoint is pure state. named_snapshot walks a module's parameters under structured paths; restore builds a new table and touches no graph.

The code that produced it

use topos::{Activation, Mlp, Tape, checkpoint, init};

let tape: Tape<f32> = Tape::new();
let mlp = Mlp::new(&tape, &[3, 4, 2], Activation::Relu, init::xavier(11));
let network = tape.into_network();
let parameters = network.parameters();

// A named snapshot is the module's parameters under structured
// paths; a restore builds a new table from one.
let snapshot = checkpoint::named_snapshot(&parameters, &mlp);
let restored = checkpoint::named_restore(&parameters, &mlp, snapshot.clone());
let mut lines: Vec<String> = snapshot
    .iter()
    .map(|(path, payload)| format!("{path:<12} {}", payload.shape()))
    .collect();
lines.push(String::new());
lines.push(format!(
    "restored == original, payload for payload: {}",
    restored.payloads() == parameters.payloads()
));
show::text(lines.join("\n"))

Where next#

Entries and plans — compiled training, and what a plan keeps.
The neural tier — the layers a training loop is built from.
Spec and state — the argument, and recurrence as feeds, what it opened.