MLP Training🔗

From the repository root, train the model on the CPU for 200 updates. The fixed seed makes this run reproducible under the same arithmetic and execution settings:

# Fix the seed and update count for the MLP trace
# interpreted below.
lake exe torchlean quickstart_mlp \
  --device cpu --steps 200 --seed 2026

== Quickstart: simple MLP training (seed=2026, steps=200) ==
target(heldout)    = [0.200000]
untrained(heldout) = [-0.088261]
dataset size = 25
mean_loss(before training) = 0.495227
step 0: loss=0.250488
step 25: loss=0.586318
step 50: loss=0.847444
step 75: loss=0.003933
step 100: loss=0.023397
step 125: loss=0.069799
step 150: loss=0.000061
step 175: loss=0.029500
mean_loss(after training) = 0.002402
steps=200 arithmetic=native scalar=Float32 loss=0.495227 -> 0.002402
trained(heldout) = [0.228325]

The source is SimpleMlpTrain.lean. Its program structure is:

-- Construction, data, and per-run options meet at the call
-- that returns the trained snapshot.
model builder
  + trainer configuration
  + dataset
  + train options
  -> trained result

The log separates individual training losses from the mean loss over the dataset. Start with these entries:

• dataset size = 25 is the number of items the trainer got after materializing the dataset, not the number of rows in a file. The data chapter is where that number comes from.

• mean_loss(before training) and mean_loss(after training) are the only two full-dataset measurements in the run. The trainer takes them once each, at the boundaries.

• The step k: loss=... lines are single-sample losses, measured on the tape that produced that step's gradient. That is why they are not monotone: 0.586318 at step 25 and 0.847444 at step 50 use different samples at different parameter points, so their increase alone does not establish that training diverged. The full-dataset mean fell from 0.495227 to 0.002402 while those numbers bounced around.

• arithmetic=native scalar=Float32 records which of TorchLean's executable semantics ran. A loss curve without that label is ambiguous, and we say why below.

• untrained(heldout) and trained(heldout) are the same held-out input (0.25, -0.75) before and after: -0.088261 moves to 0.228325 against a target of 0.200000.

Model Architecture🔗

Lean reports both the seeded construction stage and the model's input/output shapes:

tfModel : nn.Builder (nn.Sequential [2] [1])

nn.Builder waits for a seed to construct the nn.Sequential [2] [1] model. Its architecture is already fixed; the loss, optimizer, and device are still separate choices.

The intermediate width eight belongs to the architecture even though only [2] and [1] appear at the outer boundary. It determines both the number of hidden activations for each sample and the shapes of the two weight matrices. The builder delays choosing their values until a seed is supplied. This lets the same architecture be initialized repeatedly for the seed comparison below without rewriting its forward computation or changing what counts as a well-shaped input.

Objective And Optimizer Configuration🔗

Trainer.new runs the seeded builder and stores persistent choices. It still does not consume data or update a parameter. Its summary shows the initialized model's layer shapes and state layout:

model:
Sequential: [2] -> [1], layers=3, params=33, state=33
  [0] Linear(2, 8): [2] -> [8] params=24, state=24 [[8, 2], [8]]
  [1] ReLU: [8] -> [8] params=0, state=0 []
  [2] Linear(8, 1): [8] -> [1] params=9, state=9 [[1, 8], [1]]

Thirty three numbers: 8\times2 weights and 8 biases in the first affine layer, then 1\times8 weights and 1 bias in the second. params counts trainable scalars and state counts everything the checkpoint stores, so a model with normalization buffers would show the two numbers disagreeing.

For regression, the default objective is mean-squared error. For a prediction and target with n entries, where n>0, the loss is:

L(\theta;x,y) = \frac1n\sum_{i=1}^{n} \left(F_\theta(x)_i-y_i\right)^2.

Changing .meanSquaredError to .oneHotCrossEntropy axis changes the objective and target convention without changing the architecture. The zero-based axis may name any output dimension; Lean rejects an axis outside the output shape. A custom objective supplies a checked scalar loss program.

