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 4 normalization layers (batch, layer, group, instance). They all live in rustyml::neural_network::layers::regularization for one reason. Each layer behaves differently in training mode than in inference mode, and the crate factors that split into one 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.

Every layer here plugs into a Sequential model with the same .add(...) call, just like Dense Layers and Activations. Every constructor returns a Result. See Error Handling for the reason.

Import everything with use rustyml::neural_network::layers::*;. This re-exports Dropout, SpatialDropout1D/2D/3D, GaussianNoise, GaussianDropout, BatchNormalization, LayerNormalization (with its LayerNormalizationAxis), GroupNormalization, and InstanceNormalization. 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

Every mode-dependent layer carries a private training: bool flag. It exposes 2 ways to run forward. The Layer trait defines both methods:

fn forward(&mut self, input: &Tensor) -> Result<Tensor, Error>;   // records caches, honors the flag
fn predict(&self, input: &Tensor) -> Result<Tensor, Error>;       // always eval, writes no caches

forward takes &mut self. It reads the training flag and stores whatever the backward pass needs: masks, batch statistics, or the noise draw. predict takes &self. It always runs the inference path and records nothing. This is why a compiled model can serve concurrent requests across threads. Flip the flag with set_training_if_mode_dependent(is_training) (the trait method) or the inherent set_training(is_training). Layers that do not depend on the mode, such as Dense, activations, and pooling, inherit a no-op for both methods.

Inside a Sequential model, you rarely touch this flag yourself. fit and fit_with_batches set every layer to training mode and call forward. predict and evaluate call each layer’s predict method instead. Watch for one trap: calling forward by hand for inference. A BatchNormalization layer whose flag is still true computes batch statistics from your test batch and changes its running averages. For anything that is not a training step, call model.predict(...), or set the flag to false. 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. Their forward output is therefore identical whether the flag is on or off, and predict equals forward bit for bit. Only its backward pass depends on the mode: in inference mode it returns the upstream gradient unchanged instead of computing the real gradient. Batch normalization is the only layer here that keeps state (running mean and variance). It is the only one that computes something different between the 2 modes.

This split also explains why one model can report 2 different losses. The per-epoch figures in the History from fit and fit_with_batches come from training-mode forward passes. In training mode, dropout zeros a fraction of the activations and rescales the survivors. Batch norm normalizes with each batch’s own statistics and folds them into its running averages. evaluate runs the inference path instead: dropout is the identity, and batch norm 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 that batch’s own update. 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, input_shape) builds the classic layer. rate is the fraction of units to zero. It must lie in [0.0, 1.0] inclusive, or the call returns Error::InvalidParameter. input_shape is checked against every input, but only for its rank and its per-sample axes. Axis 0 is the batch axis. It varies with whoever calls forward, so the check skips it. A layer declared vec![4, 8] therefore takes a batch of any size, but still rejects rows that are not 8 wide. The wildcard vec![] turns the shape check off entirely.

The implementation is inverted dropout. During a training forward, it 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: inference becomes a pure identity operation (predict returns the input unchanged), 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 rustyml::neural_network::layers::*;
use rustyml::neural_network::traits::Layer;
use ndarray::Array;

fn main() {
    // 50% rate, expecting [4, 8] inputs. Seed the mask so this run is reproducible.
    let mut dropout = Dropout::new(0.5, vec![4, 8]).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.
    dropout.set_training_if_mode_dependent(true);
    let train_out = dropout.forward(&input).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.
    dropout.set_training_if_mode_dependent(false);
    assert_eq!(dropout.forward(&input).unwrap(), input);

    // predict() is the eval path regardless of the flag, and never caches a mask.
    assert_eq!(dropout.predict(&input).unwrap(), input);
}

The mask is drawn from a per-layer StdRng. By default Dropout::new seeds it from the global seed (or from OS entropy if none is set). with_random_state(seed) re-seeds it deterministically. Because the RNG advances with each draw, 2 consecutive training forward calls produce different masks. Seeding fixes the sequence across process runs, not equality between calls. See Reproducibility and Random Seeds for how the global seed threads through every randomized component. Dropout has no trainable parameters. Calling backward before forward (so no mask was cached) returns NnError::ForwardPassNotRun. In inference mode or at rate == 0.0, the backward 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. A dropped pixel’s value 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, which forces the network to avoid depending on any single feature map. This mirrors Keras’ SpatialDropout* layers.

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 one keep/drop value per (batch, channel) pair: 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 rustyml::neural_network::layers::*;
use rustyml::neural_network::traits::Layer;
use ndarray::Array;

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

    let input = Array::ones((1, 4, 4, 8)).into_dyn();
    let out = sd.forward(&input).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 injecting noise instead of zeroing values. GaussianNoise::new(stddev, input_shape) is additive. During training it adds zero-mean N(0, stddev^2) noise: output = input + noise. This is a data-augmentation style of regularization. It perturbs inputs without changing 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, input_shape) 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 caches the exact noise draw, so its backward pass can reuse it. Because y = x * noise, dx = grad * noise. A backward call before any training forward call returns ForwardPassNotRun. Both layers are the identity at inference, and also when stddev or rate is 0. Neither layer has trainable parameters.

