# 3.2. Dense Layers and Activations
`Dense` is the main layer of a feedforward network. It applies a linear map, `input * W + b`, then an elementwise nonlinearity. RustyML fuses that nonlinearity into the layer itself. It does not treat the nonlinearity as a separate stage.
This page covers the constructor and its fused-activation design. It also covers the exact initialization scheme and the 12 activations, with their gradients and failure modes. It covers the GEMM-backed compute path and how to read or inject weights. It assumes you have already read [the Sequential model page](./3.1._The_Sequential_Model.md). The layers here are what you collect in a `SequentialBuilder`, and the input width of each one comes from the model build.
## 3.2.1. The fused-activation design
The constructor takes the activation as its second argument, not as a separate layer:
```rust,ignore
pub fn new(units: usize, activation: impl Into<Activation>) -> Result<Dense, Error>
```
Internally the layer stores an `Activation` *value*. This is a plain `Copy` enum with 12 variants. It is not a generic type parameter `Dense<A>`, and it is not a `Box<dyn Activation>`. That choice is deliberate. A generic parameter would monomorphize `Dense` 12 ways.
It would also force weight deserialization to probe every `Dense<A>` pairing to find the concrete type on disk. A trait object would add an indirection to every elementwise call. The runtime enum keeps `Dense` a single concrete type. The persistence layer can then downcast every saved layer to exactly 1 struct (see [Saving and Loading Weights](./3.9._Saving_and_Loading_Weights.md)). The activation math stays a pure, stateless function, and the layer calls it inside its own forward and backward passes.
Keras makes the same fusion choice, with `Dense(units, activation="relu")`. RustyML has no implicit "no activation". You must always pass one. To get a pure linear layer, pass `Activation::Linear`, the identity. This is the equivalent of Keras' `activation=None`. Every example below that ends in a regression head does exactly this.
Fusing is the right default. RustyML also provides every activation as a standalone layer (`Linear`, `ReLU`, `LeakyReLU`, `ELU`, `SELU`, `Sigmoid`, `HardSigmoid`, `Tanh`, `Softplus`, `Softsign`, `Exponential`, `Softmax`). Use these when you need something between the linear map and the nonlinearity. A normalization layer is the most common case (see [Regularization and Normalization Layers](./3.8._Regularization_and_Normalization_Layers.md)). The 2 forms are equivalent:
```rust,ignore
// These two stacks compute the same thing.
builder.add(Dense::new(32, Activation::ReLU).unwrap()); // fused: one cached output
builder.add(Dense::new(32, Activation::Linear).unwrap()) // split: identity Dense ...
.add(ReLU::new()); // ... then a standalone ReLU
```
Prefer the fused form. It caches a single activated tensor instead of 2. The backward pass then differentiates the activation in terms of that cached output (see 3.2.6). Use the split form only when a layer must sit between `W*x + b` and the nonlinearity.
## 3.2.2. Constructing and sizing a layer
`Dense::new(units, activation)` returns `Result<Dense, Error>`. `units` is the number of neurons, and therefore the output width. It must be non-zero, and a `0` yields `Error::InvalidParameter` (see [Error Handling](../Chapter-01/1.6._Error_Handling.md)).
**The constructor no longer takes the input width.** A kernel extent comes from the input, and the input is not there when you write the constructor. `Layer::build` is where the layer learns it. `build` reads the last axis of the shape it receives, records it, and draws the kernel at that width. `SequentialBuilder::build` calls it once per layer, so a stack of `Dense` layers never repeats a width. 2 adjacent layers can no longer disagree about one.
The parameter count is `input_width * units + units`: 1 weight per (input, output) pair, plus 1 bias per output. An unbuilt layer holds no array, so it reports 0 for both counts and `output_shape()` reports `"Unknown"`. `param_count()` reports it as `ParamCounts::trainable(n)`, which is `n` trainable elements and 0 non-trainable ones.
The layer gives its 2 tensors the Keras names `kernel` and `bias`. The optimizer keys its state on those names, and a checkpoint addresses them by the same names (see [3.4](./3.4._Optimizers.md) and [3.9](./3.9._Saving_and_Loading_Weights.md)). `output_shape()` renders `(None, units)` once the layer is built for a free batch axis, mirroring Keras' summary. A `Dense(3, ...)` built for a last axis of 4 therefore holds a `4 x 3` weight matrix and a `1 x 3` bias. That is `12 + 3 = 15` trainable scalars.
`compute_output_shape` answers the same question without any of that state. It is pure: it replaces the last axis of the shape you hand it with `units`, and it reads no cache. It works on a layer that was never built and never run. A layer that **is** built also checks the last axis against the width it holds, and refuses a shape that disagrees.
`with_use_bias(false)` builds the layer without the bias, exactly as the Keras `use_bias` argument does. The forward pass is then `activation(input * W)`, the parameter count drops to `input_width * units`, and the layer holds 1 array in place of 2. Put the flag on a layer that a normalization layer follows. The shift of that layer makes the bias of the dense layer redundant.
Input is a tensor of rank 2 or more, whose last axis holds the built width. The layer contracts that last axis alone and replaces it with `units`. Every axis in front of it passes through untouched. So a `(batch, features)` input gives `(batch, units)`, and a `(batch, timesteps, features)` input gives `(batch, timesteps, units)`. A rank-1 input, and a last axis of the wrong width, are both rejected with `Error::InvalidInput` rather than silently reshaped. The following program constructs a layer, builds it, inspects it, and injects known weights:
```rust
use ndarray::Array2;
use rustyml::neural_network::Shape;
use rustyml::neural_network::layers::{Activation, Dense, ParamCounts};
use rustyml::neural_network::traits::{Layer, LayerBase, UnaryLayer};
fn main() {
// The constructor takes the unit count alone, and it draws nothing.
let mut dense = Dense::new(3, Activation::ReLU).unwrap();
assert_eq!(dense.param_count(), ParamCounts::trainable(0));
println!("before build: {}", dense.output_shape()); // Unknown
// `build` supplies the input shape. The layer reads 4 off the last axis and allocates.
dense.build(&Shape::with_free_batch(&[8, 4])).unwrap();
assert_eq!(dense.param_count(), ParamCounts::trainable(15)); // 4*3 + 3
println!("after build: {}", dense.output_shape()); // (None, 3)
// Read the drawn arrays by name, with no clone.
for entry in dense.weights() {
println!("{} {:?} {:?}", entry.name, entry.kind, entry.value.shape());
}
assert_eq!(dense.weight("kernel").unwrap().shape(), [4, 3]);
assert_eq!(dense.weight("bias").unwrap().shape(), [1, 3]);
// Inject known weights. The shapes are validated against what the build settled.
let weights =
Array2::from_shape_vec((4, 3), (0..12).map(|v| v as f32).collect::<Vec<f32>>()).unwrap();
let bias = Array2::zeros((1, 3));
dense.set_weights(weights, bias).unwrap();
// `with_use_bias(false)` drops the bias array, so the layer holds the kernel alone.
let mut no_bias = Dense::new(3, Activation::ReLU).unwrap().with_use_bias(false);
no_bias.build(&Shape::with_free_batch(&[8, 4])).unwrap();
assert_eq!(no_bias.param_count(), ParamCounts::trainable(12));
assert!(no_bias.weight("bias").is_none());
println!("no-bias paths: {}", no_bias.weights().len());
// The shape answer is pure, so an unbuilt twin answers for any rank.
let twin = Dense::new(3, Activation::ReLU).unwrap();
let sequence = Shape::new(vec![None, Some(5), Some(4)]);
println!("rank 3: {}", twin.compute_output_shape(&sequence).unwrap()); // (None, 5, 3)
}
```
`set_weights` needs a built layer, and it returns `Error::NeuralNetwork(NnError::NotBuilt)` on a layer that holds no array yet. There is nothing to check a shape against before the build. It then checks both shapes and returns `Error::NeuralNetwork(NnError::WeightShape)` on a mismatch. `set_weights` rejects a `(3, 3)` weight for a `4 -> 3` layer, or a `(1, 4)` bias, rather than truncating it.
Its second argument is optional. Pass `None` for a layer built with `with_use_bias(false)`, and pass the array for every other layer. The other order returns `Error::InvalidParameter`, because a layer that holds no bias would silently drop one that reaches it.
The rank costs nothing. 1 kernel serves every leading position. The layer folds all the leading axes into 1 row axis, and it runs the same single matrix product a rank-2 input runs. A rank-3 pass and the folded rank-2 pass therefore give the same values, bit for bit.
`forward_mut` is the entry point of a caller that drives 1 layer by hand. It builds the layer from the tensor it receives, and then it completes the pass. This is why the 2 layers below need no `build` call. The layer parks the shape of the input it received, so the gradient goes back at the rank it arrived at. The build shape holds the last axis alone, so `output_shape()` reports `(None, units)` at every rank. This is what lets a `Dense` layer sit after a recurrent layer that returns sequences, and transform every timestep with the same weights:
```rust
use ndarray::Array;
use rustyml::neural_network::Ctx;
use rustyml::neural_network::layers::*;
use rustyml::neural_network::traits::{ParamId, UnaryLayer};
fn main() {
// 2 samples, 5 timesteps, 4 features per step.
let x = Array::from_shape_fn((2, 5, 4), |(n, t, f)| (n * 20 + t * 4 + f) as f32 * 0.01)
.into_dyn();
// The kernel contracts the last axis, and every leading axis passes through. A
// training context holds what the backward pass needs.
let mut dense = Dense::new(3, Activation::Linear).unwrap().with_random_state(7);
let mut ctx = Ctx::training();
let out = dense.forward_mut(&x, &mut ctx).unwrap();
assert_eq!(out.shape(), &[2, 5, 3]);
// 1 kernel serves every leading position, so a rank-3 pass equals a folded rank-2
// pass. An inference context parks nothing at all, so it stays empty.
let folded = x.to_shape((10, 4)).unwrap().to_owned().into_dyn();
let mut twin = Dense::new(3, Activation::Linear).unwrap().with_random_state(7);
let mut eval = Ctx::inference();
let flat = twin.forward_mut(&folded, &mut eval).unwrap();
assert_eq!(flat.shape(), &[10, 3]);
assert_eq!(eval.pending_caches(), 0);
for row in 0..10 {
for u in 0..3 {
assert_eq!(out[[row / 5, row % 5, u]], flat[[row, u]]);
}
}
// The gradient goes back at the rank of the input, and the 2 parameter gradients
// go into the store of the context.
let grad = dense.backward(&out, &mut ctx).unwrap();
assert_eq!(grad.shape(), &[2, 5, 4]);
let kernel_grad = ctx.grads().get(ParamId::new(0, "kernel")).unwrap();
assert_eq!(kernel_grad.shape(), &[4, 3]);
println!("output shape: {:?}", out.shape());
}
```
Each layer takes its own context above. A context serves 1 pass of 1 stack, and the caches of 1 position form a stack that the matching backward pass empties. 2 layers that write into 1 context under 1 position would therefore cross their caches. A model gives every layer its own position, and a caller that drives 1 layer by hand takes position 0.
A convolutional stack still needs a `Flatten` in front of its classifier head. `Dense` would otherwise contract the channel axis alone and leave every spatial position separate, which is a `1x1` convolution and not a classifier. Insert `Flatten` when you want 1 vector per sample.
## 3.2.3. Weight initialization
Weights use **Xavier/Glorot uniform** initialization. Each element is drawn from `Uniform(-limit, +limit)`, where `limit = sqrt(6 / (fan_in + fan_out))`. For a `Dense` layer the 2 fans are the built input width and `units`. Biases start at exactly zero. It matches Keras' `Dense` default exactly (`glorot_uniform` kernel, `zeros` bias).
The rule is a value, not a hidden constant. `Initializer` is a small closed enum with 3 variants. `GlorotUniform` covers almost every kernel of the module, `Uniform { limit }` covers the embedding table, and `Orthogonal` covers the recurrent kernels. A layer holds the value it draws with, so the rule compares, copies, and serializes with the layer.
**The layer supplies the 2 fans, and nothing derives them from the shape of the array.** They arrive as a `Fans` pair that the layer builds from its own configuration. The transposed convolution is the reason. A plain convolution holds its kernel as `(spatial axes, channels, filters)`, and a transposed convolution holds the same kernel as `(spatial axes, filters, channels)`.
The last 2 axes carry opposite roles in the 2 layouts. A rule that reads them gives the transposed layer the 2 fans the wrong way round. Glorot hides that swap, because its range reads the sum of the pair and a sum is symmetric. A layer that names its 2 counts cannot make the mistake at all. The defect never waits for an initializer that reads 1 fan alone.
RustyML applies Glorot regardless of the activation. It does **not** switch to He/Kaiming initialization for `ReLU` layers, even though He is the textbook match for rectifiers. For shallow nets this rarely matters. For a deep ReLU stack, early convergence may be slightly slower than a He-initialized equivalent. If that happens, initialize with `set_weights` from your own scaled draw.
Initialization draws through the crate's shared RNG, in `build`. By default the seed comes from the process-global seed if one is set, otherwise from entropy. 2 runs then give different weights unless you fix the randomness. There are 2 ways to make it reproducible. Set a thread-local global seed with `rustyml::random::set_global_seed(...)` before you build the model. This also fixes dropout masks and the fit-time batch shuffle.
You can instead seed 1 layer explicitly, with `Dense::new(...)?.with_random_state(seed)`. On an unbuilt layer the call only **records** the seed, because there is no array to redraw yet, and `build` spends it. On a layer that is already built the call redraws at once, so the order of the 2 calls never matters. All 15 builders that draw a weight work this way.
Either way the seed leaves the global stream untouched. The full seeding model is in [Reproducibility and Random Seeds](../Chapter-07/7.1._Reproducibility_and_Random_Seeds.md). It also explains why an explicit per-layer seed does not change the seeds handed to unseeded layers.
## 3.2.4. The activations
`Activation` has 12 variants. Each variant's backward pass is expressed in terms of the *activated output* `a = f(x)`, not the pre-activation `x`. The layer caches that output, not the input. That contract admits every activation whose derivative has a closed form in `a`, which covers all 12 below.
The same contract is exactly why GELU, SiLU (Swish), and Mish are absent. Their `a = x * g(x)` shape has no closed-form inverse, so no derivative in `a` alone exists. Adding them needs a wider contract that also hands the backward pass the pre-activation.
| `Activation` | Forward `f(x)` | Backward (given upstream `g`, output `a`) | Output range |
| --- | --- | --- | --- |
| `ReLU` | `max(0, x)` | pass `g` where `a > 0`, else `0` | `[0, inf)` |
| `Sigmoid` | `1 / (1 + e^-x)` | `g * a * (1 - a)` | `(0, 1)` |
| `Tanh` | `tanh(x)` | `g * (1 - a^2)` | `(-1, 1)` |
| `Softmax { axis }` | shifted `exp`, normalized along `axis` (default `-1`, the last axis) | `a_i * (g_i - sum_j(a_j * g_j))` (lane Jacobian) | simplex, each lane sums to `1` |
| `Linear` | `x` | pass `g` through | `(-inf, inf)` |
| `LeakyReLU { negative_slope }` | `x` for `x >= 0`, else `negative_slope * x` | pass `g` where `a >= 0`, else `g * negative_slope` | `(-inf, inf)` |
| `ELU { alpha }` | `x` for `x > 0`, else `alpha * (e^x - 1)` | pass `g` where `a > 0`, else `g * (a + alpha)` | `(-alpha, inf)` |
| `SELU` | `scale * x` for `x > 0`, else `scale * alpha * (e^x - 1)` | `g * scale` where `a > 0`, else `g * (a + scale * alpha)` | `(-scale * alpha, inf)` |
| `Softplus` | `ln(1 + e^x)` | `g * (1 - e^-a)` | `(0, inf)` |
| `Softsign` | `x / (1 + abs(x))` | `g * (1 - abs(a))^2` | `(-1, 1)` |
| `HardSigmoid` | `clip(x/6 + 0.5, 0, 1)` | `g / 6` where `0 < a < 1`, else `0` | `[0, 1]` |
| `Exponential` | `e^x` | `g * a` | `(0, inf)` |
**ReLU** is the default hidden-layer choice. It is cheap, and it does not saturate on the positive side. Its failure mode is the *dead neuron*. The derivative is `0` for `x <= 0`, so a neuron whose pre-activation is negative for every example in the batch gets zero gradient.
It never updates, and it stays off for good. A high learning rate makes this worse. It pushes neurons into the dead region early. Glorot init has no leaky variant. Your options are a smaller learning rate and well-scaled inputs, or, if you inject your own weights, a better initial scale.
**LeakyReLU, ELU, and SELU** are the direct answer to that failure mode. Each one keeps a non-zero gradient below 0. A unit whose pre-activation stays negative for the whole batch still receives gradient, so it can recover. `LeakyReLU { negative_slope }` scales the negative side by a constant, and its output stays unbounded. `PReLU` in 3.2.5 goes further and learns that constant. `ELU { alpha }` instead saturates the negative side at `-alpha`, which pulls the mean activation toward 0, at a cost of 1 exponential per negative element.
`SELU` is the same shape as `ELU`, with `alpha` and `scale` fixed at `1.6732632` and `1.0507010`. Its self-normalizing property holds only under Lecun-normal initialization, and `Dense` hard-codes Glorot uniform (see 3.2.3). RustyML has no Lecun-normal initializer yet. `SELU` still trains as a plain activation under Glorot. The variance-preserving guarantee of Klambauer et al. (2017) does not hold there. To get it, inject your own Lecun-normal draw with `set_weights`.
`negative_slope` and `alpha` must both be finite and greater than 0. `Activation::validate` enforces that bound. All 14 trainable layer constructors call it: `Dense`, `Conv1D`, `Conv2D`, `Conv3D`, `Conv1DTranspose`, `Conv2DTranspose`, `Conv3DTranspose`, `DepthwiseConv1D`, `DepthwiseConv2D`, `SeparableConv1D`, `SeparableConv2D`, `SimpleRNN`, `LSTM`, and `GRU`. An unusable value therefore returns `Error::InvalidParameter` where you build the model, not on the first forward pass.
The bound is strict, and the output-only contract is the reason. The backward pass reads the branch off the sign of `a`. A value of 0 collapses the whole negative side onto `a = 0`, which erases the branch. A negative value inverts the sign, so the backward pass reads the wrong branch. Use `Activation::ReLU` when you want a slope of 0.
`LeakyReLU` uses `x >= 0` for the positive branch, so its derivative at exactly 0 is 1. `ELU` and `SELU` use `x > 0`, which puts 0 itself on the negative branch. Their derivatives at exactly 0 are therefore `alpha` and `scale * alpha`, not 1. `LeakyReLU::default()` uses a slope of `0.3`, and `ELU::default()` uses an `alpha` of `1.0`.
The following hidden stack uses `ELU` twice. The first layer fuses it into `Dense`. The second pairs an identity `Dense` with the standalone `ELU` layer, which is the split form of 3.2.1:
```rust
use ndarray::Array;
use rustyml::neural_network::Shape;
use rustyml::neural_network::layers::{Activation, Dense, ELU};
use rustyml::neural_network::losses::MeanSquaredError;
use rustyml::neural_network::optimizers::Adam;
use rustyml::neural_network::sequential::SequentialBuilder;
fn main() {
// 4 samples, 3 features, 1 continuous target each. Negative features exercise the
// saturating branch of ELU.
let x = Array::from_shape_vec(
(4, 3),
vec![-1.0, 0.1, 0.2, 1.0, -0.9, 0.8, 0.2, 0.1, -2.0, 0.9, 1.0, 0.8],
)
.unwrap()
.into_dyn();
let y = Array::from_shape_vec((4, 1), vec![0.3, 2.7, 0.3, 2.7])
.unwrap()
.into_dyn();
let mut model = SequentialBuilder::new()
// Fused: the activation runs inside the Dense layer.
.add(Dense::new(8, Activation::ELU { alpha: 1.0 }).unwrap())
// Split: an identity Dense, then the same activation as its own layer.
.add(Dense::new(8, Activation::Linear).unwrap())
.add(ELU::new(1.0).unwrap())
.add(Dense::new(1, Activation::Linear).unwrap()) // regression head
.build(&Shape::known(x.shape()))
.unwrap();
model.compile(
Adam::new(0.01, 0.9, 0.999, 1e-8, 0.0).unwrap(),
MeanSquaredError::new(),
);
model.fit(&x, &y, 20).unwrap();
let preds = model.predict(&x).unwrap();
assert_eq!(preds.shape(), &[4, 1]);
assert!(preds.iter().all(|v| v.is_finite()));
// A slope of 0 erases the negative branch, so the constructor rejects it.
let dead_slope = Activation::LeakyReLU { negative_slope: 0.0 };
assert!(Dense::new(8, dead_slope).is_err());
}
```
**Sigmoid** squashes values to `(0, 1)`. Its derivative `a * (1 - a)` peaks at `0.25` when `a = 0.5`, and it decays toward zero as the output saturates. Stacking sigmoids in a deep hidden path throttles gradients: the classic vanishing-gradient problem. Use `Sigmoid` for a single binary output, paired with binary cross-entropy, or as a gate. Do not use it as a deep hidden nonlinearity. At extreme inputs, `f32` saturates it to exactly `0.0` or `1.0`.
**Tanh** maps to `(-1, 1)`. Unlike sigmoid, it is zero-centered, which tends to make hidden-layer optimization better behaved. Its gradient `1 - a^2` still reaches `1` near the origin. It saturates the same way sigmoid does at the tails. It is the natural bounded, zero-centered hidden activation, and the recurrent layers use it internally.
**Softplus and Softsign** are the smooth option and the bounded option. `Softplus` is `ln(1 + e^x)`, a smooth approximation of ReLU. Its output is strictly positive, and its derivative `1 - e^-a` never reaches 0, so it has no dead unit at all. It costs 1 exponential and 1 logarithm per element, which makes it more expensive than `LeakyReLU`.
`Softsign` is `x / (1 + abs(x))`. It is bounded to `(-1, 1)` and zero-centered, exactly like tanh. It saturates polynomially rather than exponentially, so its tails keep more gradient than the tails of `Tanh`. It also needs no exponential. Treat it as the cheaper, slower-saturating substitute for `Tanh`.
**HardSigmoid** approximates sigmoid with a clipped straight line, and it uses no exponential. It is `clip(x/6 + 0.5, 0, 1)`, so it reaches exactly 0 at `x = -3` and exactly 1 at `x = 3`. True sigmoid only approaches both ends and never arrives. The derivative is the constant `1/6` on the linear segment, and exactly 0 outside it. That flat outer gradient is the same dead-unit risk ReLU carries. Use `HardSigmoid` as a gate where the cost of `exp` matters.
**Exponential** is `e^x`. Its output is strictly positive, and its derivative is the output itself. Use it on a head that must emit a positive quantity, such as a rate, a variance, or a scale. It has no upper bound, so a pre-activation above about 88 overflows `f32` to `inf`. Keep the head's input scaled, or emit a log-scale value through `Linear` instead.
**Softmax** turns a lane of logits into a probability distribution. `Softmax { axis }` names the axis that the lanes run along, and the default is `-1`, the last axis. The forward pass subtracts each lane's maximum before it takes the exponent. This makes the largest term `exp(0) = 1`, and the sum always `>= 1`. That shift makes softmax overflow-proof and shift-invariant: adding a constant to every logit leaves the output unchanged.
Softmax accepts an input of any rank. A rank-1 input is legal, and it normalizes its single axis, which matches Keras. The backward pass is the true Jacobian-vector product across the lane, not an elementwise multiply, and each gradient lane sums to zero. An axis that resolves outside the rank of the input returns `Error::InvalidInput`.
A negative axis counts back from the end. The value resolves against the rank of the input on each call, never when you build it. The same `axis` therefore means a different axis for inputs of different rank. An `axis` of `-1` is axis 1 of a rank-2 input, and axis 3 of a rank-4 input. The standalone `Softmax` layer takes its axis from the `with_axis` builder, and `Softmax::new()` starts at the default `-1`.
An embedded softmax, the kind you pass to `Dense` or to a convolution layer, accepts the default axis alone. Building the layer refuses any other axis, with `Error::InvalidParameter`, not on the first forward pass. `Dense` folds every leading axis into 1 row axis, and that fold keeps each last-axis lane whole while it destroys every other axis. A convolution layer instead applies the activation to its full output tensor, with no fold. Every host family still restricts an embedded softmax to axis `-1`, because that axis alone names the same lane in each family. Use the standalone `Softmax` layer for any other axis:
```rust
use ndarray::Array;
use rustyml::neural_network::Ctx;
use rustyml::neural_network::layers::{Activation, Dense, Softmax};
use rustyml::neural_network::traits::UnaryLayer;
fn main() {
// 2 samples, 4 positions, 3 channels.
let x = Array::from_shape_fn((2, 4, 3), |(n, t, c)| (n * 12 + t * 3 + c) as f32 * 0.1)
.into_dyn();
// An inference context parks nothing, because no backward pass follows here.
let mut ctx = Ctx::inference();
// The default axis is -1, the last axis, so each of the 8 positions becomes a
// distribution over the 3 channels.
let mut over_channels = Softmax::new();
let per_position = over_channels.forward_mut(&x, &mut ctx).unwrap();
let first_position: f32 = (0..3).map(|c| per_position[[0, 0, c]]).sum();
assert!((first_position - 1.0).abs() < 1e-6);
// with_axis picks another axis. Axis 1 makes each of the 6 channel columns a
// distribution over the 4 positions.
let mut over_positions = Softmax::new().with_axis(1);
let per_channel = over_positions.forward_mut(&x, &mut ctx).unwrap();
let first_channel: f32 = (0..4).map(|t| per_channel[[0, t, 0]]).sum();
assert!((first_channel - 1.0).abs() < 1e-6);
// A rank-1 input normalizes its single axis.
let logits = Array::from_vec(vec![1.0_f32, 2.0, 3.0]).into_dyn();
let probabilities = Softmax::new().forward_mut(&logits, &mut ctx).unwrap();
assert_eq!(probabilities.shape(), &[3]);
// An embedded softmax accepts the default axis alone.
assert!(Dense::new(2, Activation::Softmax { axis: -1 }).is_ok());
assert!(Dense::new(2, Activation::Softmax { axis: 1 }).is_err());
println!("lane sums: {first_position} and {first_channel}");
}
```
Softmax belongs on the output layer of a classifier, paired with the correct loss configuration. Mid-network, softmax is *mechanically* valid: the Jacobian backward is correct, so a `Dense(..., Softmax)` buried in the stack still trains. It is almost never what you want, though. It collapses the representation onto a simplex, discards all magnitude information, and saturates worse than ReLU or tanh.
**Linear** is the identity. Its gradient is `1`, and its output is unbounded. Use it for regression heads, and whenever the loss function expects raw logits.
The loss function must match the activation on the output head. There are 2 correct pairings for multi-class classification. Mixing them silently corrupts training:
- **Softmax head + `CategoricalCrossEntropy::new(false)`.** The final `Dense` emits probabilities. The loss consumes probabilities. Correct.
- **Linear head + `CategoricalCrossEntropy::new(true)`.** The final `Dense` emits raw logits. The loss applies a numerically stable log-softmax internally, and returns the fused `(softmax(z) - y)` gradient in one step. Also correct, and more numerically stable.
The 2 pairings produce the *same* gradient mathematically. The fused `from_logits = true` path avoids the intermediate `-y/p` division and the separate softmax step. It therefore degrades more gracefully when a predicted probability is tiny. Do not mix the 2 pairings.
A softmax head with `from_logits = true` applies softmax twice. A linear head with `from_logits = false` feeds raw logits to a loss that expects probabilities. Loss details live in [Loss Functions](./3.3._Loss_Functions.md).
```rust,ignore
// More stable alternative to a Softmax head: emit logits, and fuse the softmax into the loss.
let mut model = builder
.add(Dense::new(3, Activation::Linear).unwrap()) // raw logits, NOT probabilities
.build(&Shape::known(x.shape()))?;
model.compile(optimizer, CategoricalCrossEntropy::new(true)); // from_logits = true
```
The activations are pure math, with no `NaN`/`Inf` sanitization. A `NaN` propagates through untouched. `tanh` saturates to `1` or `-1` at large inputs. `ReLU` of a large negative input is `0`. A `NaN` anywhere in a softmax row contaminates the whole row through the normalizer. Non-finite values surface downstream, as a `NaN` loss, not at the activation step.
## 3.2.5. PReLU: a learned negative slope
`LeakyReLU` fixes its negative slope at construction, so you have to guess a good value. `PReLU` learns it instead. The forward transform is the same: `x` for `x >= 0` and `alpha * x` below 0. But `alpha` is a trainable array that the optimizer updates with the rest of the model.
`PReLU` is a layer, not an `Activation` variant, and it cannot become one. An `Activation` value carries no state, and this layer carries a trainable array. So it never fuses into `Dense`. Place it after the layer whose output it activates, which is the split form of 3.2.1.
`PReLU::new(alpha)` takes the starting slope alone. The build supplies the shape, exactly as it does for every other layer that owns an array. The layer then holds 1 slope per position of that shape with the batch axis removed, every one starting at `alpha`. A `[batch, 32, 32, 64]` build shape therefore gives 65,536 slopes, 1 per pixel and channel. That is almost never what you want after a convolution.
`with_shared_axes(axes)` is the fix. Each named axis drops to extent 1 in the slope array, and the slope broadcasts back over that axis. The axes count from the batch axis at 0, and the batch axis is shared already:
| Build shape | `shared_axes` | Slope shape | Slope count |
| --- | --- | --- | --- |
| `[batch, 8]` | none | `[8]` | 8 |
| `[batch, 32, 32, 64]` | none | `[32, 32, 64]` | 65,536 |
| `[batch, 32, 32, 64]` | `[1, 2]` | `[1, 1, 64]` | 64 |
| `[batch, 32, 32, 64]` | `[1, 2, 3]` | `[1, 1, 1]` | 1 |
| `[batch, 20, 16]` | `[1]` | `[1, 16]` | 16 |
`[1, 2]` on a 4-D input is the standard choice after a convolution. It gives 1 slope per channel, which is the channel-wise form of He et al. (2015). `[1, 2, 3]` gives a single learned slope for the whole layer, which is `LeakyReLU` with the constant learned rather than guessed.
The layer also stops checking a shared axis at forward time, because 1 slope covers any extent. The same layer then serves images of several sizes. Every other axis after the batch axis must match the build shape, and a mismatch is `Error::InvalidInput`. The layer never checks the batch axis, so a partial final mini-batch always passes.
**The derivative at exactly 0 is 0.** That is neither branch: not the `1` the positive side gives, and not the `alpha` the negative side gives. It is what makes an `alpha` of `0` reproduce `ReLU` exactly, in the transform and in the gradient. That is why `0` is the natural starting value. `LeakyReLU` differs here, because it takes the positive branch at `x >= 0` and its derivative at `0` is `1`. A `PReLU` with a frozen uniform slope therefore matches `LeakyReLU` everywhere except at exactly `0`.
Weight decay skips the slopes, in every optimizer that takes a `weight_decay` argument. Decay pulls a parameter toward `0`. A slope of `0` turns the layer back into `ReLU`, so decay would erase what the layer learns. The slopes count as a no-decay parameter, the same class as a bias and a normalization `gamma`.
The stack below learns 1 slope per filter on a small convolutional model:
```rust
use ndarray::Array;
use rustyml::neural_network::Shape;
use rustyml::neural_network::layers::{Activation, Conv2D, Dense, Flatten, PReLU};
use rustyml::neural_network::losses::MeanSquaredError;
use rustyml::neural_network::optimizers::Adam;
use rustyml::neural_network::sequential::SequentialBuilder;
fn main() {
// 2 single-channel 6x6 images, 1 continuous target each.
let x = Array::from_shape_vec(
(2, 6, 6, 1),
(0..72).map(|v| 0.05 * v as f32 - 1.8).collect::<Vec<_>>(),
)
.unwrap()
.into_dyn();
let y = Array::from_shape_vec((2, 1), vec![0.3, -0.4])
.unwrap()
.into_dyn();
let mut model = SequentialBuilder::new()
// A 3x3 valid convolution over 6x6 emits [2, 4, 4, 4].
.add(Conv2D::new(4, (3, 3), (1, 1), Activation::Linear).unwrap())
// Sharing the 2 spatial axes gives 4 slopes, 1 per filter, instead of 64.
.add(PReLU::new(0.25).unwrap().with_shared_axes(vec![1, 2]).unwrap())
.add(Flatten::new())
.add(Dense::new(1, Activation::Linear).unwrap())
.build(&Shape::known(x.shape()))
.unwrap();
model.compile(
Adam::new(0.01, 0.9, 0.999, 1e-8, 0.0).unwrap(),
MeanSquaredError::new(),
);
model.fit(&x, &y, 10).unwrap();
assert_eq!(model.predict(&x).unwrap().shape(), &[2, 1]);
// The slope array keeps the extent-1 place of every shared axis.
println!("alpha: {:?}", model.weight("1.alpha").unwrap().shape()); // [1, 1, 4]
}
```
The layer holds 1 array, named `alpha`. `weight("alpha")` reads it, and `set_weights(alpha)` writes it back. `alpha` is an `ArrayD<f32>`, and its rank follows the input rank, less 1 for the batch axis. A shared axis keeps its extent-1 place in that array, so a per-channel layer reports `[1, 1, 64]` and not `[64]`.
| Error | Cause |
| --- | --- |
| `Error::InvalidInput` | the build shape has rank below 2, or holds a `0` |
| `Error::InvalidParameter` | `alpha` is not finite |
| `Error::InvalidParameter` | `shared_axes` holds `0` or a repeat, or names an axis the build shape does not hold |
| `Error::InvalidInput` | a forward input whose rank differs, or whose extent on an axis that is not shared differs, from the build shape |
| `Error::NeuralNetwork(NnError::WeightShape)` | `set_weights` got an array that is not the slope shape |
## 3.2.6. Forward, backward, and the GEMM path
The forward pass is `activation(input * W + b)`, computed as a single call into the [gemmkit](https://docs.rs/gemmkit) backend. The product, the per-column bias, and, for `ReLU`, the activation, all run in 1 pass. The bias and activation apply in the kernel's epilogue while the output tile is still in registers. `ReLU` is the only one of the 12 activations with a fused epilogue. `Linear` needs no separate step, since the biased product is already the output. Every other activation runs as a separate vectorized `Activation::forward` pass over the biased product.
The fused result matches the unfused product plus scalar activation bit for bit, with 1 exception. The fused `ReLU` epilogue maps a `NaN` pre-activation to `0.0`. The standalone `ReLU` activation used elsewhere in the crate instead propagates `NaN`. A `NaN` pre-activation already means a diverged model. The cached-output `ReLU` derivative then treats that `0.0` like a dead unit. It backpropagates a zero gradient there, instead of propagating the `NaN` further back.
Whether the product threads is gemmkit's decision, not the layer's. It compares `batch * input_width * units` against a work gate, 589,824 by default. That count is the raw product, not a FLOP count with a factor of 2. Below the gate, the product stays on 1 thread. Above it, the worker count ramps with the work, rather than moving straight to the full machine.
Small layers in a tight loop stay serial by design, because the per-call dispatch would dominate. These products also nest safely inside an outer parallel region. The backend is described in full in [Matrix Multiplication](../Chapter-06/6.2._Matrix_Multiplication.md), and the tunable gates in [Performance Tuning and Parallelism](../Chapter-07/7.3._Performance_Tuning_and_Parallelism.md).
The elementwise activation that follows has its own, separate parallelism gate. Each activation falls into 1 of 2 cost classes. `ReLU`, `LeakyReLU`, `Softsign`, and `HardSigmoid` are memory-bound "cheap maps" whose crossover sits at 4,000,000 elements. At any practical layer size, they run serial. `Sigmoid`, `Tanh`, `Softmax`, `ELU`, `SELU`, `Softplus`, and `Exponential` are `exp`-dominated, and go parallel above 131,072 elements. `Linear` copies the tensor and has no gate.
The backward pass uses the same 2 classes, but it picks the class from the *derivative*, not from the forward pass. `Sigmoid`, `Tanh`, `Softmax`, `ELU`, `SELU`, and `Exponential` each compute their derivative from the cached output alone, with no exponential call. Their backward pass therefore uses the cheap-map gate, even though their forward pass uses the `exp` gate. `Softplus` is the 1 exception: its derivative still needs 1 exponential per element, so its backward pass keeps the `exp` gate too. Moving either gate only trades serial for parallel. The results are identical, and every product is run-to-run deterministic on the same machine.
`forward` takes `&self`, and it parks the input and the activated output in the context. `Ctx::training()` selects the mode that parks them, and `Ctx::inference()` selects the mode that parks nothing. The 2 modes compute the same values, and the inference mode saves the cost of the parking. The layer itself holds no cache at all, which is why several threads can run inference through 1 model. The backward pass takes the parked output back and differentiates the activation with it.
It then computes 3 quantities. The weight gradient is `input^T * grad`, a GEMM. The bias gradient is the column-sum over the batch. The input gradient is `grad * W^T`, another GEMM. An input of rank 3 or more folds its leading axes into the row axis first. The 3 quantities therefore stay the same 3 products at every rank.
`backward` returns the input gradient, and it puts the other 2 in the gradient store of the context, under the names `kernel` and `bias`. Read one back with `ctx.grads().get(ParamId::new(scope, name))`, where `scope` is the position of the layer in the model that drives it. A bias-free layer computes no bias gradient, so the store holds none. The store sums, so a layer that 2 positions of a graph model share receives the total of the 2 backward passes.
Calling `backward` before `forward` returns `Error::NeuralNetwork(NnError::ForwardPassNotRun)`. An upstream gradient whose shape does not match the parked output returns `Error::ShapeMismatch`. Both are errors, never panics.
## 3.2.7. Reading and setting weights
There is no standalone `weight()` / `bias()` getter of its own. The accessors come from the `LayerBase` trait, and every layer of the crate answers the same 3. All 3 need a built layer, because an unbuilt layer holds no array.
`weights()` gives a `Vec<WeightRef<'_>>`, 1 entry per array the layer holds. Each entry carries the `name` of the array, its `kind`, and a read view of the live storage. `weights_mut()` gives the same list for writing, under the same names and in the same order. `weight(name)` gives 1 read view, or `None` when the layer holds no array under that name. None of the 3 paths clones anything.
`kind` is a `WeightKind`, and it says whether an optimizer updates the array. `Trainable` covers a kernel, a bias, and a normalization scale or shift. `NonTrainable` covers state that the layer keeps and no optimizer writes. A `Dense` layer reports `Trainable` for both of its arrays.
`parameters_mut()` is the accessor that the optimizer takes, and the same trait holds it. It gives 1 `ParamRef` per trainable tensor: the name, a flat mutable view of the data, and a `decays` flag. The flag tells decoupled weight decay to skip a bias. The entry carries no gradient, because a gradient belongs to 1 pass and a parameter outlives every pass. The optimizer builds the address `ParamId::new(scope, name)`, reads the gradient of that address out of the store, and skips a tensor that holds none.
`kernel` has shape `(input_width, units)`, where the width is the one the build settled. `bias` has shape `(1, units)`, as the construction example in 3.2.2 shows. To write parameters, use `set_weights(weights, bias)`, which validates both shapes. These same names are the on-disk addresses. A checkpoint writes each array as `<scope>.<name>`, where `scope` is the position of the layer. The kernel of the first layer of a model is therefore `0.kernel`. Anything you can read here is what round-trips through save and load, in [Saving and Loading Weights](./3.9._Saving_and_Loading_Weights.md).
## 3.2.8. 2 worked models
A regression net ends in a `Linear` head, and trains against mean squared error. The hidden layer fuses `ReLU`. The output layer fuses `Linear`, because a regression target is unbounded:
```rust
use ndarray::Array;
use rustyml::neural_network::Shape;
use rustyml::neural_network::layers::{Activation, Dense};
use rustyml::neural_network::losses::MeanSquaredError;
use rustyml::neural_network::optimizers::SGD;
use rustyml::neural_network::sequential::SequentialBuilder;
fn main() {
// 4 samples, 3 features, 1 continuous target each.
let x = Array::from_shape_vec(
(4, 3),
vec![0.0, 0.1, 0.2, 1.0, 0.9, 0.8, 0.2, 0.1, 0.0, 0.9, 1.0, 0.8],
)
.unwrap()
.into_dyn();
let y = Array::from_shape_vec((4, 1), vec![0.3, 2.7, 0.3, 2.7])
.unwrap()
.into_dyn();
let mut model = SequentialBuilder::new()
.add(Dense::new(8, Activation::ReLU).unwrap()) // hidden, fused ReLU
.add(Dense::new(1, Activation::Linear).unwrap()) // regression head: identity
.build(&Shape::known(x.shape()))
.unwrap();
model.compile(SGD::new(0.05, 0.9, false, 0.0).unwrap(), MeanSquaredError::new());
model.fit(&x, &y, 20).unwrap();
let preds = model.predict(&x).unwrap();
println!("prediction shape: {:?}", preds.shape()); // [4, 1]
}
```
A classifier ends in a `Softmax` head over one-hot targets, paired with `CategoricalCrossEntropy::new(false)`, because the head emits probabilities:
```rust
use ndarray::Array;
use rustyml::neural_network::Shape;
use rustyml::neural_network::layers::{Activation, Dense};
use rustyml::neural_network::losses::CategoricalCrossEntropy;
use rustyml::neural_network::optimizers::Adam;
use rustyml::neural_network::sequential::SequentialBuilder;
fn main() {
// 4 samples, 3 features, 2 classes (one-hot targets).
let x = Array::from_shape_vec(
(4, 3),
vec![0.0, 0.1, 0.2, 1.0, 0.9, 0.8, 0.1, 0.0, 0.2, 0.8, 1.0, 0.9],
)
.unwrap()
.into_dyn();
let y = Array::from_shape_vec((4, 2), vec![1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0])
.unwrap()
.into_dyn();
let mut model = SequentialBuilder::new()
.add(Dense::new(8, Activation::ReLU).unwrap())
.add(Dense::new(2, Activation::Softmax { axis: -1 }).unwrap()) // probability head
.build(&Shape::known(x.shape()))
.unwrap();
model.compile(
Adam::new(0.01, 0.9, 0.999, 1e-8, 0.0).unwrap(),
CategoricalCrossEntropy::new(false),
);
model.fit(&x, &y, 20).unwrap();
let probs = model.predict(&x).unwrap();
// Each row is a distribution over the 2 classes, summing to 1.
println!("class-probability shape: {:?}", probs.shape()); // [4, 2]
}
```
Swap the head to `Activation::Linear`, and the loss to `CategoricalCrossEntropy::new(true)`. This classifier then trains on the more stable fused-logits path, and `predict` produces raw scores instead of probabilities. Which optimizer to compile with, and how learning rate and momentum interact with the dead-neuron and saturation behavior above, is the subject of [Optimizers](./3.4._Optimizers.md).