Here the output has shape [1], so the mean contains one squared residual. In the batched version, [5, 1] contains five residuals and the same objective averages all five. This reduction is part of the loss definition: it fixes the scale of the gradients passed to Adam. Changing a sum to a mean would therefore change more than the number printed in the log, even if the model and target tensors were otherwise identical.

The optimizer can print its configuration in the syntax you would use to construct it:

"optim.adam { learningRate := 2.9999999999999999e-2, beta1 := 9.0000000000000002e-1, beta2 := 9.9900000000000000e-1, epsilon := 1.0000000000000000e-8 }"

The moment coefficients and epsilon match PyTorch's defaults. The corresponding PyTorch configuration prints:

lr, betas, eps = 0.03 (0.9, 0.999) 1e-08

FloatLib supplies the scientific decimal text. describe checks it with Lean's literal decoder and compares the reconstructed binary64 bits before returning it. When the check fails, it emits a Float.ofBits expression instead. This protects small stabilizers and coefficients near one from changing when their descriptions are copied.

Adam is Kingma and Ba (2015); the single-sample step underneath it is older than deep learning and goes back to Robbins and Monro (1951).

Initial Parameter Values🔗

Before any data arrives, building the model has fixed all 33 parameter values. The first loss depends on these values together with the data, objective, mode, and arithmetic. Reading the initial state separates those choices from the later optimizer updates:

input weights = [[-0.337801, 0.526946],
 [-0.767575, 0.519218],
 [0.116565, 0.347734],
 [-0.242321, 0.244720],
 [-0.419680, 0.002862],
 [-0.461676, -0.360324],
 [-0.510277, 0.741238],
 [-0.108099, 0.091942]]
input bias    = [0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000, 0.000000]
output bias   = [0.000000]

Two decisions are visible. Biases start at exactly zero. Weights are drawn from Xavier uniform initialization (Glorot and Bengio, 2010), whose bound for a 2\to8 layer is

0.774597

and every entry above sits inside \pm0.774597, the largest magnitude being 0.767575.

The four shapes identify which values belong to which computation. The first matrix forms eight weighted combinations of the two input coordinates; the first bias shifts those combinations before ReLU. The second matrix combines the eight resulting activations, and the final bias shifts the prediction. Zero biases do not make the initial model zero, because the two weight matrices already contain nonzero values. The negative initial held-out prediction is consistent with this layout.

PyTorch makes both decisions differently. Its Linear uses Kaiming uniform initialization (He et al., 2015) with bound 1/\sqrt{n_{\text{in}}} and gives the bias a random value from the same interval:

== default Linear(2,8) initialization, seed 2026 ==
weight max |.| = 0.705845
kaiming bound  = 0.707107
bias           = [0.582698, 0.278374, 0.009535, 0.614985, -0.093329,
                  0.385542, 0.295646, 0.047354]
xavier bound   = 0.774597

Different initialization explains why the two runs begin with different losses on the same architecture and seed: 0.495227 here against 0.510072 in PyTorch. The generators and initialization schemes are not matched, so equal seed values do not establish equal parameters. Where we do claim agreement with PyTorch, as in the round-trip chapter, the agreement is on imported weights and forward values rather than on initialization.

Training Dataset🔗

The quickstart uses a deterministic grid over [-1,1]^2 and labels it with a small piecewise linear function, so that the target is something a 2\to8\to1 ReLU network can actually represent:

y(x) = 0.8\,\mathrm{relu}(x_1+x_2) -0.4\,\mathrm{relu}(x_2-x_1) +0.2.

The mathematical coordinates x_1,x_2 correspond to Lean's zero-based x[0] and x[1]. The label function applies the two ReLUs independently, then combines them:

Data.Synthetic.squareGrid builds a 5\times5 grid as a Tensor Float [25, 2], and Tensor.mapLeading applies tfTarget to each of the 25 leading rows. Materializing the dataset shows what the trainer will actually see:

samples = 25
input  = [-1.000000, -1.000000]
target = [0.200000]

The corner (-1,-1) lands on the constant term: both ReLUs are zero there, so the label is exactly 0.2. That is the same dataset size = 25 the command-line transcript reports, and the same first row PyTorch computes from the same grid.

