Trainer¶

Trainer runs everything that learns. A script calls set_seed, train, then optionally train_out_proj, train_finetune and prune, and infer to predict (see Training and Heads and fine-tuning). The other methods are the steps train runs and the checkpoint helpers.

Trainer ¶

Trainer(config: Any, batch_callback: Callable[[Any], tuple[Tensor, Tensor]] | None = None, feature_names: list[str] | None = None, class_weights: Tensor | None = None)

Grows and trains a SONN, layer by layer.

train runs the structural search: it creates a layer of candidate neurons, calibrates them on the layer's inputs (fit_layer_inputs), fits all candidates in parallel and keeps the best ones (train_layer), and adds layers until the growth rule (GrowthCriterion) stops it. The layers after the best one are dropped and the model is saved. After that, train_out_proj fits the optional output head, train_finetune trains every parameter end to end, prune removes the neurons that do not reach the output and infer predicts over a loader.

Every train call writes to a run folder of its own, run_dir: <train.checkpoint_dir>/<YYYY-MM-DD-HH-MM-SS>/. It holds the step checkpoints written as training goes, model_last.ckpt and the run's train.log. load_model_checkpoint restores a saved model.

Parameters:

Name	Type	Description	Default
`config`	`SONNConfig - like`	The full configuration (`model` and `train` sections).	required
`batch_callback`	`callable`	Maps a loader batch to an `(x, y)` pair. Default: batches already are `(x, y)`.	`None`
`feature_names`	`list of str`	Feature names, used in logs and model descriptions.	`None`
`class_weights`	`tensor`	Per-class loss weights, passed to the model at the start of `train`.	`None`

run_dir `property` ¶

run_dir: Path | None

The current run folder, <checkpoint_root>/<YYYY-MM-DD-HH-MM-SS>.

train creates a new one at the start of every run, or reopens the run it resumes, and load_model_checkpoint sets it to the run it loads from. The run's step checkpoints, model_last.ckpt and train.log live there. None before the first run or load.

checkpoint_root `property` ¶

checkpoint_root: Path

The folder that holds the run folders: train.checkpoint_dir of the trainer's configuration, checkpoints in the working directory by default.

set_seed `classmethod` ¶

set_seed(seed: int = 42) -> None

Seed random, NumPy and PyTorch (CPU and every GPU), and make cuDNN deterministic.

run_folders `classmethod` ¶

run_folders(root: str | Path) -> list[Path]

List the run folders in root, oldest first.

A run folder's name is its start time, YYYY-MM-DD-HH-MM-SS, with a -<n> suffix for runs started in the same second. They sort by the time, then by the suffix as a number. Other entries are left out.

Parameters:

Name	Type	Description	Default
`root`	`str or Path`	The folder to look in, usually `checkpoint_root`.	required

Returns:

Type	Description
`list of Path`	The run folders; empty when `root` does not exist.

parse_checkpoint_step `classmethod` ¶

parse_checkpoint_step(fname: str) -> tuple[int, int, int]

Parse a step checkpoint file name.

Parameters:

Name	Type	Description	Default
`fname`	`str`	File name of the form `model_layer_<L>_neuron_<N>_step_<S>[_last].ckpt`.	required

Returns:

Name	Type	Description
`key`	`tuple of int`	`(layer, neuron_model, step)`, or `(-1, -1, -1)` when the name does not match, so that such files sort below every real checkpoint.

from_checkpoint ¶

from_checkpoint(run_dir: str | Path | None = None) -> dict[str, Any] | None

Load the latest step checkpoint of a run, for resuming train.

The latest one is the highest (layer, neuron_model, step) among the run folder's step checkpoints, _last copies included. Other files in the folder, such as train.log and model_last.ckpt, are ignored.

Parameters:

Name	Type	Description	Default
`run_dir`	`str or Path`	The run folder. Default: `run_dir`.	`None`

Returns:

Name	Type	Description
`checkpoint`	`dict or None`	The checkpoint contents, or None if there is no run folder or it holds no step checkpoint.

save_checkpoint ¶

save_checkpoint(ckpt: StepCheckpoint, suffix: str = '') -> None

Save a step checkpoint in the run folder and delete the older ones.

