3.8. Regularization and Normalization Layers

This page covers the layers that reshape the signal in the network instead of learning a mapping. These are dropout and its spatial variants, the 2 Gaussian noise layers, and the 5 normalization layers: batch, layer, group, instance, and unit. It closes with Rescaling, the fixed affine map that conditions an input before the first weight. Every layer here except Rescaling lives in rustyml::neural_network::layers::regularization, and Rescaling lives beside them in rustyml::neural_network::layers.

Almost every layer here behaves differently in training mode than in inference mode. The crate factors that split into 1 shared mechanism. Learn that mechanism first, because the rest of the page is only detail. Skip it, and batch normalization in particular can quietly compute the wrong numbers. UnitNormalization and Rescaling are the 2 exceptions, and sections 3.8.9 and 3.8.10 say why.

Every layer here plugs into a SequentialBuilder with the same .add(...) call, just like Dense Layers and Activations. No constructor here takes an input shape, because the model build supplies it. Every constructor except Rescaling::new returns a Result. See Error Handling for the reason.

Import the layers with use rustyml::neural_network::layers::*;. This re-exports Dropout, SpatialDropout1D/2D/3D, GaussianNoise, GaussianDropout, BatchNormalization, LayerNormalization (with its LayerNormalizationAxis), GroupNormalization, InstanceNormalization, UnitNormalization (with its UnitNormalizationAxis), and Rescaling. The context type is rustyml::neural_network::Ctx, and the layer methods live in rustyml::neural_network::traits. As everywhere in this crate, a Tensor is ndarray::ArrayD<f32>. It uses single precision and a dynamic rank.

3.8.1. Training mode versus inference mode

No layer holds a training flag. The mode travels with the pass, inside a context that the caller builds. The UnaryLayer trait gives every layer with 1 input these 3 methods:

fn forward(&self, input: &Tensor, ctx: &mut Ctx) -> Result<Tensor, Error>;
fn backward(&self, grad_output: &Tensor, ctx: &mut Ctx) -> Result<Tensor, Error>;
fn forward_mut(&mut self, input: &Tensor, ctx: &mut Ctx) -> Result<Tensor, Error>;

forward and backward take &self. The layer computes, and the context remembers. A Ctx holds 4 channels. These are the training flag, the cache that a forward pass writes for its backward pass, the gradient store, and the state channel. The state channel carries the non-trainable values that a training pass changes, such as running statistics and a random stream.

Ctx::training() builds the context of a training pass, and Ctx::inference() builds the context of an inference pass. A training pass parks what its backward pass needs, and it proposes every non-trainable value it changes. An inference pass parks nothing and proposes nothing. Every mode-dependent layer then takes its inference behavior.

forward_mut is the entry point for a caller that drives 1 layer by hand. It builds the layer from the tensor when the layer holds no build, runs the forward pass, and moves the proposed state into the layer. Plain forward leaves that last step undone, and it refuses an unbuilt layer with NnError::NotBuilt. Without the step, the random stream of a dropout layer never advances, and 2 calls draw the same mask.

No gradient lives in a layer. The backward pass adds each one to ctx.grads(), under the address ParamId::new(scope, name). The scope is the position of the layer in the model, and the name is what the layer calls the tensor, such as "gamma". LayerBase::parameters_mut yields the trainable tensors themselves, and it holds no gradient at all.

Inside a Sequential model, none of this needs attention. fit and fit_with_batches build a training context for each batch, and they move the state of every layer back after the forward pass. predict and evaluate build an inference context instead. predict borrows &self and writes nothing, which is why 1 compiled model serves concurrent requests across threads. Watch for 1 trap: a training context used for inference. A BatchNormalization layer then reads statistics from your test batch and proposes new running averages. For anything that is not a training step, call model.predict(...), or pass Ctx::inference(). The table below shows what each layer family does in each mode:

Note the bottom row. Layer, group, and instance normalization compute their statistics from the current input in every mode. The output of an inference pass therefore equals the output of a training pass, bit for bit. Only the backward pass depends on the mode. An inference context makes it return the upstream gradient unchanged instead of the real gradient. Batch normalization is the only layer here that keeps state, which is the running mean and the running variance. It is also the only one that computes something different between the 2 modes.

This split also explains why 1 model can report 2 different losses. The per-epoch figures in the History from fit and fit_with_batches come from training passes. In a training context, dropout zeros a fraction of the activations and rescales the survivors. Batch normalization normalizes with the statistics of each batch and folds them into its running averages. evaluate runs the inference path instead. Dropout is the identity there, and batch normalization reads the running statistics without changing them. On a model that holds either layer, the 2 numbers do not agree, and neither number is wrong. The training figure is the loss of a deliberately handicapped network, measured before the own update of that batch. evaluate scores the network you actually hold. Use evaluate to select checkpoints and to stop training early. Call it at any point in training. It borrows &self, updates nothing, and draws no random numbers. So it cannot consume a dropout mask, and it cannot shift the shuffle order of the run it measures.