The labels are computed once from the grid and do not depend on the trainer's parameters. During training, only the model's attempt to reproduce those labels changes. The target's two ReLU terms also explain the architecture choice: a hidden unit can form each required linear combination before applying ReLU, and the output layer can combine them. Representability makes this a useful small regression problem; it does not tell Adam which representation to find or how many updates it needs.

The dataset type matches the model map [2]\to[1]. A batched model would require a batched dataset with a leading dimension in both item shapes, which we do further down. Everything about where samples come from, how they are shuffled, and what happens to a short final batch belongs to the data chapter. Fixing the grid here lets the experiments below vary initialization, optimizer, or update grouping while using the same samples.

Training And Results🔗

A complete run is one IO action. This checked example trains the model and evaluates the same held-out input used by the command-line program:

dataset size = 25
mean_loss(before training) = 0.495227
step 0: loss=0.250488
step 50: loss=0.847444
step 100: loss=0.023397
step 150: loss=0.000061
mean_loss(after training) = 0.002402
steps=200 arithmetic=native scalar=Float32 loss=0.495227 -> 0.002402
trained(heldout) = [0.228325]

The initial mean 0.495227, the shared logged steps, the final mean 0.002402, and the prediction [0.228325] agree with the command-line transcript. Only the logging interval changed, from 25 updates to 50. The checked block recomputes these values when the page builds.

The returned Trainer.Result retains:

• final parameters and buffers, readable as Float tensors through trained.state and writable to a checkpoint file through trained.save path;

• the runtime model state;

• the completed step count and numeric before/after losses;

• prediction and verification closures.

Trainer.load trainer path data restores a saved result for the same trainer; it evaluates data once so that the restored report carries the mean loss of the restored parameters rather than a copied number.

trained.report.loss.before and trained.report.loss.after are ordinary host Float values. The model still executes with the arithmetic selected by the trainer, including .ieee ; Session.meanLoss converts per-item losses to host Float before averaging; the result stores those before/after means. Callers can therefore compare, serialize, or plot losses directly without parsing printed backend values.

Prediction accepts Float tensors and performs conversion into the arithmetic representation selected by the trainer. It does not rebuild or reinitialize the model.

There is also an ownership boundary at the result. Session.finish freezes the model state used by the result's prediction and verification closures. Further steps or a later checkpoint load on the live session cannot silently change an already returned result. This matters when keeping several training milestones in memory: a reported loss and the model saved beside it must continue to refer to the same parameter snapshot while the experiment proceeds.

dataset size = 25
mean_loss(before training) = 0.495227
mean_loss(after training) = 0.125425
state shapes = [[8, 2], [8], [1, 8], [1]]
initial output bias = [0.000000]
trained output bias = [0.279321]

Gradient Descent Updates🔗

Let \theta_t be the parameter pack at update t, (x_t,y_t) the selected sample, and \eta the learning rate. 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_t is a shape-indexed pack of differently shaped tensors sharing element type α;

• g_t has exactly the same pack structure;

• L is a scalar tensor program;

• the subtraction and scaling occur coordinatewise in the selected arithmetic semantics.

An optimizer is therefore not merely a function from a flat vector to a flat vector. It owns shape-aligned state.

For this model, reverse mode returns gradients with shapes [[8, 2], [8], [1, 8], [1]], in the same order as the parameters. A bias gradient cannot accidentally be used as a weight gradient merely because both are stored in the same pack. The optimizer can traverse corresponding tensors and then corresponding scalar coordinates. That alignment is what allows the same update implementation to work for this MLP and for a model with a different number of layers or tensor ranks.

You can watch a single update in isolation. Session.step returns the loss it measured before applying the update, so feeding it the same sample twice shows the effect of exactly one step:

loss before update 1 = 0.250488
loss before update 2 = 0.110670

One Adam step on the corner sample cut its loss from 0.250488 to 0.110670. The first value matches step 0 of the longer run because the seed, sample, and arithmetic match. The second call measures the effect of the first update before applying another update of its own. These two values concern the corner sample; a step that helps it can increase losses on other samples. Dataset loss must be measured separately. For the stochastic approximation setting, see Robbins and Monro (1951).