Writes model_layer_<L>_neuron_<N>_step_<S>[_<suffix>].ckpt with the snapshot and the model config, then keeps the train.keep_last_n most recent step checkpoints of the run folder (see cleanup_checkpoints).

Parameters:

Name	Type	Description	Default
`ckpt`	`StepCheckpoint`	The training state to save.	required
`suffix`	`str`	Optional suffix for the file name, such as "last". Default "".	`''`

load_model_checkpoint ¶

load_model_checkpoint(model: SONN, device: device | str = 'cpu', run: str | Path | None = None) -> None

Restore a model saved as model_last.ckpt.

Reads it from a run folder, rebuilds the layer stack, loads the weights and moves the model to device. The run folder becomes run_dir, so a later save_model_checkpoint writes to the same run.

Parameters:

Name	Type	Description	Default
`model`	`SONN`	The model to restore into, built with the same config.	required
`device`	`device or str`	Device to move the restored model to. Default "cpu".	`'cpu'`
`run`	`str or Path`	The run folder, by its name in `checkpoint_root` or by its path. Default: the trainer's current run, or, for a trainer that has none, the newest run folder in `checkpoint_root` that holds a `model_last.ckpt`.	`None`

Raises:

Type	Description
`FileNotFoundError`	If there is no such run folder, or no run folder holds a `model_last.ckpt`.

save_model_checkpoint ¶

save_model_checkpoint(model: SONN) -> None

Save the model's state dict as model_last.ckpt in the run folder (a new one if the trainer has none yet).

cleanup_checkpoints ¶

cleanup_checkpoints(checkpoint_dir: str | Path, keep_last_n: int = 10) -> None

Delete a run folder's oldest step checkpoints, keeping the keep_last_n most recent.

Only the files named model_layer_<L>_neuron_<N>_step_<S>.ckpt count, ordered by parse_checkpoint_step. Everything else in the folder (model_last.ckpt, the _last copies, train.log, other files and subfolders) is left alone.

Parameters:

Name	Type	Description	Default
`checkpoint_dir`	`str or Path`	The run folder.	required
`keep_last_n`	`int`	Number of most recent step checkpoints to keep. Default 10.	`10`

cleanup_layer_checkpoints ¶

cleanup_layer_checkpoints(layer_idx: int) -> None

Delete every step checkpoint of one layer from the run folder, _last copies included.

Parameters:

Name	Type	Description	Default
`layer_idx`	`int`	The layer whose checkpoints are deleted.	required

init_distributed `classmethod` ¶

init_distributed(backend: str = 'nccl') -> None

Initialize torch.distributed for multi-GPU / multi-node training.

Call once before Trainer.train(), then launch with torchrun: torchrun --nproc_per_node= --nnodes= script.py

Sets LOCAL_RANK as the active CUDA device so each rank uses a distinct GPU within the node.

train ¶

train(model: SONN, train_dl: DataLoader, dev_dl: DataLoader, test_dl: DataLoader, verbose: bool = True, resume: bool | str = False, val_dl: DataLoader | None = None) -> SONN

Grow and train the network, one layer at a time.

Each call is a run with a folder of its own, run_dir: a new <train.checkpoint_dir>/<YYYY-MM-DD-HH-MM-SS>/, or the folder of the run it resumes. The run's step checkpoints, model_last.ckpt and train.log (every record of the root logger from the start of the run until the next run starts) are written there.

Each iteration creates a layer of candidate neurons, calibrates them on the layer's inputs, trains them and keeps the best ones (train_layer). Growth stops when GrowthCriterion sees no improvement for train.criterion_minimum_width layers, or at train.max_layer_count. The layers after the one with the smallest error are then dropped, and the model is saved as model_last.ckpt in the run folder. The output head and the end-to-end fine-tune are separate steps (train_out_proj, train_finetune).

Parameters:

Name	Type	Description	Default
`model`	`SONN`	The model to train. Its existing layers are discarded.	required
`train_dl`	`DataLoader`	Training split: fits the neuron coefficients.	required
`dev_dl`	`DataLoader`	Dev split: scores and selects the neurons, and drives the growth rule unless `train.stop_source` is 'val'.	required
`test_dl`	`DataLoader`	Not used during training.	required
`verbose`	`bool`	Log the time each layer takes. Default True.	`True`
`resume`	`bool or str`	False starts a new run. True continues the newest run folder in `checkpoint_root` that holds a step checkpoint, from its latest one, and starts a new run when there is none. A run folder's name continues that run. Default False.	`False`
`val_dl`	`DataLoader`	Validation split that selects nothing. Its error is logged per layer, and it drives the growth rule when `train.stop_source` is 'val'.	`None`

Returns:

Name	Type	Description
`model`	`SONN`	The trained model (the same object).

train_layer ¶

train_layer(model: SONN, layer: SONNLayer, train_dl: DataLoader, dev_dl: DataLoader, checkpoint_data: dict[str, Any] | None) -> None

Calculate neuron weights

neuron_selection ¶

neuron_selection(model: SONN, layer: SONNLayer, accumulator: LayerAccumulator, trained_params_per_model: dict[int, dict], device: device | str) -> tuple[torch.Tensor, torch.Tensor]

Pick the nbest survivors across the layer's neuron-model ensemble and prune the rest in place.

Concatenates per-neuron-model errors and module indices from the accumulator, drops NaNs, takes top-k (lowest error) up to model.nbest_neurons, re-sorts by module index, prunes each neuron_model down to its surviving rows, and restores model.shared_proj from the best neuron's trained params (shared-proj path only).

Mutates layer.neuron_models, layer.d_model, and optionally model.shared_proj weight/bias in place. Returns the (err_values, module_idxs) tuple the caller attaches to layer.err_values / layer.module_idxs.

train_model_ensemble ¶

train_model_ensemble(model: SONN, neuron_model: BasePolynomNeuron, neuron_model_idx: int, layer_idx: int, train_dl: DataLoader, dev_dl: DataLoader, checkpoint_data: dict[str, Any] | None, accumulator: LayerAccumulator, loss_fn_vmapped: Callable, eval_loss_fn_vmapped: Callable, params_batch: dict[str, Tensor], buffers_batch: dict[str, Tensor], shared_param_names: list[str], features_precomputed: bool = False) -> tuple[LayerAccumulator, dict[str, torch.Tensor], StepCheckpoint | None]

Train one neuron module's candidate ensemble in parallel.

Runs the configured batched optimizer over the stacked parameters of all candidates, with per-candidate early stopping and learning-rate drops (on the dev loss, or on the training loss when train.early_stop_source is 'train'). Step checkpoints are saved every train.save_interval steps and at the end, and the trained parameters are written back to the module. Under distributed training each rank trains its own slice of the ensemble.

Parameters:

Name	Type	Description	Default
`model`	`SONN`	The model being grown.	required
`neuron_model`	`BasePolynomNeuron`	The neuron module whose candidates are trained.	required
`neuron_model_idx`	`int`	Index of that module in the layer.	required
`layer_idx`	`int`	Index of the layer.	required
`train_dl`	`DataLoader`	Training and dev splits.	required
`dev_dl`	`DataLoader`	Training and dev splits.	required
`checkpoint_data`	`dict or None`	Step checkpoint to resume from, or None to start fresh.	required
`accumulator`	`LayerAccumulator`	The layer's running state, stored in the checkpoints.	required
`loss_fn_vmapped`	`callable`	Gradient and evaluation functions from `create_loss_functions`.	required
`eval_loss_fn_vmapped`	`callable`	Gradient and evaluation functions from `create_loss_functions`.	required
`params_batch`	`dict of str to tensor`	Stacked candidate parameters and buffers.	required
`buffers_batch`	`dict of str to tensor`	Stacked candidate parameters and buffers.	required
`shared_param_names`	`list of str`	Parameters shared by all candidates (the shared projection).	required
`features_precomputed`	`bool`	The loaders already yield the layer's input features, so the model forward is skipped. Default False.	`False`

Returns:

Name	Type	Description
`accumulator`	`LayerAccumulator`	The accumulator passed in.
`params_batch`	`dict of str to tensor`	The trained candidate parameters.
`last_checkpoint`	`StepCheckpoint`	The snapshot of the finished ensemble, as saved.

write_back_params `classmethod` ¶

