SONN

SONN is the model: a PyTorch module that starts with no layers and that Trainer.train grows. Training shows the calls in order, and Inspecting a model how to read a trained one.

SONN

SONN(config: Any, d_model: int, feature_names: list[str] | ndarray | None = None, preprocessing: Module | None = None, class_weights: Tensor | None = None)

Bases: SONNModule

Self-organizing deep learning neural network (GMDH) as a PyTorch module.

The model starts with no layers; Trainer.train grows them, and infer predicts with the trained network.

Parameters:

Name	Type	Description	Default
config	DictConfig or mapping	The configuration. It is merged into default_config(), so it may hold only the keys that differ from the defaults.	required
d_model	int	Number of input features.	required
feature_names	list of str or ndarray	Feature names, used in logs, by get_selected_features and in the network diagrams.	None
preprocessing	Module	A module applied to every input batch before the first layer.	None
class_weights	tensor	Per-class loss weights for a multi-class model, or the weight of the positive class for a binary one. Trainer sets them when it is given class weights.	None

device property

device: device

Device of the first neuron module, or the configured train.device before any layer exists.

need_bias_err property

need_bias_err: bool

True when the selection criterion uses the bias error ('bias', 'validate_bias').

need_regularity_err property

need_regularity_err: bool

True when the selection criterion uses the regularity (validation) error ('validate', 'validate_bias').

concat_inputs property

concat_inputs: bool

True when a layer's input is more than the previous layer's output.

False only for shortcut: false (no raw features, no older layers), the configuration in which the last layer's LayerNorm also applies to its own output at inference — see forward.

set_class_weights

set_class_weights(class_weights: Tensor) -> None

Replace the loss function's class weights without rebuilding the model.

Safe to call after model.to(device) — the new loss module is moved to the model's current device automatically.

state_dict

state_dict(*args: Any, **kwargs: Any) -> dict[str, Any]

Record the layer class names, then return the module state dict.

restore_from_checkpoint_metadata

restore_from_checkpoint_metadata(checkpoint_data: dict[str, Any]) -> None

Rebuild the layer stack from checkpoint metadata.

Replaces self.layers with new layers holding neuron modules of the saved classes and shapes. Only the architecture is restored: the weights are loaded afterwards with load_state_dict (see Trainer.load_model_checkpoint).

Parameters:

Name	Type	Description	Default
checkpoint_data	dict	The model's state dict as saved in a checkpoint, including the params_metadata entries of the model, its layers and their neuron modules.	required

default_config classmethod

default_config() -> Any

Return the configuration schema with every default filled in.

The result is OmegaConf.structured(SONNConfig): a typed DictConfig with the sections model and train and a few top-level flags, a fresh copy on every call. Merge changes into it:

config = OmegaConf.merge(SONN.default_config(), {"model": {"nbest_neurons": 8}})

A key the schema does not have, or a value of the wrong type, is rejected when it is merged. SONNConfig and its sections document every key.

Returns:

Type	Description
DictConfig	The default configuration.

get_selected_features_indices

get_selected_features_indices() -> list[int]

Return features that was selected as useful for neuron during fit

get_unselected_features_indices

get_unselected_features_indices() -> list[int]

Return features that was not selected as useful for neuron during fit

get_unselected_features

get_unselected_features() -> str

Return names of features that was not selected as useful for neuron during fit

get_selected_features

get_selected_features() -> str

Return names of features that was selected as useful for neuron during fit

plot_layer_error

plot_layer_error() -> None

Plot layer error on validation set vs layer index

matplotlib is imported here rather than at module scope so it can stay an optional dependency — this is the only place in the package that touches it, and a base install must still be able to import SONN.

compute_loss staticmethod

compute_loss(neuron_model: Module, loss_fn: Module | None, model_type: str, soft_binner: Module | None, ridge_alpha: float, params: dict[str, Tensor], buffers: dict[str, Tensor], x: Tensor, y: Tensor) -> torch.Tensor

Loss of one candidate neuron, written for torch.func transforms.

Runs neuron_model with the given parameters and buffers through functional_call, so the trainer can vmap it over a candidate ensemble and take grad with respect to params. For multi-class models the scalar output is turned into logits by the soft binner, the shared projection or the neuron's own projection, whichever the parameters provide.

Parameters:

Name	Type	Description	Default
neuron_model	Module	The neuron module whose forward is called.	required
loss_fn	Module or None	Per-sample loss. None returns the raw predictions (logits for multi-class) instead of a loss.	required
model_type	str	model.type; "multi-class" selects the logits path.	required
soft_binner	Module or None	Maps the scalar output to class logits, if the model uses one.	required
ridge_alpha	float	L2 penalty on params["weight"], added to the loss when > 0. Projection weights are not penalized.	required
params	dict of str to tensor	Parameters of one candidate (one slice of the ensemble under vmap).	required
buffers	dict of str to tensor	Buffers of the neuron module.	required
x	(B, D) tensor	Neuron input.	required
y	(B,) tensor	Targets (class indices for multi-class).	required

Returns:

Name	Type	Description
out	tensor	Mean loss plus the ridge penalty as a scalar, or the predictions when loss_fn is None.

get_error

get_error(criterion_type: CriterionType, regularity_err: Tensor | None, bias_err: Tensor | None) -> torch.Tensor

Compute error of the neuron according to specified criterion

new_layer_sources