Mixed-Dtype Indexed Training🔗

One tensor remains homogeneous, but a training problem need not use one element type everywhere. Language models are the standard case: learned parameters, logits, and the scalar loss use a floating-point type α, while token inputs and targets use bounded indices such as Fin vocab. The public objective API keeps that distinction in its type:

@Module.Objective.dataStep : {α β : Type} →
  [inst : Storage α] →
    [inst_1 : Storage β] →
      [inst_2 : Context α] →
        [Runtime.FromFloat α] →
          {stateShapes : List Shape} →
            {firstShape secondShape : Shape} →
              Module.Objective α β stateShapes [] [firstShape, secondShape] →
                optim.Optimizer → IO (Tensor β firstShape → Tensor β secondShape → IO Unit)

The outer IO action creates a reusable step function taking two β tensors. Those tensors contain tokens; the step updates module state and returns Unit. The floating-point type α belongs to the objective and to the values returned by loss evaluation. A character-level GPT uses this interface in CharGpt.lean:

-- Keep discrete token tensors at the data boundary and
-- floating-point logits at the prediction
-- boundary.
let trainStep ← module.dataStep <|
  optim.adamW { learningRate := train.training.learningRate }
let evaluateLoss ← module.dataLossEvaluator evalDef
let predict ← module.indexedPredictor model

inputTokens and targetTokens keep their discrete element type; the returned loss and logits use α. dataLossEvaluator creates a reusable evaluator when repeated loss measurements should share the trained state. These functions hide the heterogeneous runtime context and its session lifecycle. TensorPack still exists as the typed internal representation of differently shaped parameters and graph inputs, but normal indexed-model code does not construct or deconstruct it.

The distinction between α and β is computational as well as notational. A token chooses a row of an embedding table; the chosen row contains floating-point parameters that can receive gradients. The integer identifying that row is not itself a floating-point parameter to optimize. Returning Unit from the step makes its purpose explicit: the observable effect is the updated module state. Use the loss evaluator when the caller also needs a scalar measurement, or the predictor when it needs the full output tensor.

Training State🔗

The next Adam update depends on its accumulated moments and step counter as well as the current parameters. Changing the next sample or dropout mask changes the gradient it will consume:

The live session owns model and optimizer state; the loop manages sample order, and external generators remain the caller's responsibility. A result is not a snapshot of every row in this table. The manual API exposes them when an experiment needs a custom loop or a checkpoint must record more than parameter tensors. Saving only weights is enough for inference, but it is not enough to resume Adam at the same update.

The buffers row is not hypothetical. Dropout consumes randomness and denotes a different function in training and evaluation mode (Srivastava et al., 2014), and batch normalization keeps running statistics that the forward pass updates without the optimizer touching them (Ioffe and Szegedy, 2015). Our state=33 above equals params=33 only because this model has neither.

Adam Moments And Step Counter🔗

Adam (Kingma and Ba, 2015) maintains an exponentially weighted mean m_t of the gradients and a mean v_t of their coordinatewise squares. Here t=1 denotes the first update: \theta_0 is the initial parameter pack, m_0=v_0=0, and g_t is the gradient evaluated at \theta_{t-1}. The decay factors \beta_1,\beta_2\in[0,1) determine how much of each previous estimate survives:

m_t=\beta_1m_{t-1}+(1-\beta_1)g_t,

v_t=\beta_2v_{t-1}+(1-\beta_2)g_t^2.

At the first update, the zero initial moments give m_1=(1-\beta_1)g_1 and v_1=(1-\beta_2)g_1^2. Dividing by 1-\beta_1 and 1-\beta_2 recovers g_1 and g_1^2. After t updates, the weights assigned to the observed gradients sum to (1-\beta_1)\sum_{k=0}^{t-1}\beta_1^k=1-\beta_1^t; the corresponding weights on their squares sum to 1-\beta_2^t. Bias correction normalizes each estimate by this total weight:

\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 = \theta_{t-1}-\eta \frac{\widehat m_t}{\sqrt{\widehat v_t}+\epsilon}.

This counter records completed updates, so the first iteration in a zero-based loop has Adam counter t=1. The square, square root, and division act coordinatewise. The first moment supplies a smoothed gradient, while the square root of the second moment rescales each coordinate; the positive \epsilon keeps a zero second-moment estimate from producing a zero denominator. Both moment packs have the same dependent tensor shapes as \theta. Restoring only parameters from a checkpoint loses these estimates and the step counter used in bias correction, so it does not resume the same Adam update.

PyTorch exposes the corresponding state for each parameter:

per-parameter state keys: ['exp_avg', 'exp_avg_sq', 'step']
step = 1.0 exp_avg shape [8, 2]

exp_avg is m_t, exp_avg_sq is v_t, and step is t. The shape of exp_avg matches the weight it belongs to, which is the untyped version of the statement we get from the type: a moment pack has the same shape layout as the parameter pack. Momentum (Polyak, 1964) is the same story with one buffer instead of two, and decoupled weight decay (Loshchilov and Hutter, 2019) changes the update without changing the state layout.

TorchLean also provides SGD, momentum SGD, AdamW, AdaGrad, RMSProp, Adadelta, and Muon-related module-level runtime configuration. Muon requires an explicit orthogonalization backend and is not one of the opaque optim.Optimizer constructors accepted by every high-level trainer path. Their state and laws remain optimizer-specific.

Update Counts And Gradient Accumulation🔗

For the unbatched model [2]\to[1], this configuration:

-- Four samples contribute to each of the 200 optimizer
-- updates.
steps := 200
samplesPerStep := 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.

Fifty updates with one sample each and fifty updates with four samples each are two different amounts of work, and they land in different places:

-- samplesPerStep = 1
dataset size = 25
mean_loss(before training) = 0.495227
mean_loss(after training) = 0.335392
-- samplesPerStep = 4
dataset size = 25
mean_loss(before training) = 0.495227
mean_loss(after training) = 0.008248

The second run saw four times as many samples, so this is not evidence that averaging gradients is better per unit of work. It is evidence of what the option means: steps counts updates, and the starting point of both runs is identical because samplesPerStep does not touch initialization.

Averaging four gradients at one parameter point is also different from taking four consecutive updates. In the consecutive case, the second gradient sees parameters changed by the first update; in the accumulated case, all four gradients see the original parameters. Adam then updates its moments once from their mean. The option therefore changes both how much data contributes to each update and how often the optimizer's memory advances, even before considering a vectorized implementation.

For a true vectorized minibatch, the batch axis goes into the model type, and nn.linear takes it as batchShape:

Training it is the same call, and the summary tells us that the batch axis cost us no parameters:

batched model:
Sequential: [5, 2] -> [5, 1], layers=3, params=33, state=33
  [0] Linear(2, 8): [5, 2] -> [5, 8] params=24, state=24 [[8, 2], [8]]
  [1] ReLU: [5, 8] -> [5, 8] params=0, state=0 []
  [2] Linear(8, 1): [5, 8] -> [5, 1] params=9, state=9 [[1, 8], [1]]
dataset size = 5
mean_loss(before training) = 0.495227
mean_loss(after training) = 0.013059
steps=40 arithmetic=native scalar=Float32 loss=0.495227 -> 0.013059

The parameter count remains 33 because all five samples share the layer weights. The dataset now contains five items, each holding five samples; no incomplete group remains in this example. The initial mean loss again prints as 0.495227: this MLP applies the same initial parameters to the same 25 samples, grouped differently. The output lets us check the grouping separately from the subsequent optimizer trajectory.

With the default TrainOptions.samplesPerStep := 1, an update consumes one item and therefore runs one vectorized forward and backward pass over five samples. Setting samplesPerStep 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.

The maintained CSV example is already batched this way:

# This CSV experiment uses tensor batches of five and its
# own regression data.
python3 NN/Examples/Data/generate_small_data.py
lake exe torchlean data_csv \
  --device cpu --batch 5 --steps 5 --seed 2026