write_back_params(model: Module, new_params: dict[str, Tensor], new_buffers: dict[str, Tensor] | None = None) -> None

Copy updated batched parameters & buffers into the model.

Parameters:

Name	Type	Description	Default
`model`	`Module`	The module to update in place.	required
`new_params`	`dict of str -> torch.Tensor`	Updated batched parameters.	required
`new_buffers`	`dict of str -> torch.Tensor`	Updated buffers.	`None`

create_loss_functions ¶

create_loss_functions(model: SONN, module: BasePolynomNeuron) -> tuple[Callable, Callable, Callable, dict[str, torch.Tensor], dict[str, torch.Tensor], list[str]]

Build the vmapped loss functions for one neuron module's candidates.

Parameters:

Name	Type	Description	Default
`model`	`SONN`	The model being grown.	required
`module`	`BasePolynomNeuron`	The neuron module whose candidates are trained.	required

Returns:

Name	Type	Description
`loss_fn_vmapped`	`callable`	Per-candidate gradient of the training loss, including the `train.ridge_alpha` penalty.
`eval_loss_fn_vmapped`	`callable`	Per-candidate loss without the penalty.
`pred_fn_vmapped`	`callable`	Per-candidate predictions.
`params_batch`	`dict of str to tensor`	The module's parameters, plus the shared projection for multi-class models that use one.
`buffers_batch`	`dict of str to tensor`	The module's buffers and `src_idxs`.
`shared_param_names`	`list of str`	Names of the parameters shared by all candidates (not vmapped).

fit_layer_inputs ¶

fit_layer_inputs(model: SONN, layer: SONNLayer, train_dl: DataLoader) -> None

Per-layer input pass: calibrate every neuron model on the layer's actual inputs before the layer trains.

The inputs have to be the features that reach this layer: the (preprocessed) raw inputs for layer 0, the frozen prefix's outputs — plus the shortcut-concatenated originals — deeper down. So this runs once per layer, after create_layer and before any of the layer's neurons train, and measures model(x, skip_last_layer=True) exactly as train_layer will feed it. Three things come out of it, each only when some neuron model in the layer asks (see the hooks on BasePolynomNeuron):

mean / std per input feature (needs_input_stats): the sigma squash of the orthogonal families. Streamed in float64.
a seeded reservoir sample of up to train.input_sample_rows rows (needs_input_sample), the whole split when it is smaller, taken in the same pass. Handed to fit_input_sample(sample, stream).
a streaming pass (needs_input_stream, and only when the mode resolves to 'stream': train.rbf_kmeans_mode 'stream', or 'auto' with more rows than the cap): train.rbf_kmeans_passes further passes over the split feeding stream_input_batch(x_batch), then finish_input_stream() once.

Costs one forward pass over the training split per layer, plus one per streaming pass. Skipped outright when no neuron in the layer asks for anything, so configs that use none of it pay nothing.

ds_loss ¶

ds_loss(model: SONN, eval_loss_fn_vmapped: Callable, params_batch: dict[str, Tensor], buffers_batch: dict[str, Tensor], dl: DataLoader, device: device | str, skip_model_fwd: bool = False) -> torch.Tensor

Mean evaluation loss of every candidate over a loader.

Parameters:

Name	Type	Description	Default
`model`	`SONN`	The model being grown.	required
`eval_loss_fn_vmapped`	`callable`	Per-candidate loss, from `create_loss_functions`.	required
`params_batch`	`dict of str to tensor`	Stacked candidate parameters and buffers.	required
`buffers_batch`	`dict of str to tensor`	Stacked candidate parameters and buffers.	required
`dl`	`DataLoader`	The split to evaluate on.	required
`device`	`device or str`	Device the batches are moved to.	required
`skip_model_fwd`	`bool`	The loader already yields the layer's input features, so the model forward is skipped. Default False.	`False`

Returns:

Name	Type	Description
`loss`	`(num_candidates,) tensor`	Per-candidate loss, averaged over the batches.

regularity_err ¶

regularity_err(model: SONN, pred_fn_vmapped: Callable, params_batch: dict[str, Tensor], buffers_batch: dict[str, Tensor], dl: DataLoader, device: device | str, skip_model_fwd: bool = False) -> torch.Tensor