3.8.2. Dropout

Dropout::new(rate) builds the classic layer, and rate is its only argument. rate is the fraction of units to zero. It must lie in [0.0, 1.0] inclusive, or the call returns Error::InvalidParameter.

Dropout accepts an input of any shape, at rank 1 or higher. It owns no array, and it reads no extent of its input. Nothing is left for a shape to size, and nothing can disagree with a shape. 1 instance therefore serves a rank-2 batch and a rank-4 feature map alike. GaussianNoise and GaussianDropout are the same in this respect. The build records the shape it receives, and the layer checks no later input against it. The old shape argument, and the empty-vector wildcard that turned its check off, are both gone together.

The implementation is inverted dropout. A training pass samples a uniform mask and keeps each unit with probability 1 - rate. It then scales the survivors by 1 / (1 - rate), which keeps the expected activation unchanged. That scaling is the whole point. An inference pass becomes a pure identity operation, and serving code needs no rescaling step. The 2 boundary values short-circuit this logic. rate == 0.0 is the identity. rate == 1.0 zeros every unit.

use ndarray::Array;
use rustyml::neural_network::Ctx;
use rustyml::neural_network::layers::*;
use rustyml::neural_network::traits::UnaryLayer;

fn main() {
    // 50% rate. Seed the mask so this run is reproducible.
    let mut dropout = Dropout::new(0.5).unwrap().with_random_state(7);
    let input = Array::ones((4, 8)).into_dyn();

    // Training: about half the units become 0, the survivors become 1/(1-0.5) = 2.0.
    // forward_mut builds the layer from the tensor and completes the pass.
    let mut ctx = Ctx::training();
    let train_out = dropout.forward_mut(&input, &mut ctx).unwrap();
    let kept = train_out.iter().filter(|&&v| v != 0.0).count();
    println!("kept {kept}/{} units, each rescaled to 2.0", input.len());

    // Inference: inverted dropout makes the layer the identity, and nothing is cached.
    let mut eval = Ctx::inference();
    assert_eq!(dropout.forward(&input, &mut eval).unwrap(), input);
    assert_eq!(eval.pending_caches(), 0);

    // The layer reads no extent, so the same instance takes any shape.
    let volume = Array::ones((2, 3, 5, 9)).into_dyn();
    let out = dropout.forward(&volume, &mut eval).unwrap();
    assert_eq!(out.shape(), &[2, 3, 5, 9]);
}

with_noise_shape(shape) makes the draw coarser than 1 value per element. It takes 1 entry per axis. An entry of Some(1) gives the whole axis 1 shared draw, so the same units drop at every position of that axis. An entry of None takes the extent of the input on that axis, which keeps the draws there independent. Any other entry must equal the extent of the input on that axis. A vector shorter than the rank of the input lines up against the last axes. That is the usual right-aligned broadcast rule, so the axes it leaves out share 1 draw. A caller that omits the batch axis therefore shares 1 mask across the whole batch.

The layer samples the mask at that shape, keeps it at that shape for the backward pass, and broadcasts it in both directions. It never builds a full-size copy. Some(0) gives Error::InvalidParameter, and an empty vector gives Error::EmptyInput. Use the method on sequence data, where vec![None, Some(1), None] drops the same features at every timestep and gives the recurrent variant of dropout. SpatialDropout1D is the equivalent for a feature map, and it needs no argument.

use ndarray::Array;
use rustyml::neural_network::Ctx;
use rustyml::neural_network::layers::*;
use rustyml::neural_network::traits::UnaryLayer;

fn main() {
    // (batch = 2, timesteps = 4, features = 8).
    let input = Array::ones((2, 4, 8)).into_dyn();

    // 1 shared draw along the time axis: the same features drop at every timestep.
    let mut dropout = Dropout::new(0.5)
        .unwrap()
        .with_noise_shape(vec![None, Some(1), None])
        .unwrap()
        .with_random_state(11);
    let mut ctx = Ctx::training();
    let out = dropout.forward_mut(&input, &mut ctx).unwrap();
    assert_eq!(out.shape(), &[2, 4, 8]);

    for n in 0..2 {
        for f in 0..8 {
            let first = out[[n, 0, f]];
            assert!((0..4).all(|t| out[[n, t, f]] == first));
        }
    }

    // A vector shorter than the rank lines up against the last axes, so the axes it
    // leaves out share 1 draw. This one shares 1 mask across the whole batch.
    let mut shared = Dropout::new(0.5)
        .unwrap()
        .with_noise_shape(vec![Some(1), None])
        .unwrap()
        .with_random_state(11);
    let mut ctx = Ctx::training();
    let batch_out = shared.forward_mut(&input, &mut ctx).unwrap();
    for t in 0..4 {
        for f in 0..8 {
            assert_eq!(batch_out[[0, t, f]], batch_out[[1, t, f]]);
        }
    }

    println!("noise-shape masks behaved as documented");
}