use rustyml::neural_network::layers::*;
use rustyml::neural_network::traits::Layer;
use ndarray::Array;

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, vec![2, 4]).unwrap().with_random_state(1);
    noise.set_training_if_mode_dependent(true);
    let noisy = noise.forward(&input).unwrap();

    // Multiplicative: output = input * N(1, sqrt(rate/(1-rate))). E[output] stays at input.
    let mut gdrop = GaussianDropout::new(0.3, vec![2, 4]).unwrap().with_random_state(1);
    gdrop.set_training_if_mode_dependent(true);
    let scaled = gdrop.forward(&input).unwrap();

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

    // Both are the identity at serving time.
    noise.set_training_if_mode_dependent(false);
    gdrop.set_training_if_mode_dependent(false);
    assert_eq!(noise.forward(&input).unwrap(), input);
    assert_eq!(gdrop.forward(&input).unwrap(), input);
}

3.8.5. The normalization layers at a glance

The 4 normalization layers share one skeleton. Each layer subtracts a mean and divides by a standard deviation computed over some set of axes. It then applies a learnable per-channel affine transform: 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. Normalizing scale and shift should not get pulled toward zero. Every layer also takes an epsilon value (typically 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 the 4 layers is only which axes the statistics reduce over. 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 choosing a layer. Batch norm couples samples together through shared statistics. This makes it effective, but also fragile at small batch sizes. The other 3 layers normalize each sample independently, so batch size does not affect them. All 4 layers preserve the input shape. The shape guards below are not optional. Group and instance normalization require rank 3 or higher (a 2-D input returns Error::InvalidInput). Batch norm accepts rank 2 and higher.

3.8.6. BatchNormalization

BatchNormalization::new(input_shape, momentum, epsilon) takes the shape first, then the 2 scalars. input_shape must be non-empty, or the call returns Error::EmptyInput. momentum must be in [0.0, 1.0]. epsilon must be positive. Dimension 0 is the batch. The last dimension is the channel or feature axis. The per-channel parameters have length input_shape.last().

For a 2-D [N, C] input, this is ordinary per-feature batch norm. For a rank-3-or-higher [N, ...spatial, C] input, the statistics reduce over the batch and every spatial position. This gives 1 mean, variance, gamma, and beta per channel: “spatial” batch norm, matching Keras’ axis=-1 default. 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. Collapsing the leading axes reinterprets the same bytes instead of reshaping them.

A 1-D input_shape, such as vec![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 norm mode-dependent is a pair of running statistics. During a training forward call, the layer first normalizes with the current batch’s mean and variance. It then updates the running statistics as running = running * momentum + batch * (1 - momentum). This is the Keras convention. A high momentum, such as 0.99, weights the history heavily and moves the running estimate slowly. PyTorch uses the opposite convention: its momentum weights the new batch. A value copied directly from PyTorch code produces the opposite behavior here.

momentum == 0.0 discards history entirely, so the running statistics equal the last batch’s. momentum == 1.0 freezes the running statistics at their initial values (mean 0, variance 1). This quietly turns eval-mode normalization into a near-identity operation, usually by mistake. At inference, forward and predict normalize with the running statistics instead of the batch. This is exactly why feeding a test batch through forward in training mode corrupts the model.

Placement relative to the activation matters. The fused-activation design of Dense shapes your options. The original batch-norm paper places it before the nonlinearity. Dense::new(.., .., Activation::ReLU) folds the activation into the linear layer. So “BN 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)
model
    .add(Dense::new(4, 8, Activation::Linear).unwrap())
    .add(BatchNormalization::new(vec![6, 8], 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, using the simpler activation-in-the-Dense form:

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

fn main() {
    // The 6 in the BatchNormalization shape [6, 8] is not enforced. Only the trailing 8 is.
    let x = Array::ones((6, 4)).into_dyn();
    let y = Array::ones((6, 1)).into_dyn();

    let mut model = Sequential::new();
    model
        .add(Dense::new(4, 8, Activation::ReLU).unwrap())
        .add(BatchNormalization::new(vec![6, 8], 0.9, 1e-5).unwrap())
        .add(Dropout::new(0.3, vec![6, 8]).unwrap())
        .add(Dense::new(8, 1, Activation::Linear).unwrap())
        .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 eval path (running stats, dropout as identity).
    model.fit(&x, &y, 3).unwrap();
    println!("prediction shape: {:?}", model.predict(&x).unwrap().shape());
}

The leading dimension of input_shape records a batch size, but the layer does not check it. A BatchNormalization layer declared for [6, 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 epoch’s reported loss either. fit_with_batches weights each batch by the number of samples it holds, instead of giving every batch one vote. So the History entry stays the dataset-wide mean per-sample loss, however the split lands. This matches the accounting Keras uses.

The check does enforce the rank and every per-sample axis. So an [n, 16] input, or an input of the wrong rank, still returns Error::ShapeMismatch. This also means the vec![] wildcard is not the way to make a layer survive mini-batching. Use it only when the per-sample shape itself varies.

You can also inject known weights directly with set_weights(gamma, beta, running_mean, running_var). A shape mismatch returns NnError::WeightShape. 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(input_shape, epsilon) normalizes across features within each sample. 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. Batch norm’s statistics get noisy or meaningless with a handful of samples. Layer norm is unaffected, because every sample stands alone. Its forward output is the same in training and inference, so predict equals forward. 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 is Keras-style: the statistics span the combined elements of those axes, and gamma/beta become 1-D over their product. Changing the axis resizes gamma and beta to match, so call this method 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 rustyml::neural_network::layers::*;
use rustyml::neural_network::traits::Layer;
use ndarray::Array;

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(vec![2, 4], 1e-5).unwrap();
    let out = ln.forward(&input).unwrap();

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

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

    // Statistics come from the current input, so predict() reproduces forward() exactly.
    assert_eq!(ln.predict(&input).unwrap(), out);
}

3.8.8. GroupNormalization and InstanceNormalization

These 2 layers fill the middle ground for convolutional models where the batch is too small for batch norm to work well. Examples include detection and segmentation backbones, and generative models. GroupNormalization::new(input_shape, num_groups, epsilon) splits the channels into num_groups contiguous groups and normalizes within each group, per sample. With 1 group, it becomes layer norm over all channels. With as many groups as channels, it becomes instance norm. InstanceNormalization::new(input_shape, epsilon) normalizes each (sample, channel) plane on its own. This is the standard choice for style transfer, because it strips per-instance contrast while leaving 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 norm, 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 norm is group norm 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 rustyml::neural_network::layers::*;
use rustyml::neural_network::traits::Layer;
use ndarray::Array;

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(vec![1, 4, 4], 1e-5).unwrap();
    let out_in = inn.forward(&input).unwrap();

    // ...which is exactly GroupNorm with 1 group per channel.
    let mut gn = GroupNormalization::new(vec![1, 4, 4], 4, 1e-5).unwrap();
    let out_gn = gn.forward(&input).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);
}

One divisibility rule matters for GroupNormalization: num_groups must evenly divide the channel count. The crate checks this at forward time, not in the constructor. So a mismatched layer builds without error, but fails on its first forward call with Error::InvalidParameter. GroupNormalization::new also catches num_groups == 0 at construction. Both constructors catch an empty input_shape (Error::EmptyInput) and a non-positive or non-finite epsilon (Error::InvalidParameter) up front. Both layers accept set_weights(gamma, beta).

3.8.9. 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 crate’s global-seed machinery. new seeds from the global seed or from entropy. with_random_state(seed) pins the seed. The normalization layers are fully deterministic. Even the internal parallel and serial thresholds these layers use are bit-for-bit invariant: the parallel path produces the same result as the serial path. So enabling threads never changes 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 one 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. LayerNormalization, GroupNormalization, and InstanceNormalization serialize only their gamma and beta values. This is complete, because they recompute statistics from each input. BatchNormalization is the exception. Its weight record carries gamma, beta, and running_mean and running_var. Inference uses those running statistics. If the crate dropped them, a loaded model would normalize with garbage values. Because the crate persists them, a saved and reloaded batch-norm model predicts identically to the original. See Saving and Loading Weights and Model Persistence in Depth for the mechanics of the round trip.