Per-candidate regularity criterion, computed streaming per batch.

Regression / binary: normalized MSE (matches regularity_error). Multi-class: normalized cross-entropy (matches regularity_error_ce).

The sums are accumulated batch by batch instead of concatenating the predictions into one (ensemble, N, K) tensor first: for Otto-sized ensembles (~17k candidates after shortcut) that tensor would take over 6 GB of GPU memory, twice that inside log_softmax. GPU memory stays at one batch's predictions.

bias_err ¶

bias_err(model: SONN, pred_fn_vmapped_a: Callable, pred_fn_vmapped_b: Callable, params_batch_a: dict[str, Tensor], params_batch_b: dict[str, Tensor], buffers_batch_a: dict[str, Tensor], buffers_batch_b: dict[str, Tensor], dl_a: DataLoader, dl_b: DataLoader, device: device | str, bias_method: str = 'l2', skip_model_fwd: bool = False) -> torch.Tensor

Bias error between two neurons fitted on disjoint splits.

Regression / binary: bias_error — normalized squared disagreement on scalars. Multi-class, bias_method="l2": bias_error_l2 on logits (ensemble, N, K). Multi-class, bias_method="js": bias_error_js on logits (ensemble, N, K). pred_fn_vmapped returns (ensemble, N) for regression and (ensemble, N, K) for multi-class, so no separate raw-pred function is needed. Both splits are concatenated before calling the criterion function.

train_out_proj ¶

train_out_proj(model: SONN, train_dl: DataLoader, dev_dl: DataLoader) -> None

Train model.out_proj with all other parameters frozen.

Runs a standard PyTorch training loop — no vmapped ensemble, no custom optimizer — because out_proj is a single shared nn.Linear whose gradient comes from the full batch in one shot.

Stops at cfg.max_steps, or after early_stop_patience evaluations in a row without an improvement of the dev loss. ReduceLROnPlateau lowers the learning rate when the dev loss stalls, down to lr_min. With optimizer: lbfgs the head is fitted full-batch instead, without the plateau schedule.

The head ends on its last step, or with cfg.keep_best_weights on the weights of its lowest dev loss (a last step between two evaluations is evaluated too), copied back in place from where cfg.best_weights_copy keeps them. The model is then saved to the run folder's model_last.ckpt.

train_finetune ¶

train_finetune(model: SONN, train_dl: DataLoader, dev_dl: DataLoader, cfg: Any = None, val_dl: DataLoader | None = None) -> None

End-to-end fine-tuning of all SONN parameters.

After SONN's structural search is done every layer is a normal differentiable module — gradients flow through BasePolynomNeuron.weight, shared_proj / soft_binner / neuron_proj, and out_proj all at once.

The step is launch-bound on these models (hundreds of tiny kernels), so on a CUDA device with adam / adamw it is captured once as a CUDA graph over a static buffer of the loader's batch shape and replayed per step, with the next batch copied into the buffer; see _FinetuneStep and OutProjTrainConfig.cuda_graph / data_on_device. Evaluation, the plateau schedule and the early stop run outside the graph every eval_interval steps.

cfg is an OutProjTrainConfig-shaped block of LR / scheduler / early-stop hyperparameters; it defaults to model.param.train.finetune_train, which exists so an end-to-end pass can use a different (usually much smaller) learning rate than the out_proj / layer_finetune passes that share out_proj_train.

The model ends on its last step, or with cfg.keep_best_weights on the parameters of its lowest loss on the split the early stop reads (a last step between two evaluations is evaluated too), copied back in place, so a captured graph stays valid, from where cfg.best_weights_copy keeps them. The model is then saved to the run folder's model_last.ckpt.

infer ¶

infer(model: SONN, test_dl: DataLoader, verbose: bool = True, use_compile: bool = False) -> tuple[torch.Tensor, torch.Tensor]

Predict over test_dl and return (predictions, targets), both in loader order.

The predictions are one value per row for regression, a logit per row for binary models and a row of log-probabilities per sample for multi-class models (see SONN.infer). Each batch must be an (x, y) pair: batch_callback is not applied here.