== CSV loader training tutorial ==
model:
Sequential: [5, 2] -> [5, 1], layers=3, params=33, state=33
  [0] Linear(2, 8): [5, 2] -> [5, 8] params=24, state=24 [[8, 2], [8]]
  [1] ReLU: [5, 8] -> [5, 8] params=0, state=0 []
  [2] Linear(8, 1): [5, 8] -> [5, 1] params=9, state=9 [[1, 8], [1]]
data_dir = NN/Examples/Data
csv_path  = NN/Examples/Data/small_regression.csv
seed      = 2026
train     = Adam(lr=0.05), steps=5, batch_size=5, shuffle=true, drop_last=true
dataset size = 5
mean_loss(before training) = 0.210192
mean_loss(after training) = 0.055578
steps=5 arithmetic=native scalar=Float32 loss=0.210192 -> 0.055578
predict(batch=heldout) = [[0.303611], [0.303611], [0.303611], [0.303611], [0.303611]]

The five identical predictions are not a bug: the held-out batch repeats one row five times, so a batched model answers it five times.

The CSV loss cannot be read as the next stage of the earlier grid experiment. Its starting mean is 0.210192, not 0.495227, because this command uses the CSV example's data and configuration. Within that run, the before and after values measure the same dataset and are directly comparable. Across the two examples, the matching architecture and parameter count explain the shared API, while the different target values and learning rate explain why the logs describe separate experiments.

Eager And Typed Graph Execution🔗

Compare:

# Keep the run configuration fixed except for eager versus
# typed graph execution.
lake exe torchlean quickstart_mlp \
  --device cpu --execution eager --steps 20 --seed 2026

lake exe torchlean quickstart_mlp \
  --device cpu --execution typed-graph --steps 20 --seed 2026

Both print the same three numbers:

mean_loss(before training) = 0.495227
mean_loss(after training) = 0.401184
trained(heldout) = [0.365380]

The trainer method remains train; execution := .typedGraph changes how that method runs without changing the model's forward definition. The execution chapter develops the graph reuse and proof boundaries in detail.

Agreement on twenty updates of a 33-parameter model is evidence, not proof. It is worth having anyway, because the two execution modes take different code paths through the runtime, and a disagreement deserves investigation of operations, reductions, and numeric tolerances.

Typed graph trainer execution is currently CPU-only. A non-CPU typed graph request is rejected rather than silently falling back to a different semantics.

Device Selection And Backend Profiles🔗

Programmatically, Trainer.Config extends Trainer.RunConfig, so the device can be set in the same record as the objective and seed. RunConfig.withDevice checks the device against the table of maintained backend profiles and returns an Except:

cpu: accepted
cuda: accepted
metal: device `metal` has no maintained runtime profile; provide an explicit backend profile

The CPU-only guide build still accepts cuda at this stage. withDevice checks whether a maintained profile exists; opening a session later checks whether the binary can execute it. Metal, ROCm, TPU, and Trainium have configuration names but no maintained profile accepted here. CUDA execution additionally requires a binary linked with the native runtime:

# The CUDA request needs a native-runtime build; the report
# records selected operation
# providers.
lake -R -K cuda=true exe torchlean quickstart_mlp \
  --device cuda --steps 20 --seed 2026 --show-backend

The two successful lookup rows attach device profiles without allocating tensors or executing kernels. Opening a session validates execution and rejects CUDA when the native runtime is unavailable. --show-backend prints the selected implementation for each operation; the backend chapter explains the profile, provider, VJP, and evidence fields in that report.

Training Arithmetic🔗

The common executable selections are:

# These alternative flags choose the scalar arithmetic used
# by a training run.
--arithmetic native
--arithmetic ieee

Native arithmetic uses Lean's builtin Float32 operations; .ieee now uses FloatLib binary32, including its finite and exceptional cases. These recorded results predate that migration and retain their original labels. On that run the two agree to the last printed digit:

steps=20 arithmetic=native scalar=Float32 loss=0.495227 -> 0.401184
steps=20 arithmetic=ieee scalar=IEEE32Exec loss=0.495227 -> 0.401184