The layer draws the mask from a per-layer StdRng. By default Dropout::new seeds it from the global seed, or from OS entropy when no global seed is set. with_random_state(seed) re-seeds it deterministically. A forward pass takes &self, so it draws from a copy of the stream and leaves the advanced copy in the state channel. forward_mut moves that copy into the layer, and a Sequential training step does the same. 2 plain forward calls with fresh contexts therefore draw the same mask, and 2 forward_mut calls draw different masks. Seeding fixes the sequence across process runs. See Reproducibility and Random Seeds for how the global seed threads through every randomized component. Dropout has no trainable parameters. A backward pass in a training context with no forward pass behind it returns NnError::ForwardPassNotRun. In an inference context, or at rate == 0.0, the backward pass simply passes the gradient through.

3.8.3. Spatial dropout for feature maps

Plain dropout is a poor fit for convolutional feature maps. Adjacent pixels in a channel correlate strongly, so zeroing scattered individual elements removes almost no information. The value of a dropped pixel is nearly recoverable from its neighbors, so the regularizing effect washes out. SpatialDropout1D, SpatialDropout2D, and SpatialDropout3D fix this by dropping an entire channel at a time. When a channel drops, all of its spatial positions go to zero together. This forces the network to avoid a dependence on any single feature map.

SpatialDropout1D::new(rate), SpatialDropout2D::new(rate), and SpatialDropout3D::new(rate) each take the rate alone. The 3 layers differ only in the expected rank. All 3 use the channels-last layout. SpatialDropout1D expects a 3-D (batch, length, channels) input. SpatialDropout2D expects a 4-D (batch, height, width, channels) input. SpatialDropout3D expects a 5-D (batch, depth, height, width, channels) input. A wrong rank returns Error::InvalidInput.

Internally, each layer samples 1 keep/drop value per (batch, channel) pair, which is a tiny [batch, channels] mask. It applies the same 1 / (1 - rate) inverted-dropout scale across the whole channel, so the layer never builds a full-size mask. Construction, seeding with with_random_state, the boundary behaviors, and the ForwardPassNotRun contract match plain Dropout. The error even names the concrete layer. A premature backward call on a SpatialDropout2D layer returns ForwardPassNotRun("SpatialDropout2D").

use ndarray::Array;
use rustyml::neural_network::Ctx;
use rustyml::neural_network::layers::*;
use rustyml::neural_network::traits::UnaryLayer;

fn main() {
    // (batch=1, height=4, width=4, channels=8): whole channels drop as a unit.
    let mut sd = SpatialDropout2D::new(0.5).unwrap().with_random_state(3);

    let input = Array::ones((1, 4, 4, 8)).into_dyn();
    let mut ctx = Ctx::training();
    let out = sd.forward_mut(&input, &mut ctx).unwrap();

    // Every position within a channel shares 1 value: 0.0 (dropped) or 2.0 (kept & rescaled).
    for c in 0..8 {
        let first = out[[0, 0, 0, c]];
        assert!((0..4).all(|h| (0..4).all(|w| out[[0, h, w, c]] == first)));
        println!("channel {c}: all 16 spatial positions hold {first}");
    }
}

3.8.4. Gaussian noise and Gaussian dropout

GaussianNoise and GaussianDropout regularize by an injection of noise instead of by zeroing values. GaussianNoise::new(stddev) is additive. A training pass adds zero-mean N(0, stddev^2) noise, so output = input + noise. This is a data-augmentation style of regularization. It perturbs inputs without a change of their expected value, but it does inflate their variance. With a large enough stddev, a positive input can become negative. stddev must be non-negative and finite. Construction rejects a negative, NaN, or Inf value. This check is deliberate. Without it, the sampler would panic on the first forward call. The backward pass is a pure pass-through in every mode, because the noise does not depend on the input, so d(x + noise)/dx = 1. It caches nothing and never returns ForwardPassNotRun.

GaussianDropout::new(rate) is multiplicative. It is the closer analogue to plain dropout. It multiplies each input by a sample from N(1, sigma), where sigma = sqrt(rate / (1 - rate)). The mean-1 noise leaves E[output] = input. This is the same expectation-preserving idea as inverted dropout, but it uses continuous multipliers instead of hard zeros. Here rate must be in [0.0, 1.0), exclusive of 1, because sigma diverges as rate approaches 1. Unlike GaussianNoise, this layer parks the exact noise draw in the context, so its backward pass can reuse it. Because y = x * noise, dx = grad * noise. A backward call in a training context before any training forward call returns ForwardPassNotRun. Both layers are the identity in an inference context, and also when stddev or rate is 0. Neither layer has trainable parameters, and neither one reads an extent of its input. So both take an input of any shape, at rank 1 or higher.

use ndarray::Array;
use rustyml::neural_network::Ctx;
use rustyml::neural_network::layers::*;
use rustyml::neural_network::traits::UnaryLayer;