use_compile=True wraps model.infer with torch.compile before the loop, fusing the sequential per-layer kernel launches into a single optimized graph. First batch is slower (compilation), subsequent batches are faster — worthwhile when test_dl has many batches or is called repeatedly.

prune ¶

prune(model: SONN) -> None

Strip the network down to the neurons that actually reach the output.

The last layer is cut to the columns the model reads, and every earlier layer is then cut to whatever the layers reading it still reference — its successor, and with shortcut.prev_layers later layers too. A layer nobody reads any more is deleted.

How many columns the model reads is SONN._readout_width. Without a head, SONN.infer returns the best-error column and a single neuron suffices, except for a multi-class model with per-neuron projections (use_neuron_proj): its prediction sums every last-layer neuron's projected output, so the whole last layer stays. With out_proj it reads in_features columns, chosen by _best_neuron_columns — so collapsing the layer to one neuron would leave the head consuming zero-padding on every other input and silently destroy the model rather than raising. Keep exactly the columns the head consumes instead.

The head needs no reindexing: the retained neurons keep their err_values, so re-running _best_neuron_columns on the pruned layer re-derives the same neurons in the same ascending-error order, which is the order the head's inputs were fitted in. Every kept layer's d_model is set to its new width.

Raises:

Type	Description
`ValueError`	If the model was built with `model.use_layer_norm`. Its LayerNorms normalize each row over every column of a layer's input, so removing columns would change the predictions; the model is left as it is.

GrowthCriterion `dataclass` ¶

GrowthCriterion(width: int, epsilon: float, min_delta: float = 0.0, min_error: float = sys.float_info.max, ref_error: float = sys.float_info.max, best_index: int = 0, last_improved_index: int = 0, last_gain: float = 0.0, last_margin: float = 0.0, last_since: int = 0)

The layer-growth stop rule of Trainer.train, kept out of the loop so it can be tested on a scripted sequence of layer errors.

A layer improves when it lowers the error of the last accepted improvement (ref_error) by at least max(min_delta, epsilon * best), so several small steps down add up instead of each failing the margin on its own; the search stops after width consecutive layers without an improvement. best_index (the layer with the lowest error, margin or not) is what the trainer keeps; the window is counted from last_improved_index. Layer 0 always improves.

update ¶

update(layer_index: int, err: float) -> bool

Record a layer's error; returns True when the search should stop.

describe ¶

describe(layer_index: int, err: float) -> str

Return a one-line log message on the last update: error, best layer, gain and margin, layers without improvement.

setup_logger ¶

setup_logger(log_path: str | Path | None = None) -> logging.Logger

Configure the root logger for a training script.

Replaces every handler on the root logger (closing them, a run's train.log handler included) with a tqdm-aware console handler at INFO level, plus a file handler when log_path is given. Each training run also writes its own train.log in its run folder (see Trainer.run_dir), so a file here is needed only for lines outside the runs.

Parameters:

Name	Type	Description	Default
`log_path`	`(str, Path or None)`	A log file to write as well, overwritten on each call. Default None: the console only.	`None`

Returns:

Type	Description
`Logger`	The configured root logger.

Trainer¶

Trainer ¶

run_dir property ¶

checkpoint_root property ¶

set_seed classmethod ¶

run_folders classmethod ¶

parse_checkpoint_step classmethod ¶

from_checkpoint ¶

save_checkpoint ¶

load_model_checkpoint ¶

save_model_checkpoint ¶

cleanup_checkpoints ¶

cleanup_layer_checkpoints ¶

init_distributed classmethod ¶

train ¶

train_layer ¶

neuron_selection ¶

train_model_ensemble ¶

write_back_params classmethod ¶

create_loss_functions ¶

fit_layer_inputs ¶

ds_loss ¶

regularity_err ¶

bias_err ¶

train_out_proj ¶

train_finetune ¶

infer ¶

prune ¶

GrowthCriterion dataclass ¶

update ¶

describe ¶

setup_logger ¶

run_dir `property` ¶

checkpoint_root `property` ¶

set_seed `classmethod` ¶

run_folders `classmethod` ¶

parse_checkpoint_step `classmethod` ¶

init_distributed `classmethod` ¶

write_back_params `classmethod` ¶

GrowthCriterion `dataclass` ¶