# 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`](./3.1._The_Sequential_Model.md) with the same `.add(...)` call, just like [Dense Layers and Activations](./3.2._Dense_Layers_and_Activations.md). No constructor here takes an input shape, because the model build supplies it. Every constructor except `Rescaling::new` returns a `Result`. See [Error Handling](../Chapter-01/1.6._Error_Handling.md) 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`](./3.1._The_Sequential_Model.md) trait gives every layer with 1 input these 3 methods:
```rust,ignore
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:
| Layer | `forward` with `Ctx::training()` | `forward` with `Ctx::inference()` | `backward` with `Ctx::inference()` |
| --- | --- | --- | --- |
| `Dropout`, `SpatialDropout*` | drop a fraction, rescale survivors | identity (pass through unchanged) | gradient passed through |
| `GaussianNoise` | add `N(0, stddev^2)` | identity | gradient passed through |
| `GaussianDropout` | multiply by `N(1, sigma)` | identity | gradient passed through |
| `BatchNormalization` | normalize with **batch** stats, propose new running stats | normalize with **running** stats | gradient passed through |
| `LayerNormalization`, `GroupNormalization`, `InstanceNormalization` | normalize with stats from the current input | **same** stats from the current input | gradient passed through |
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.
```rust
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. `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.
```rust
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](../Chapter-07/7.1._Reproducibility_and_Random_Seeds.md) 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")`.
```rust
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. 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.
```rust
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:
| Layer | Mean/variance reduced over | 1 statistic per | Depends on batch? | `gamma`/`beta` length |
| --- | --- | --- | --- | --- |
| `BatchNormalization` | batch `N` (and all spatial) | channel | **yes** | `C` |
| `LayerNormalization` | the normalized axis (last, by default) | everything but that axis | no | normalized axis size |
| `GroupNormalization` | a group of channels + spatial, per sample | (sample, group) | no | `C` |
| `InstanceNormalization` | spatial only, per sample and channel | (sample, channel) | no | `C` |
| `UnitNormalization` | no mean, an L2 norm over the chosen axes | everything but the chosen axes | only with axis 0 | none |
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)`. `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:
```rust
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:
```rust,ignore
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:
```rust
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.9._Saving_and_Loading_Weights.md).
## 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` picks the last axis, `Custom(axis)` picks a single other axis, and `Multiple(vec![...])` normalizes 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.
```rust
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:
```rust
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.
```rust
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](../Chapter-04/4.2._Standardization_and_Normalization.md), 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`.
```rust
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](../Chapter-07/7.3._Performance_Tuning_and_Parallelism.md). Seed everything together through [Reproducibility and Random Seeds](../Chapter-07/7.1._Reproducibility_and_Random_Seeds.md).
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](./3.9._Saving_and_Loading_Weights.md) and [Model Persistence in Depth](../Chapter-07/7.2._Model_Persistence_in_Depth.md) for the mechanics of the round trip.