fn main() {
    let input = Array::from_elem((2, 4), 3.0_f32).into_dyn();

    // Additive: output = input + N(0, 0.5^2). Mean preserved, variance added.
    let mut noise = GaussianNoise::new(0.5).unwrap().with_random_state(1);
    let mut ctx = Ctx::training();
    let noisy = noise.forward_mut(&input, &mut ctx).unwrap();

    // Multiplicative: output = input * N(1, sqrt(rate/(1-rate))). E[output] stays at input.
    let mut gdrop = GaussianDropout::new(0.3).unwrap().with_random_state(1);
    let mut ctx = Ctx::training();
    let scaled = gdrop.forward_mut(&input, &mut ctx).unwrap();

    println!("additive[0] = {}, multiplicative[0] = {}", noisy[[0, 0]], scaled[[0, 0]]);

    // Both are the identity at serving time.
    let mut eval = Ctx::inference();
    assert_eq!(noise.forward(&input, &mut eval).unwrap(), input);
    assert_eq!(gdrop.forward(&input, &mut eval).unwrap(), input);
}

3.8.5. The normalization layers at a glance

4 of the 5 normalization layers share 1 skeleton. Each of those 4 subtracts a mean and divides by a standard deviation computed over some set of axes. It then applies a learnable per-channel affine transform, which is gamma * x_normalized + beta. gamma initializes to ones, and beta initializes to zeros. The optimizer treats both as trainable parameters, marked no-decay, so weight decay skips them. Weight decay should not pull the scale and the shift toward zero.

Every layer also takes an epsilon value, such as 1e-5, added under the square root for numerical stability. This epsilon also makes a zero-variance input produce a finite all-zero output instead of a NaN.

What separates those 4 layers is only which axes the statistics reduce over. UnitNormalization is the odd one out. It subtracts no mean, holds no parameter, and takes no epsilon argument. Consider an input shaped [N, ...spatial, C], with batch N leading and channels C trailing:

Use the “depends on batch?” column as the practical guide for the choice of a layer. Batch normalization couples samples together through shared statistics. This makes it effective, but also fragile at small batch sizes. The other 4 layers normalize each sample independently, so batch size does not affect them, unless you point UnitNormalization at the batch axis.

All 5 layers preserve the input shape. The shape guards below are not optional. Group and instance normalization need rank 3 or higher, and a 2-D input returns Error::InvalidInput. Batch normalization accepts rank 1 and higher, with the rank-1 case as a special channel-free fallback. Section 3.8.6 gives the detail.

3.8.6. BatchNormalization

BatchNormalization::new(momentum, epsilon) takes the 2 scalars and nothing else. momentum must be in [0.0, 1.0]. epsilon must be positive. The build supplies the shape. Dimension 0 is the batch, and the last dimension is the channel or feature axis. The per-channel parameters take the extent of that last axis, so a build shape of (None, 8) gives 4 arrays of length 8.

For a 2-D [N, C] input, this is ordinary per-feature batch normalization. For a rank-3-or-higher [N, ...spatial, C] input, the statistics reduce over the batch and every spatial position. This gives 1 mean, 1 variance, 1 gamma, and 1 beta per channel, which is spatial batch normalization. Both cases run the same code path. Because the channel axis is innermost, [N, ...spatial, C] already is the [M, C] matrix the per-channel folds read. A collapse of the leading axes reinterprets the same bytes instead of a reshape of them.

A rank-1 build shape, such as Shape::known(&[4]), is a special case with no channel axis. It uses length-1 scalar parameters broadcast over the whole input.

The state that makes batch normalization mode-dependent is a pair of running statistics. A training forward pass first normalizes with the mean and the variance of the current batch. It then proposes running = running * momentum + batch * (1 - momentum) through the state channel of the context. LayerBase::apply_state moves the 2 new arrays into the layer, and a Sequential training step calls it after every forward pass. forward_mut calls it too, for a layer that a caller drives by hand. A high momentum, such as 0.99, weights the history heavily and moves the running estimate slowly. Some other libraries define the constant the other way round, and weight the new batch instead. A value copied from such a library produces the opposite behavior here.

BatchNormalization is the 1 layer in the crate that reports both kinds of parameter. For C channels it reports 2 * C trainable elements, which are gamma and beta, and 2 * C non-trainable elements, which are moving_mean and moving_variance. Those 4 names are also the names a checkpoint addresses the arrays by. param_count() therefore returns ParamCounts::new(2 * C, 2 * C), and summary prints a total of 4 * C for the layer, split evenly between its 2 columns. The moving statistics move on every training forward pass, but no optimizer ever touches them.

with_center(false) and with_scale(false) drop beta and gamma. All 4 normalization layers with an affine transform carry the pair. A dropped array leaves the layer, so it leaves the parameter count and the checkpoint with it. The 2 moving arrays stay in every case, because inference reads them whatever the 2 flags say:

use rustyml::neural_network::Shape;
use rustyml::neural_network::layers::ParamCounts;
use rustyml::neural_network::layers::regularization::normalization::batch_normalization::BatchNormalization;
use rustyml::neural_network::traits::{LayerBase, UnaryLayer};

fn main() {
    // The build settles the 8 channels: gamma and beta are trainable, the 2 moving
    // arrays are not.
    let mut full = BatchNormalization::new(0.99, 1e-5).unwrap();
    full.build(&Shape::with_free_batch(&[6, 8])).unwrap();
    assert_eq!(full.param_count(), ParamCounts::new(16, 16));
    for entry in full.weights() {
        println!("{:<16} {:?} {:?}", entry.name, entry.kind, entry.value.shape());
    }

    // `with_scale(false)` drops gamma, and `with_center(false)` drops beta.
    let mut plain = BatchNormalization::new(0.99, 1e-5)
        .unwrap()
        .with_scale(false)
        .with_center(false);
    plain.build(&Shape::with_free_batch(&[6, 8])).unwrap();
    assert_eq!(plain.param_count(), ParamCounts::new(0, 16));
    let paths: Vec<&str> = plain.weights().iter().map(|w| w.name).collect();
    println!("paths without gamma and beta: {paths:?}");
}

Where a ReLU follows the layer, set scale to false, because the rectifier absorbs a positive scale. Where the next layer has a bias of its own, set center to false.

momentum == 0.0 discards history entirely, so the running statistics equal the statistics of the last batch. momentum == 1.0 freezes the running statistics at their initial values, which are mean 0 and variance 1. This quietly turns inference-mode normalization into a near-identity operation, usually by mistake. In an inference context, the forward pass normalizes with the running statistics instead of the batch. This is exactly why a test batch pushed through a training context corrupts the model.

Placement relative to the activation matters. The fused-activation design of Dense shapes your options. The original batch-normalization paper places it before the nonlinearity. Dense::new(.., Activation::ReLU) folds the activation into the linear layer. So “normalization before activation” needs a linear Dense layer, then BatchNormalization, then a separate activation layer:

use rustyml::neural_network::layers::activation::relu::ReLU;
// Dense(linear) -> BatchNorm -> ReLU  (BN before the nonlinearity, paper ordering)
builder
    .add(Dense::new(8, Activation::Linear).unwrap())
    .add(BatchNormalization::new(0.99, 1e-5).unwrap())
    .add(ReLU::new());

A common alternative applies the activation first, then normalizes. This is just Dense::new(.., Activation::ReLU) followed by BatchNormalization. Both orderings train fine. Pick one and stay consistent. The example below shows a full runnable model, with the simpler activation-in-the-Dense form:

use rustyml::neural_network::Shape;
use rustyml::neural_network::layers::*;
use rustyml::neural_network::sequential::SequentialBuilder;
use rustyml::neural_network::optimizers::Adam;
use rustyml::neural_network::losses::MeanSquaredError;
use ndarray::Array;

fn main() {
    // The build threads the 8 channels of the Dense layer into the normalization layer.
    let x = Array::ones((6, 4)).into_dyn();
    let y = Array::ones((6, 1)).into_dyn();

    let mut model = SequentialBuilder::new()
        .add(Dense::new(8, Activation::ReLU).unwrap())
        .add(BatchNormalization::new(0.9, 1e-5).unwrap())
        .add(Dropout::new(0.3).unwrap())
        .add(Dense::new(1, Activation::Linear).unwrap())
        .build(&Shape::with_free_batch(x.shape()))
        .unwrap();
    model.compile(
        Adam::new(0.01, 0.9, 0.999, 1e-8, 0.0).unwrap(),
        MeanSquaredError::new(),
    );

    // fit() runs the training path (batch stats, running-stat updates, dropout on).
    // predict() runs the inference path (running stats, dropout as identity).
    model.fit(&x, &y, 3).unwrap();
    println!("prediction shape: {:?}", model.predict(&x).unwrap().shape());
}

The batch axis of the build shape is free, and the layer never checks it. A BatchNormalization layer built for (None, 8) accepts a batch of any size. So fit_with_batches trains at any batch_size, including the short final chunk left over when the dataset does not divide evenly. That short chunk does not skew the reported loss of the epoch either. fit_with_batches weights each batch by the number of samples it holds, instead of 1 vote per batch. So the History entry stays the dataset-wide mean per-sample loss, however the split lands.

The check does enforce the rank and every per-sample axis. So an [n, 16] input into a layer built for 8 channels, and an input of the wrong rank, are both refused. This layer owns 4 arrays sized from the channel count, which is why it is not 1 of the 3 layers that accept any shape.