The two implementations agree at the printed precision on this workload. Comparing their bits would require more than these decimal loss values. 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 and no-overflow hypotheses when relating it to FloatLib binary32. The floating-point chapter is where those boundaries are drawn.

A loss curve without its arithmetic semantics is incomplete. The same architecture and seed may round differently in native binary32, the bit-level reference, a fused CUDA kernel, or an external provider.

Model And Optimizer Checkpoints🔗

Native Float32 modules use an exact binary32 checkpoint on CPU and CUDA. The binary64 Float path retains its exact-bit JSON format on CPU. Neither representation passes through decimal text. The expected state shapes come from the model, and every tensor must have the right shape and scalar count before the checkpoint is accepted:

-- Use this model definition as the expected tensor manifest
-- for the loaded state.
def loadForThisModel (path : System.FilePath) :=
  Checkpoint.State.load (nn.build 2026 model) path

The result is an IO action returning tensors whose dependent shape list is exactly nn.stateShapes (nn.build 2026 model). The runtime checkpoint loader turns such a checked pack into runtime state handles, while Checkpoint.load and Checkpoint.save 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 state is installed. This matters when two models share an initial state prefix: a checkpoint for the larger model cannot be accepted as a checkpoint for the smaller one merely because the first few shapes agree.

Here is the round trip, end to end, in a temporary directory that vanishes when the block finishes:

dataset size = 25
mean_loss(before training) = 0.495227
mean_loss(after training) = 0.125425
steps=0 arithmetic=native scalar=Float32 loss=0.125425 -> 0.125425
trained  = [0.189974]
restored = [0.189974]

The restoring trainer uses seed 999, so matching predictions show that loading replaced its initial state. steps=0 says the restored result performed no updates. The two losses in the restored report are both 0.125425, which equals the trained run's final loss, because Trainer.load evaluates the supplied dataset once rather than copying a number it was handed.

The restored model reproduces the displayed prediction. Resuming update forty-one of the old Adam process would also require its optimizer history, which this model checkpoint does not contain.

The supervised training paths wire the loadCheckpoint? and saveCheckpoint? fields of Trainer.TrainOptions into the checkpoint format selected by the runtime arithmetic. For example:

-- The per-call save path records model state after the
-- requested training updates.
def saveClassifier : Trainer.TrainOptions :=
  { steps := 200
    logEvery := 25
    saveCheckpoint? := some "artifacts/classifier.state.json" }

The same file can be written after the fact with trained.save and read back with Trainer.load. Loading through loadCheckpoint? continues training from the stored parameters; loading through Trainer.load produces a result whose steps is zero and whose two losses are both the mean loss of the restored parameters on the supplied data, exactly as printed above.

This is a model-state checkpoint, not a complete training snapshot. The eager CUDA runtime can save Adam or AdamW moments and step counters separately with Checkpoint.Optimizer.save, then restore them with Checkpoint.Optimizer.load. 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, arithmetic 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.

Custom Training Loops🔗

trainer.train owns the optimizer loop. For custom control at the same API level, open the trainer as a session and own the loop in the calling program:

step 0: loss = 0.250488
step 1: loss = 0.031549
step 2: loss = 0.001341
step 3: loss = 0.031191
step 4: loss = 0.021362
step 5: loss = 0.003199
mean loss after 6 updates = 0.465379

The six update lines report pre-update losses on neighboring grid points. The final eval instead measures all 25 samples at the parameters reached after those updates, returning 0.465379 against an initial mean of 0.495227. The low losses on several visited samples therefore coexist with a much larger dataset mean. A custom loop can measure both by choosing when to call eval.

Trainer.Session exposes:

-- State and measurements can be queried between updates
-- before finish freezes a result.
Session.step
Session.stepBatch
Session.steps
Session.predict
Session.loss
Session.eval
Session.state
Session.save
Session.load
Session.finish

Use it when the program needs a custom accumulation policy, multiple losses, generated batches, reinforcement-learning interaction, or detailed instrumentation. Generated batches simply become the samples passed to step; PINN collocation points and simulator batches fit this directly.

