2.4. Training, One State Transition At A Time
Training repeatedly transforms parameters, optimizer memory, data order, and runtime state. TorchLean's high-level trainer packages that transition, while the lower manual API exposes each piece.
We will train the running 2\to8\to1 MLP and then unpack what happened.
2.4.1. The Smallest Complete Run
Execute:
lake exe torchlean quickstart_mlp \ --device cpu --steps 200 --seed 2026
The current checkout reports:
dataset size = 25 mean_loss(before) = 0.761530 mean_loss(after) = 0.003234 heldout x=(0.25,-0.75), target=0.2, prediction(after)=[0.210239]
The source is
SimpleMlpTrain.lean.
Its program structure is:
model builder + trainer configuration + dataset + train options -> trained handle
The loss values are measurements from this run. The model and optimizer definitions, by contrast, are reusable objects that can appear in theorem statements or another runtime profile.
2.4.2. Declare The Architecture
import NN.API
open TorchLean
def model :
nn.M (nn.Sequential (shape![2]) (shape![1])) :=
nn.Sequential![
nn.linear 2 8,
nn.relu,
nn.linear 8 1
]
At this point we have:
-
input and output shapes;
-
layer structure;
-
parameter shapes;
-
seeded initialization actions.
We do not yet have a loss, optimizer, concrete parameter values, or device.
2.4.3. Attach A Training Problem
def trainer (seed : Nat) :=
Trainer.new model
{ task := .regression
optimizer := optim.adam { lr := 0.03 }
dtype := .float
backend := .eager
seed := seed }
Trainer.new runs the seeded builder and stores persistent choices. It still does not consume data
or update a parameter.
For regression, the default objective is mean-squared error. If a prediction and target each have
n entries:
L(\theta;x,y)
=
\frac1n\sum_{i=1}^{n}
\left(F_\theta(x)_i-y_i\right)^2.
Changing .regression to .crossEntropy changes the objective and target convention without
changing the architecture. A custom task supplies a checked scalar loss program.
2.4.4. Build The Dataset
The quickstart uses a deterministic grid. A four-point example can be written directly:
def xs : Tensor.T Float (shape![4, 2]) :=
tensor! [
[0.0, 0.0],
[0.0, 1.0],
[1.0, 0.0],
[1.0, 1.0]
]
def ys : Tensor.T Float (shape![4, 1]) :=
tensor! [[0.0], [1.0], [1.0], [0.0]]
def data : Trainer.Dataset (shape![2]) (shape![1]) :=
Data.tensorDataset xs ys
The dataset type matches the model map [2] → [1]. A batched model would require a batched dataset
with a leading dimension in both item shapes.
2.4.5. Call Train And Keep The Result
def run : IO Unit := do
let trained ← (trainer 2026).train data
{ steps := 200
batchSize := 4
logEvery := 25 }
trained.printSummary
let heldout : Tensor.T Float (shape![2]) :=
tensor! [0.25, -0.75]
let yhat ← trained.predict heldout
IO.println s!"prediction={Tensor.pretty yhat}"
The returned handle retains:
-
final parameters;
-
runtime model state;
-
before/after summary;
-
prediction and verification closures.
Prediction accepts Float tensors and performs the runtime scalar conversion selected by the trainer. It does not rebuild or reinitialize the model.
2.4.6. What One Update Computes
Let \theta_t be the parameter pack at update t. A plain SGD update is:
g_t=\nabla_\theta L(\theta_t;x_t,y_t),
\qquad
\theta_{t+1}=\theta_t-\eta g_t.
Every symbol is structured:
-
\theta_tis a shape-indexed pack of differently shaped tensors sharing scalar typeα; -
g_thas exactly the same pack structure; -
Lis a scalar tensor program; -
the subtraction and scaling occur coordinatewise in the selected scalar semantics.
An optimizer is therefore not merely a function from a flat vector to a flat vector. It owns shape-aligned state.
2.4.7. Adam's Hidden State Is Explicit
Adam maintains first and second moment estimates:
m_t=\beta_1m_{t-1}+(1-\beta_1)g_t,
v_t=\beta_2v_{t-1}+(1-\beta_2)g_t^2.
With bias correction:
\widehat m_t=\frac{m_t}{1-\beta_1^t},
\qquad
\widehat v_t=\frac{v_t}{1-\beta_2^t},
and the parameter update is:
\theta_{t+1}
=
\theta_t-\eta
\frac{\widehat m_t}{\sqrt{\widehat v_t}+\epsilon}.
The two moment packs have the same dependent tensor shapes as \theta. The step counter matters
because it changes the bias correction. Restoring only parameters from a checkpoint but not Adam
state is not the same continuation of training.
TorchLean also provides SGD, momentum SGD, AdamW, AdaGrad, RMSProp, Adadelta, and Muon-related runtime configuration. Their constructors share one trainer interface; their state and laws remain optimizer-specific.
2.4.8. What Does steps Count?
For the unbatched model [2] → [1], this configuration:
steps := 200 batchSize := 4
means 200 optimizer updates. Each update consumes four samples, differentiates them at the same parameter point, and averages their gradient packs. Logging reports the mean pre-update loss from those same forward tapes. The trainer does not run an extra forward pass merely to print that loss. Before building those tapes, it advances mutable model buffers once per item; trainable parameters stay fixed until the mean gradient is ready. This matters for dropout, BatchNorm, and any operation whose tape retains data needed by backward: the displayed scalar, saved state, and gradient come from the same step.
For a true vectorized minibatch, define:
def batchedModel :
nn.M
(nn.Sequential
(shape![2, 2])
(shape![2, 1])) :=
nn.Sequential![
nn.linear 2 8,
nn.relu,
nn.linear 8 1
]
def batchedData :
Trainer.Dataset (shape![2, 2]) (shape![2, 1]) :=
Data.batchDataset 2 data
(shuffle := true)
(seed := 2026)
Each item in batchedData holds two samples in one tensor. With the default
TrainOptions.batchSize := 1, an update consumes one item and therefore runs one vectorized
forward/backward operation over two samples. The four-row example produces two full items; asking
Data.batchDataset for groups of five would produce none because typed batching drops incomplete
groups. Setting TrainOptions.batchSize above one would accumulate gradients across several of
these tensor minibatches. Vectorization can change reduction order and performance; it is not an
optimization flag applied to the first model.
Run the maintained version:
lake exe torchlean quickstart_minibatch_mlp \ --device cpu --batch 5 --steps 5 --seed 2026
The printed model confirms:
Sequential: [5, 2] -> [5, 1], layers=3, params=33
The parameter count remains 33 because linear layers preserve the batch prefix; the batch axis does not create separate weights per sample.
2.4.9. Eager And Compiled Are Execution Choices
Compare:
lake exe torchlean quickstart_mlp \ --device cpu --backend eager --steps 20 --seed 2026 lake exe torchlean quickstart_mlp \ --device cpu --backend compiled --steps 20 --seed 2026
Eager mode records operations and local reverse rules as they execute. Compiled mode builds a typed forward/derivative graph and replays it with current parameters and inputs.
The trainer method remains train; compilation is a property of the configured runner, analogous to
wrapping a PyTorch model with an execution transform rather than renaming the model's forward
method.
Compiled trainer execution is currently CPU-only. A non-CPU compiled request is rejected rather than silently falling back to a different semantics.
2.4.10. Device Selection Is A Profile
The execution profile records more than a device enum:
-
target device and platform;
-
provider preference;
-
numerical and assurance policy;
-
forward and VJP ownership;
-
available kernel capsules.
Programmatically:
def cudaRun : Trainer.RunConfig :=
({ optimizer := optim.adam { lr := 0.03 }
dtype := .float
backend := .eager } :
Trainer.RunConfig).cuda
def cudaTrainer :=
Trainer.new model
(Trainer.Config.fromRunConfig
cudaRun .regression
(seed := 2026))
A maintained CUDA run requires a CUDA-enabled build:
lake -R -K cuda=true exe torchlean quickstart_mlp \ --device cuda --steps 20 --seed 2026 --show-backend
--show-backend prints the selected capsules as operations first execute. This is how a run reports
which provider owned matmul, ReLU, loss, and their VJPs.
Named future devices such as Metal, ROCm, TPU, or Trainium may be represented in configuration, but
withDevice fails if the current build has no maintained profile. A name in an enum is not an
implementation.
2.4.11. Scalar Semantics Are Part Of The Run
The common executable selections are:
--dtype float --dtype ieee754exec
Host Float is fast and relies on the platform runtime. IEEE32Exec is TorchLean's bit-level
binary32 reference and is much slower, but exposes precise finite and exceptional behavior.
Proof-level Real and rounded-real FP32 are not executable trainer choices. FP32 has binary32
precision and gradual-underflow parameters, but no upper exponent bound, NaN, infinity, or signed
zero; bridge theorems therefore require finite/no-overflow hypotheses when relating it to
IEEE32Exec.
A loss curve without its scalar semantics is incomplete. The same architecture and seed may round differently in binary32, binary64, a fused CUDA kernel, or an external provider.
2.4.12. A Checkpoint Is A Particular Slice Of State
TorchLean's JSON parameter checkpoint records each host Float with Float.toBits, so saving and
loading does not pass through a decimal approximation. The expected parameter shapes come from the
model. Each expected tensor must be present with the right shape and scalar count before it becomes
a parameter pack:
def loadForThisModel (path : System.FilePath) := Checkpoint.loadModelParamBits (nn.run 2026 model) path
The result is an IO action returning tensors whose dependent shape list is exactly
nn.paramShapes (nn.run 2026 model). Checkpoint.toRuntimeParams turns such a checked pack into
runtime parameter handles, while Checkpoint.loadModuleParams and
Checkpoint.saveModuleParams work with an already instantiated runtime module.
The loader requires an exact tensor manifest. Missing tensors, extra tensors, shape mismatches, and malformed scalar payloads are rejected before any parameter is installed. This matters when two models share an initial parameter prefix: a checkpoint for the larger model cannot be accepted as a checkpoint for the smaller one merely because the first few shapes agree.
Classification and cross-entropy training also wire the loadParams? and saveParams? fields of
Trainer.TrainOptions into this exact-bits format. For example, the following options ask that path
to save parameters after 200 updates:
def saveClassifierParams : Trainer.TrainOptions :=
{ steps := 200
logEvery := 25
saveParams? := some "artifacts/classifier.params.json" }
The current high-level regression and custom-loss training paths do not consume those two fields;
use the direct Checkpoint/TorchLean.Checkpoint helpers with a manual module for those tasks.
This limitation is worth spelling out because accepting a shared options record is not evidence
that every task dispatch implements every optional field.
Even on the classification path, this is a parameter checkpoint, not a complete training
snapshot. The eager CUDA runtime can save Adam or AdamW moments and step counters separately with
Checkpoint.saveOptimizerState, then restore them with
Checkpoint.loadOptimizerState. That binary file records the optimizer kind, the
moment-defining hyperparameters, every parameter shape, and the requiresGrad mask. Loading rejects
a different optimizer configuration or parameter schema instead of silently attaching moments to
the wrong model. Integer metadata and float32 payloads use explicit little-endian encodings, and a
save is written to a fresh sibling file before it replaces the destination.
That optimizer file is still not a complete training snapshot. Replaying the next batch also needs the loader or stream position; stochastic layers need generator state; and interpreting the result needs the model, preprocessing, scalar semantics, backend profile, and device. A parameter-only checkpoint remains appropriate for inference or a fresh optimizer run. Pairing it with native optimizer state resumes more of an Adam trajectory, but only the state explicitly present in those two files.
2.4.13. Manual Training
The high-level trainer is intended for common runs. Trainer.Manual exposes:
stepper step trainMode evalMode callbacks loader loops prediction
Use it when the program needs a custom accumulation policy, custom scheduling, multiple losses, generated batches, reinforcement-learning interaction, or detailed instrumentation.
Trainer.Manual.StepBatchStream α shapes supplies already collated tensors as a function of the
step. PINN collocation points and simulator batches naturally fit this interface.
The lower API does not change the model or autograd semantics. It exposes the runner state that the high-level trainer normally manages.
2.4.13.1. Learning-rate schedules
The manual configuration has constant, step-decay, exponential, and warmup-cosine schedules. Here
a rate of 0.1 is halved after every three step indices:
def decay := _root_.TorchLean.Trainer.Scheduler.step 0.1 3 0.5 #eval (List.range 10).map (_root_.TorchLean.Trainer.Scheduler.lrAt decay)
Lean prints:
[0.100000, 0.100000, 0.100000, 0.050000, 0.050000, 0.050000, 0.025000, 0.025000, 0.025000, 0.012500]
Attach the same schedule to a manual step configuration with Trainer.stepLR:
def scheduledManualRun :=
Trainer.stepLR
(Trainer.steps 10 (optim.sgd { lr := 0.1 }))
0.1 3 0.5
The counter is zero-indexed, which is why the first decay occurs at index three. A step schedule
with stepSize := 0 deliberately stays at its base rate instead of dividing by zero. The compact
high-level Trainer.TrainOptions does not currently carry a scheduler; use this manual
configuration when schedule state must be part of the loop.
Transformer runs commonly warm up from a small rate and then decay toward a nonzero floor. This
configuration reaches 0.0006 after 2,000 updates and follows a cosine curve toward 0.00006:
def pretrainingSchedule :=
_root_.TorchLean.Trainer.Scheduler.warmupCosine 0.0006 0.00006 2000 162761
def scheduledPretraining :=
Trainer.warmupCosineLR
(Trainer.steps 162761 (optim.adamw { lr := 0.0006 }))
0.0006 0.00006 2000 162761
The first update uses peak / warmupSteps; the last warm-up update reaches peak. At and after
totalSteps, lrAt returns the floor exactly. The scheduler changes optimizer state only. It does
not depend on a particular model, loss, dataset, or device. If the requested warm-up is longer than
the run, TorchLean clamps it to totalSteps.
2.4.14. Four Useful Experiments
2.4.14.1. Initialization only
lake exe torchlean quickstart_mlp \ --device cpu --steps 0 --seed 2026
This isolates initialization and the initial loss.
2.4.14.2. Seed sensitivity
Run 20 steps with seeds 2026, 2027, and 2028. Keep the dataset order fixed if you want to
study only initialization.
2.4.14.3. Optimizer sensitivity
In the quickstart source, replace Adam with SGD while keeping the model, seed, and steps fixed. Compare both the initial and final losses; the initial values should agree when initialization and data order agree.
2.4.14.4. Backend report
Add --show-backend. On CPU, inspect the reference capsules. On a CUDA build, inspect which native
capsules are selected and whether any trusted external provider appears.
2.4.15. What Training Establishes
A successful run establishes that one configured pipeline executed:
data -> forward -> loss -> reverse pass -> optimizer -> parameters
It can produce valuable evidence:
-
loss and prediction traces;
-
exact parameter artifacts;
-
capsule audit rows;
-
reproducible configuration;
-
runtime errors or successful completion.
It does not automatically prove:
-
convergence for all initializations;
-
generalization to unseen data;
-
robustness to an input region;
-
equality of eager, compiled, CUDA, and LibTorch paths;
-
correctness of every native instruction.
TorchLean gives the run enough structure for a theorem, numerical bound, backend contract, or verification certificate to refer to the same model without erasing the boundary between them.