You can also inject known weights directly with set_weights(gamma, beta, moving_mean, moving_variance). You must build the layer first, because there is nothing to check a shape against before that, and an unbuilt layer returns NnError::NotBuilt. A shape mismatch returns NnError::WeightShape. The first 2 arguments are optional. Pass None for an array the layer does not hold, and pass the array for one it does, or the call returns Error::InvalidParameter. This is how a loaded model restores its inference statistics. See the next section and Saving and Loading Weights.

3.8.7. LayerNormalization

LayerNormalization::new(epsilon) normalizes across features within each sample, and the build supplies the shape it sizes gamma and beta from. It has no batch coupling and no running statistics. This is exactly why it is the normalization of choice for recurrent models. It also suits any setting where the batch is small, variable, or size 1. The statistics of batch normalization get noisy or meaningless with a handful of samples. Layer normalization is unaffected, because every sample stands alone. Its forward output is the same in a training context and in an inference context. Only its backward pass depends on the mode.

By default, it normalizes the last (feature) dimension. with_normalized_axis(...) changes this. It takes a LayerNormalizationAxis value: Default (the last axis), Custom(axis) for a single other axis, or Multiple(vec![...]) to normalize jointly over several axes at once. Multiple takes the statistics over the combined elements of those axes, and gamma/beta become 1-D over their product. The build sizes gamma and beta from the axis the layer holds. Call this method before you build the model, and before you assign weights. The layer rejects an empty, duplicated, or out-of-bounds axis list.

A non-trailing Custom axis still works correctly, but it runs on a slower strided path instead of the fused row path. A Multiple list whose axes are not already trailing and in order also works correctly. It first transposes the input to bring the normalized axes together, runs the fused row path, then transposes the result back. Both cases cost more than the default trailing-axis path, but neither changes the result.

use ndarray::Array;
use rustyml::neural_network::Shape;
use rustyml::neural_network::layers::*;
use rustyml::neural_network::traits::UnaryLayer;
use rustyml::neural_network::Ctx;

fn main() {
    let input = Array::from_shape_vec((2, 4), vec![1.0, 3.0, 5.0, 7.0, 2.0, -2.0, 0.0, 4.0])
        .unwrap()
        .into_dyn();

    // Default: each row (the last axis) is normalized to mean 0, variance 1.
    let mut ln = LayerNormalization::new(1e-5).unwrap();
    ln.build(&Shape::known(&[2, 4])).unwrap();
    let mut ctx = Ctx::training();
    let out = ln.forward(&input, &mut ctx).unwrap();

    // Custom(0): normalize down each column (across the batch axis) instead.
    let mut ln_cols = LayerNormalization::new(1e-5)
        .unwrap()
        .with_normalized_axis(LayerNormalizationAxis::Custom(0))
        .unwrap();
    ln_cols.build(&Shape::known(&[2, 4])).unwrap();
    let mut ctx = Ctx::training();
    let out_cols = ln_cols.forward(&input, &mut ctx).unwrap();

    println!("row-norm {:?}, col-norm {:?}", out.shape(), out_cols.shape());

    // Statistics come from the current input, so an inference pass reproduces the
    // training output exactly, and it caches nothing.
    let mut eval = Ctx::inference();
    assert_eq!(ln.forward(&input, &mut eval).unwrap(), out);
    assert_eq!(eval.pending_caches(), 0);
}

3.8.8. GroupNormalization and InstanceNormalization

These 2 layers fill the middle ground for convolutional models where the batch is too small for batch normalization to work well. Examples include detection and segmentation backbones, and generative models. GroupNormalization::new(num_groups, epsilon) splits the channels into num_groups contiguous groups and normalizes within each group, per sample. With 1 group, it becomes layer normalization over all channels. With as many groups as channels, it becomes instance normalization. InstanceNormalization::new(epsilon) normalizes each (sample, channel) plane on its own. This is the standard choice for style transfer, because it strips per-instance contrast and leaves batch relationships alone.

Both layers require rank 3 or higher. Both take the channel axis as the trailing axis, like every other spatial layer. Like layer normalization, both carry no running statistics and are mode-independent in the forward direction. Their per-group mean and variance come from a single pass over the data. The 2 accumulators gather sums of deviations from a value taken out of the data itself. This lets the variance fall out without a second walk of the sample. It also avoids the catastrophic cancellation that a plain E[x^2] - E[x]^2 formula would hit when the mean dwarfs the spread.

Instance normalization is group normalization with num_groups set to the channel count, and the crate implements it that way. The 2 layers produce identical output when configured to match:

use ndarray::Array;
use rustyml::neural_network::Shape;
use rustyml::neural_network::layers::*;
use rustyml::neural_network::traits::UnaryLayer;
use rustyml::neural_network::Ctx;