trainer.train itself opens a session, loops over step, and calls finish. Writing that loop gives us access to the live parameters between updates, with the same model and autograd semantics. finish packages them as the usual Trainer.Result.

Loading into an existing session has a different lifecycle from constructing a restored result. Session.load replaces the model parameters and buffers but keeps that session's optimizer history and completed-step counter. This is useful only when those retained values are intentional: old Adam moments will affect the next update of the loaded parameters. Open a fresh session when the desired experiment is a fresh optimizer run from saved weights. Saving and loading model state alone does not decide that policy for the caller.

[0.100000, 0.100000, 0.100000, 0.050000, 0.050000, 0.050000, 0.025000, 0.025000, 0.025000, 0.012500]

[0.1, 0.1, 0.1, 0.05, 0.05, 0.05, 0.025, 0.025, 0.025, 0.0125]

-- The same schedule can be attached to a managed run or a
-- caller-driven live session.
def scheduledOptions : Trainer.TrainOptions :=
  { steps := 10
    scheduler := some tfDecay }

def scheduledSession := trainer.open (scheduler := some tfDecay)

-- constant 0.03 from optim.adam
dataset size = 25
mean_loss(before training) = 0.495227
mean_loss(after training) = 0.004741
-- tfDecay, starting from 0.1
dataset size = 25
mean_loss(before training) = 0.495227
mean_loss(after training) = 0.213333

[0.000250, 0.000500, 0.000750, 0.001000, 0.001000, 0.000966, 0.000868, 0.000722, 0.000550, 0.000378, 0.000232, 0.000134,
 0.000100]

[0.000000, 0.000001, 0.000300, 0.000600, 0.000239, 0.000060]

Training Option Validation🔗

An update that consumes zero samples has no meaning, so TrainOptions.validate refuses it before any tape is built:

training: samplesPerStep must be positive

The quickstart command also requires a positive step count:

$ lake exe torchlean quickstart_mlp --device cpu --steps 0 --seed 2026
error: quickstart_mlp: --steps must be > 0

TorchLean.CLI.orThrow converts a parse failure to an IO exception, adding the command name only if the message does not already include it. To inspect initialization without updates, use a session evaluation as shown next.

Training Experiments🔗

untrained(heldout) = [-0.088261]
trained(heldout)   = [0.228325]

initial loss = 0.495227

-- seed 2026
dataset size = 25
mean_loss(before training) = 0.495227
mean_loss(after training) = 0.004741
-- seed 2027
dataset size = 25
mean_loss(before training) = 0.607745
mean_loss(after training) = 0.064443
-- seed 2028
dataset size = 25
mean_loss(before training) = 0.371631
mean_loss(after training) = 0.005225

-- adam, lr 0.03
dataset size = 25
mean_loss(before training) = 0.495227
mean_loss(after training) = 0.004741
-- sgd, lr 0.03
dataset size = 25
mean_loss(before training) = 0.495227
mean_loss(after training) = 0.037862

Training Evidence And Limits🔗

Each loss in the log comes from a particular traversal of this pipeline:

-- The log records this configured execution; each arrow
-- still has its own semantic contract.
data -> forward -> loss -> reverse pass -> optimizer -> parameters

The seed experiment shows sensitivity over a fixed number of updates; it leaves asymptotic convergence open. The six-step session measures visited-sample and dataset losses, with no unseen-data evaluation. To ask about convergence for all initializations or robustness throughout an input region, we need the statements and hypotheses developed in the optimization chapter and the certificate chapters.

The saved parameters and backend audit rows identify what ran. Establishing equality of eager, typed graph, CUDA, and LibTorch execution would require relating their operations; a successful training run alone cannot establish that relation or the correctness of every native instruction.

For a useful comparison, decide which quantity is being held fixed before interpreting a smaller loss. The accumulation example fixes update count but changes samples consumed; the optimizer example fixes samples and updates but changes the rule and its memory; the seed example changes the starting parameters. Their logs answer those particular comparisons. None supplies a timing measurement, and a displayed held-out prediction is one evaluation point rather than an estimate of performance over a held-out population.