new_layer_sources(position: int) -> tuple[list[int], bool]

(input_layers, input_raw) the config gives a layer created at position.

layer_sources

layer_sources(position: int) -> tuple[list[int], bool]

(input_layers, input_raw) of the layer at position, as stored on it.

input_blocks

input_blocks(position: int) -> list[tuple[Optional[int], int]]

(source, width) per block of the input of the layer at position, in concatenation order. source is a layer position, or None for the raw model inputs.

locate

locate(position: int, flat_idx: int) -> tuple[Optional[int], int]

Map an index into the input of the layer at position to (source, index inside that source's block); source None = raw inputs.

next_input_width

next_input_width() -> int

Input width of a layer created next, at position len(self.layers).

forward

forward(x: Tensor, skip_last_layer: bool = False) -> torch.Tensor

Run the input through the layer stack.

Applies the preprocessing module, if any, then each layer on the input its layer_sources describe (outputs of earlier layers and, under shortcut.raw_features, the raw features). Layer outputs are clamped to model.output_clamp_value. No output head is applied: infer does that.

Parameters:

Name	Type	Description	Default
x	(B, d_model) tensor	Input features.	required
skip_last_layer	bool	Stop at the input of the last layer and return it. The trainer uses this while the last layer is being fitted. Default False.	False

Returns:

Name	Type	Description
out	tensor	The last layer's output, shape (B, number of neurons in it), or its input when skip_last_layer is True. With no layers, the preprocessed input.

get_best_neuron_model

get_best_neuron_model(layer: SONNLayer) -> tuple[torch.Tensor, torch.Tensor]

Locate the neuron with the smallest error in a layer.

Parameters:

Name	Type	Description	Default
layer	SONNLayer	A trained layer, with err_values and module_idxs set.	required

Returns:

Name	Type	Description
module_idx	tensor	Index of the neuron module that holds the best neuron.
neuron_idx	tensor	Index of the best neuron inside that module.

create_layer

create_layer(layer_index: int) -> SONNLayer

Generate new layer with all possible neurons

infer

infer(x: Tensor) -> torch.Tensor

Predict with the trained model.

Moves x to the model's device, runs forward and reads the prediction off the last layer: through the output head (out_proj) if the model has one, otherwise from the best-error neuron (regression and binary), the per-neuron projections or the shared projection / soft binner (multi-class).

Parameters:

Name	Type	Description	Default
x	(B, d_model) tensor	Input features.	required

Returns:

Name	Type	Description
pred	tensor	(B,) predictions for regression and binary models (logits for binary), or (B, C) log-probabilities for multi-class models.

SONNLayer

SONNLayer(d_model: int, nbest_neurons: int, layer_index: int, use_layer_norm: bool = False)

Bases: SONNModule

Layer class

setup_layer_norm

setup_layer_norm(dim: int) -> None

Build the per-layer LayerNorm with the supplied dim as its normalized_shape. Idempotent and a no-op when the flag is off.

Called from Trainer.neuron_selection after pruning (passing the width of the next layer's input, SONN.next_input_width) and from from_checkpoint_metadata when restoring a model.

fit_input_stats

fit_input_stats(mean: Tensor, std: Tensor) -> None

Hand this layer's input statistics to every neuron model.

mean / std are length-num_feat vectors over the features that feed this layer, measured on the training set. Called by Trainer.fit_layer_inputs after the layer is created and before it trains; a no-op for neuron families that calibrate nothing.

state_dict

state_dict(*args: Any, **kwargs: Any) -> dict[str, Any]

Record the neuron class names, then return the module state dict.

from_checkpoint_metadata classmethod

from_checkpoint_metadata(metadata: dict[str, Any]) -> SONNLayer

Rebuild an empty layer from its checkpoint metadata.

Restores the constructor arguments, the LayerNorm at its saved width and the input layout (input_layers, input_raw). The neuron modules are not created here; the model restores them separately.

Parameters:

Name	Type	Description	Default
metadata	dict	The layer's entry in the checkpoint metadata.	required

Returns:

Name	Type	Description
layer	SONNLayer	The restored layer, without neuron modules.

to

to(*args: Any, **kwargs: Any) -> SONNLayer

Move the layer, including its plain tensor attributes, like nn.Module.to.

forward

forward(x: Tensor) -> torch.Tensor

Run every neuron module on the layer input.

Parameters:

Name	Type	Description	Default
x	(B, D_in) tensor	The layer input.	required

Returns:

Name	Type	Description
out	(B, len(self)) tensor	The outputs of all neuron modules, concatenated along dim 1 in module order.

get_parent_neron_module

get_parent_neron_module(idx: int) -> tuple[int, BasePolynomNeuron]

Find the neuron module that holds output column idx.

Parameters:

Name	Type	Description	Default
idx	int	Absolute neuron (output column) index in the layer.	required

Returns:

Name	Type	Description
idx	int	The same absolute index.
module	BasePolynomNeuron	The neuron module that produces that column.

Raises:

Type	Description
ValueError	If idx is not below the layer's neuron count.

set_neuron_module

set_neuron_module() -> None

Build the column map of the layer.

Sets neuron_idxs, an (N, 2) tensor whose row i holds the module index and the neuron index inside that module for output column i, and sets d_model to the number of columns N.

LayerCreationError

LayerCreationError(message: str, layer_index: int)

Bases: Exception

Raised when layer creation fails (e.g. no reference functions configured).