fn main() {
    // [batch=1, positions=4, channels=4]
    let input = Array::from_shape_vec((1, 4, 4), (0..16).map(|v| v as f32).collect::<Vec<_>>())
        .unwrap()
        .into_dyn();

    // InstanceNorm normalizes every (sample, channel) plane independently...
    let mut inn = InstanceNormalization::new(1e-5).unwrap();
    inn.build(&Shape::known(&[1, 4, 4])).unwrap();
    let mut ctx = Ctx::training();
    let out_in = inn.forward(&input, &mut ctx).unwrap();

    // ...which is exactly GroupNorm with 1 group per channel.
    let mut gn = GroupNormalization::new(4, 1e-5).unwrap();
    gn.build(&Shape::known(&[1, 4, 4])).unwrap();
    let mut ctx = Ctx::training();
    let out_gn = gn.forward(&input, &mut ctx).unwrap();

    let max_diff = out_in
        .iter()
        .zip(out_gn.iter())
        .map(|(a, b)| (a - b).abs())
        .fold(0.0_f32, f32::max);
    println!("max |InstanceNorm - GroupNorm(groups=channels)| = {max_diff}");
    assert!(max_diff < 1e-6);
}

1 divisibility rule matters for GroupNormalization. num_groups must evenly divide the channel count. A constructor cannot check it, because the channel count arrives with the shape. UnaryLayer::build checks it, and it refuses a model whose channel count and group count disagree, on the line that assembles the model. The Error::InvalidParameter names the 2 numbers. GroupNormalization::new still catches num_groups == 0 on its own. Both constructors catch a non-positive or non-finite epsilon with Error::InvalidParameter up front. Both layers accept set_weights(gamma, beta), and both need a built layer to accept it.

3.8.9. UnitNormalization

UnitNormalization scales each group of elements to an L2 norm of 1. It divides by a length and subtracts no mean, so it changes how long each group is and never which way it points. That is what makes it the layer for cosine similarity, for a metric or retrieval head. It also fits anywhere a downstream dot product has to read as an angle. Among the 5 normalization layers, it is also the only one with no parameter and no epsilon argument.

UnitNormalization::new(axis) takes 1 argument, a UnitNormalizationAxis. Default is the last axis, which under the channels-last layout is the channel axis. Custom(a) names 1 axis, counting from 0 with the batch axis included. Multiple(axes) takes 1 norm over the combined elements of several axes, which need not be next to each other.

The constructor rejects an empty list and a repeated axis with Error::InvalidParameter. The constructor cannot catch an axis past the end of the input, because it sees no shape. It becomes an Error::InvalidInput at the first forward pass, as does a rank below 2. This layer owns no array, so the build allocates nothing for it and reads no extent. The forward pass accepts a layer that holds no build.

The forward output is the same in both modes, and only a training pass parks a cache. The backward pass reads that cache in every mode, unlike the other layers of this page. A backward call that follows an inference pass therefore returns NnError::ForwardPassNotRun.

The scale is 1 / sqrt(sum of squares), capped at 1e12. The cap decides the answer for a group that is all zero. The cap multiplies such a group by 1e12, and it stays all zero, rather than a division by zero. A group whose sum of squares overflows f32 goes the other way. Its reciprocal is 0, so the output of that group is also all zero.

The derivative of the reciprocal square root itself overflows at a group the cap took over. A direct differentiation through it would therefore return NaN for every element of that group. This layer instead returns the derivative of the function it actually computed, which is the cap itself. The 2 approaches agree everywhere else.

use ndarray::Array2;
use rustyml::neural_network::layers::*;
use rustyml::neural_network::traits::UnaryLayer;
use rustyml::neural_network::Ctx;

fn main() {
    // 2 samples of 3 features: the 3-4-5 and the 5-12-13 right triangles.
    let input = Array2::from_shape_vec((2, 3), vec![3.0, 4.0, 0.0, 0.0, 5.0, 12.0])
        .unwrap()
        .into_dyn();

    let unit = UnitNormalization::new(UnitNormalizationAxis::Default).unwrap();
    let mut eval = Ctx::inference();
    let out = unit.forward(&input, &mut eval).unwrap();

    // Each row is now a unit vector: 0.6, 0.8 and then 5/13, 12/13.
    for row in 0..2 {
        let length: f32 = (0..3).map(|c| out[[row, c]] * out[[row, c]]).sum();
        assert!((length - 1.0).abs() < 1e-6);
    }
    println!("first row {:?}", (0..3).map(|c| out[[0, c]]).collect::<Vec<_>>());

    // The 2 modes give the same output. Only the training pass parks a cache.
    let mut ctx = Ctx::training();
    assert_eq!(unit.forward(&input, &mut ctx).unwrap(), out);
    assert_eq!(ctx.pending_caches(), 1);
}

Which axes you pick decides how the layer reads the data. When they are the trailing block of the shape, which the default axis always is, each group is 1 contiguous run. The layer then walks equal-length rows, which costs 2 linear walks of the input, 1 for the norm and 1 for the scale.

Any other choice leaves the groups as strided lanes. The layer reduces them in place rather than a payment for a transpose in and a transpose back out. That choice costs 1 more tensor of the input size while it folds the squares. Both paths give the same numbers.

