Neurons¶
Each entry of model.ref_functions builds one of these classes in every
layer; a module holds all of a family's candidates as rows of its
tensors. See Neuron families for the
formulas and Neuron family options
for the options.
BasePolynomNeuron(num_feat: int, num_src_feat: int, activation: ActivationLike, layer_index: int, start_index: int, dim: int = 2, max_neuron_models: int | None = None, init_method: str = 'xavier')
Bases: SONNModule, ABC
Base class of every neuron family: an ensemble of candidate neurons of one type, each reading its own tuple of layer inputs.
One module holds all num_neurons candidates of a family in a layer.
src_idxs (num_neurons, dim) lists the input columns each candidate
reads, and weight (num_neurons, num_w) holds their coefficients, the
constant term included. forward gathers each candidate's inputs,
expands them into its design row with get_args and returns the
weighted sum through activation.
Subclasses set num_w and implement get_args, get_name and
get_short_name. Every subclass is registered by class name, so
from_checkpoint_metadata can rebuild it from a checkpoint.
property
¶
Number of candidate neurons in this module (rows of weight).
property
¶
Number of candidate neurons in this module; same as num_neurons.
property
¶
Whether fit_input_stats would do anything for this neuron.
Lets the trainer skip the per-layer statistics pass entirely when no neuron in the layer consumes it.
property
¶
Whether this neuron wants a row sample of the layer input.
property
¶
Whether this neuron can refine its initialization from a streaming pass over the split (mini-batch k-means for RBF centers).
Attach per-neuron output projection parameters (xavier weight, zero bias).
proj_num_classes is stored in params_metadata so from_checkpoint_metadata can reconstruct the right parameter shape before load_state_dict runs.
classmethod
¶
Rebuild a neuron module of the saved class from checkpoint metadata.
Looks the class up in the registry by metadata["cls"] and constructs
it with the saved shapes, including the per-neuron projection when the
checkpoint has one. The tensors are not copied here; load_state_dict
fills them in afterwards.
Move or cast the module, its src_idxs and created_neuron_idxs included.
Re-initialize weight with the module's init_method.
'xavier' uses nn.init.xavier_uniform_; 'uniform' draws from
U(-0.1, 0.1). Any other method raises NotImplementedError.
abstractmethod
¶
Human-readable description of the neuron's formula, for logs and plots.
abstractmethod
¶
Short family name, for logs and plot labels.
Whether the bias-error machinery is needed for criterion_type.
False only for the plain 'validate' criterion.
Evaluate every candidate neuron on a batch of layer input.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
inp
|
Tensor
|
Layer input, (batch, num_feat). |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
(batch, num_neurons) neuron outputs, or (batch,) when the weights are 1-D (inside the trainer's vmap over candidates). |
Choose the input pairs the candidate neurons read.
With max_neuron_models set, draws that many distinct unordered pairs
at random (fewer if fewer exist); otherwise enumerates all
num_feat * (num_feat - 1) / 2 pairs in order.
Returns:
| Type | Description |
|---|---|
tuple of (torch.Tensor, int)
|
The (num_neurons, 2) index tensor and |
abstractmethod
¶
Expand gathered inputs into the design row the weights multiply.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
x
|
Tensor
|
(batch, num_neurons, dim) inputs of each candidate. |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
(batch, num_neurons, num_w) design rows. |
Calibrate anything data-dependent from layer-input statistics.
mean / std are per-layer-input-feature vectors of length
num_feat, measured over the training set before this layer trains.
The base implementation does nothing — see
BaseOrthogonalNeuron.fit_input_stats for the sigma squash.
fit_input_sample(x_sample: Tensor, stream: bool = False, seed: int | None = None, iters: int = 20) -> None
Initialize from a (N, num_feat) sample of the layer input.
With stream=True the split exceeded train.input_sample_rows and
the trainer will follow with stream_input_batch over the whole
split, so only the start should be taken from the sample. seed
(from train.seed and the layer index) makes any sampling inside
reproducible; iters is train.rbf_kmeans_iters.
One batch (B, num_feat) of layer input during the streaming pass.
One-line summary of this family's survivors after selection, for the trainer's log (e.g. how far learnable centers moved). None when the family has nothing to report.
Keep only the candidates at positions idxs, dropping the rest.
Selects the matching rows of src_idxs, created_neuron_idxs,
weight and the projection parameters, then calls _prune_extra for
subclass state.
BaseTupleNeuron(num_feat: int, num_src_feat: int, activation: ActivationLike, layer_index: int, start_index: int, max_neuron_models: int | None = None, init_method: str = 'xavier', dim: int = 2)
Bases: BasePolynomNeuron
A neuron over an unordered dim-tuple of inputs (dim >= 2).
The pair ("binary") families hard-code dim = 2 through the base
create_src_idxs; this class generalizes the candidate enumeration to
dim-tuples (pairs at dim=2, triplets at dim=3, ...) for families whose
design row is symmetric over its input slots — permuted tuples reach the
same least-squares fit, so unordered tuples suffice (cap C(n, dim), not
P(n, dim)). It carries no state of its own: per-slot calibration (the
orthogonal families' squash, an RBF family's centers) lives in the
subclass and is fed through the input-pass hooks of BasePolynomNeuron,
with _prune_extra keeping every (num_neurons, ...) tensor aligned.
Subclasses must set self.num_w before calling __init__ (the base
allocates weight from it) and may read self.dim before that, since
it is assigned here first.
Choose the unordered dim-tuples of inputs the candidates read.
With max_neuron_models set, draws that many distinct tuples at random
(fewer if fewer exist); otherwise enumerates every combination with
itertools.combinations, which is empty when num_feat < dim.
Returns:
| Type | Description |
|---|---|
tuple of (torch.Tensor, int)
|
The (num_neurons, dim) index tensor and |
LinearPolynomNeuron(num_feat: int, num_src_feat: int, activation: ActivationLike, layer_index: int, start_index: int, dim: int = 2, max_neuron_models: int | None = None, init_method: str = 'xavier')
Bases: BasePolynomNeuron
Pair neuron y = w0 + w1xi + w2xj (reference function 'linear').
LinearCovPolynomNeuron(num_feat: int, num_src_feat: int, activation: ActivationLike, layer_index: int, start_index: int, dim: int = 2, max_neuron_models: int | None = None, init_method: str = 'xavier')
Bases: BasePolynomNeuron
Pair neuron y = w0 + w1xi + w2xj + w3xixj (reference function 'linear_cov').
QuadraticPolynomNeuron(num_feat: int, num_src_feat: int, activation: ActivationLike, layer_index: int, start_index: int, dim: int = 2, max_neuron_models: int | None = None, init_method: str = 'xavier')
Bases: BasePolynomNeuron
Pair neuron with every monomial up to degree 2 (reference function 'quadratic').
CubicPolynomNeuron(num_feat: int, num_src_feat: int, activation: ActivationLike, layer_index: int, start_index: int, dim: int = 2, max_neuron_models: int | None = None, init_method: str = 'xavier')
Bases: BasePolynomNeuron
Pair neuron with powers up to 3 and one cross term (reference function 'cubic').
The row has no mixed cubic terms (xi^2xj, xixj^2).
PolyQuadratic(num_feat: int, num_src_feat: int, activation: ActivationLike, layer_index: int, start_index: int, max_neuron_models: int | None = None, init_method: str = 'xavier', dim: int = 2, squares: bool = True)
Bases: BasePolynomNeuron
Quadratic neuron over dim inputs (reference function 'polyquad').
y = w0 + sum_i w_i x_i + sum_{i<=j} w_ij x_i x_j
With squares=False the squares x_i^2 are left out, keeping only the
cross terms i < j. Candidates read unordered dim-tuples of inputs.
Compute the degree-2 polynomial of each neuron's dim inputs:
y = w0 + sum_i w1_ix_i + sum_{i<=j} w2_ijx_i*x_j
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
inp
|
Tensor
|
Layer input, (..., num_feat). Each neuron reads its |
required |
Returns:
| Type | Description |
|---|---|
Tensor
|
(..., num_neurons) neuron outputs, or (...,) when the weights are 1-D (inside the trainer's vmap over candidates). |
Describe the polynomial: full, or cross terms only when squares are off.
Not used: forward expands the terms itself. Raises NotImplementedError.
Choose the unordered dim-tuples of inputs the candidates read.
With max_neuron_models set, draws that many distinct tuples at random
(fewer if fewer exist). Without it, enumerates all pairs; full
enumeration is only implemented for dim == 2.
Returns:
| Type | Description |
|---|---|
tuple of (torch.Tensor, int)
|
The (num_neurons, dim) index tensor and |
BaseOrthogonalNeuron(num_feat: int, num_src_feat: int, activation: ActivationLike, layer_index: int, start_index: int, max_neuron_models: int | None = None, init_method: str = 'xavier', degree: int = 3, cross: bool = True, squash: bool = True, dim: int = 2, squash_method: str = 'sigma', squash_n_sigma: float = 2.0, squash_core_range: float = 0.75)
Bases: BaseTupleNeuron
Shared machinery for neurons built on an orthogonal-polynomial basis.
For dim inputs (xi, xj, ...) — optionally squashed to (ui, uj, ...) —
expanded to degree d, the design row is
[ 1,
P_1(u1), ..., P_1(un), P_2(u1), ..., P_d(un), # n*d univariate columns
u_i*u_j for each input pair (i, j) ] # only when cross=True
At dim=2 this is exactly CubicPolynomNeuron's layout when d=3, with the power
basis replaced by the orthogonal one. Subclasses supply the family purely
through _recurrence_coeffs, the three-term recurrence multipliers.
property
¶
True when the neuron squashes with the sigma method and so needs input statistics.
Calibrate the sigma squash from this layer's input statistics.
mean / std are per-layer-input-feature vectors of length
num_feat, measured over the training set. Indexing them with
src_idxs turns them into the (num_neurons, dim) per-input-slot stats
the squash needs, so neuron k's slot j is normalized by the statistics
of exactly the feature column it consumes.
No-op unless this neuron actually squashes with the sigma method.
Design row: 1, P_1..P_degree of every (squashed) input, then the
pairwise products u_i*u_j when cross is on.
LegendrePolynomNeuron(num_feat: int, num_src_feat: int, activation: ActivationLike, layer_index: int, start_index: int, max_neuron_models: int | None = None, init_method: str = 'xavier', degree: int = 3, cross: bool = True, squash: bool = True, dim: int = 2, squash_method: str = 'sigma', squash_n_sigma: float = 2.0, squash_core_range: float = 0.75)
Bases: BaseOrthogonalNeuron
Neuron on the Legendre basis P_k (uniform-weight orthogonal on [-1, 1]).
Bonnet's recursion: P_k = ((2k-1)/k) u P_{k-1} - ((k-1)/k) P_{k-2}, with P_0 = 1, P_1 = u.
ChebyshevPolynomNeuron(num_feat: int, num_src_feat: int, activation: ActivationLike, layer_index: int, start_index: int, max_neuron_models: int | None = None, init_method: str = 'xavier', degree: int = 3, cross: bool = True, squash: bool = True, dim: int = 2, squash_method: str = 'sigma', squash_n_sigma: float = 2.0, squash_core_range: float = 0.75)
Bases: BaseOrthogonalNeuron
Neuron on the Chebyshev-(first-kind) basis T_k (minimax on [-1, 1]).
Recurrence: T_k = 2 u T_{k-1} - T_{k-2}, with T_0 = 1, T_1 = u.
RBFNeuron(num_feat: int, num_src_feat: int, activation: ActivationLike, layer_index: int, start_index: int, max_neuron_models: int | None = None, init_method: str = 'xavier', dim: int = 2, centers: int = 16, placement: str = 'kmeans', seeding: str = 'pca_quantiles', width: float = 1.0, learn_centers: bool = True, learn_widths: bool = True, center_radius: float | None = 2.0, width_band: float = 4.0, normalize: bool = True, linear: bool = True, standardize: bool = True)
Bases: BaseTupleNeuron
Gaussian RBF neuron over a dim-tuple of inputs, see the module docstring.
property
¶
Always True: the inputs are standardized and the width floor is in raw units.
property
¶
Always True: the centers are placed on a sample of the layer input.
property
¶
True for k-means placement, which can be refined by a streaming pass.
exp(L * tanh(log_width / L)), the factor on width0, always
inside (1 / band, band).
Effective bump widths width0 * width_scales(), shape
(num_neurons, M) eager or (M,) under vmap.
Bounded displacement of every center from its start: the raw shift
center_shift rescaled so its norm is r * tanh(|shift| / r), i.e.
the direction is kept and the length squashed below the radius r.
The norm is computed as sqrt(|shift|^2 + eps^2) so the map is smooth
at zero shift (the start), where its Jacobian is the identity. Shape
(num_neurons, M, dim) eager or (M, dim) under vmap. The raw shift
when center_radius is None.
Design row: the M bump activations of the standardized inputs, then
the inputs themselves when linear is on, then a constant 1 when
normalize is off.
Store the per-slot input mean and std used for standardization.
mean / std are per layer-input feature; indexing them with
src_idxs gives each candidate's slot statistics. A constant slot
(std <= 1e-8) gets std 1.
fit_input_sample(x_sample: Tensor, stream: bool = False, seed: int | None = None, iters: int = 20) -> None
Place the centers on a sample of the layer input.
Grid placement builds the quantile grid. k-means placement seeds with
k-means++ or PCA quantiles, then runs iters Lloyd iterations on the
sample, batched over the candidates. With stream=True only the seeds
are stored, and stream_input_batch refines them over the whole split.
Mini-batch k-means update of the centers from one batch of layer input.
Each center moves toward the mean of its assigned rows with the
per-center rate 1 / n_m (Sculley). No-op unless fit_input_sample ran
with stream=True.
Finish the streaming pass: fix the centers, set widths and radii, and log how far the centers moved from their seeds.
One line after selection: how far the survivors' centers moved from their placed start (and how many sit at their radius) and where their width scales sit, so a fit that never moves a center, or pins everything at its bound, is visible.
Describe the centers, their placement, what is learnable, width and inputs.
generate_unique_pairs(n: int, max_neuron_models: int, seed: int | None = None, allow_self: bool = False, ordered: bool = True) -> list[tuple[int, int]]
Draw max_neuron_models distinct random index pairs from range(n).
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
n
|
int
|
Number of available inputs. |
required |
max_neuron_models
|
int
|
Number of pairs to draw. Clamped, with a warning, to the number of
distinct pairs that exist for |
required |
seed
|
int
|
Reseeds Python's global |
None
|
allow_self
|
bool
|
Allow pairs |
False
|
ordered
|
bool
|
When False, |
True
|
Returns:
| Type | Description |
|---|---|
list of tuple of int
|
The pairs, in set iteration order. |
generate_unique_combinations(n: int, num_inputs: int, max_neuron_models: int, seed: int | None = None, allow_self: bool = False, ordered: bool = True) -> list[tuple[int, ...]]
Generate unique k-tuples (pairs, triplets, etc.) of indices.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
n
|
int
|
Number of available inputs. |
required |
num_inputs
|
int
|
Size of each tuple (e.g. 2=pair, 3=triplet, etc.). |
required |
max_neuron_models
|
int
|
Maximum number of unique combinations to generate. |
required |
seed
|
int
|
Random seed for reproducibility. |
None
|
allow_self
|
bool
|
Whether repeated indices in a tuple are allowed. |
False
|
ordered
|
bool
|
Whether order matters (permutations) or not (combinations). |
True
|
Returns:
| Type | Description |
|---|---|
list of tuple of int
|
List of unique tuples. |