3.8.10. Rescaling: a fixed affine map

Rescaling applies y = x * scale + offset to every element. It is the 1 layer on this page that learns nothing and holds nothing. It keeps the input shape at every rank, it holds no parameter, and it reads no training mode. A training context and an inference context therefore give exactly the same output.

Rescaling::new(scale) sets the factor and leaves the offset at 0. with_offset(offset) sets the constant the layer adds after the multiplication. Neither call returns a Result, because no value is invalid here. A scale of 0 and a negative scale are both legal, and the layer applies them like any other value.

Its use is input normalization in front of a model. An 8-bit image scales into [0, 1] with Rescaling::new(1.0 / 255.0), and into [-1, 1] with Rescaling::new(1.0 / 127.5).with_offset(-1.0). The step inside the model matters more than it looks. Inference then applies exactly the map training applied, and a serving path cannot forget it. Compare this with the fitted scalers of 4.2. Standardization and Normalization, which learn their constants from the training data. Rescaling takes constants you already know.

The backward pass multiplies the incoming gradient by scale, because the derivative of the map is scale at every element. The offset is a constant, so it has no part in the gradient. The layer needs no build and stores no cache, not even the shape of the last input. So backward runs correctly before any forward pass, and output_shape() keeps the default answer of Unknown. A rank-0 input gives Error::InvalidInput, and an axis of extent 0 gives Error::EmptyInput.

use ndarray::Array2;
use rustyml::neural_network::layers::*;
use rustyml::neural_network::traits::UnaryLayer;
use rustyml::neural_network::Ctx;

fn main() {
    // 8-bit pixel values, 2 samples of 3 values each.
    let x = Array2::from_shape_vec((2, 3), vec![0.0, 51.0, 102.0, 153.0, 204.0, 255.0])
        .unwrap()
        .into_dyn();
    let mut ctx = Ctx::inference();

    // Into [0, 1]: the layer holds no parameter, so training never moves the map.
    let unit = Rescaling::new(1.0 / 255.0).forward(&x, &mut ctx).unwrap();
    assert!(unit.iter().all(|v| (0.0..=1.0).contains(v)));
    assert!((unit[[0, 1]] - 0.2).abs() < 1e-6);

    // Into [-1, 1]: y = x / 127.5 - 1.
    let signed = Rescaling::new(1.0 / 127.5)
        .with_offset(-1.0)
        .forward(&x, &mut ctx)
        .unwrap();
    assert!((signed[[0, 0]] + 1.0).abs() < 1e-6);
    assert!((signed[[1, 2]] - 1.0).abs() < 1e-6);

    // The backward pass multiplies the incoming gradient by scale, and needs no cache.
    let layer = Rescaling::new(0.5);
    let grad = layer
        .backward(&Array2::from_elem((2, 3), 4.0f32).into_dyn(), &mut ctx)
        .unwrap();
    assert!(grad.iter().all(|v| (*v - 2.0).abs() < 1e-6));

    println!("rescaled shape: {:?}", unit.shape());
}

Put the layer first in the stack, before the first layer that carries weights. A Rescaling anywhere else still runs, but it then scales an activation rather than the raw input. A following weight matrix can absorb the same factor for free.

3.8.11. Seeding, determinism, and what gets saved

Every layer on this page also shares 2 further properties. The first is randomness. Only the dropout and Gaussian layers draw random numbers. Each draws from its own StdRng, seeded through the global-seed machinery of the crate. new seeds from the global seed or from entropy. with_random_state(seed) pins the seed. A pass draws from a copy of the stream in the context, and the advanced copy reaches the layer through apply_state. The normalization layers are fully deterministic. Even the internal parallel and serial thresholds these layers use are bit-for-bit invariant, and the parallel path produces the same result as the serial path. So enabled threads never change your numbers. Threading is a performance knob only, covered in Performance Tuning and Parallelism. Seed everything together through Reproducibility and Random Seeds.

The second property is persistence. It has 1 asymmetry to remember before you save a model. The dropout and noise layers hold no trainable parameters. They serialize as empty, and the crate does not save their RNG state. This is harmless, because these layers are the identity at inference anyway.

UnitNormalization and Rescaling also serialize as empty, and neither holds any state to lose. LayerNormalization, GroupNormalization, and InstanceNormalization serialize gamma and beta alone, and only the ones the layer holds. This is complete, because they recompute statistics from each input.

BatchNormalization is the exception. It writes 4 named arrays: gamma, beta, moving_mean, and moving_variance. Inference uses the 2 moving statistics. If the crate dropped them, a loaded model would normalize with garbage values. Because the crate persists them, a saved and reloaded batch-normalization model predicts identically to the original. Every record carries its kind as well. So the crate refuses a file that offers a trainable array where the layer keeps state, rather than applying it. See Saving and Loading Weights and Model Persistence in Depth for the mechanics of the round trip.