rustyml 0.14.0

A high-performance machine learning & deep learning library in pure Rust, offering ML algorithms and neural network support
Documentation
# 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 five 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 stack inside a `Sequential` model.

## 3.2.1. The fused-activation design

The constructor takes the activation as its third argument, not as a separate layer:

```rust,ignore
pub fn new(
    input_dim: usize,
    units: usize,
    activation: impl Into<Activation>,
) -> Result<Dense, Error>
```

Internally the layer stores an `Activation` *value*. This is a plain `Copy` enum with 5 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` 5 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 one 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 the activations as standalone layers (`ReLU`, `Sigmoid`, `Tanh`, `Softmax`, `Linear`). 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 two forms are equivalent:

```rust,ignore
// These two stacks compute the same thing.
model.add(Dense::new(64, 32, Activation::ReLU).unwrap());          // fused: one cached output

model.add(Dense::new(64, 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 two. The backward pass then differentiates the activation in terms of that cached output (see 3.2.5). 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(input_dim, units, activation)` returns `Result<Dense, Error>`. Both dimensions must be non-zero. Passing `0` for either yields `Error::InvalidParameter` (see [Error Handling](../Chapter-01/1.6._Error_Handling.md)). `input_dim` is the number of features per row. `units` is the number of neurons, and therefore the output width.

The parameter count is `input_dim * units + units`: one weight per (input, output) pair, plus one bias per output. `param_count()` reports it as `TrainingParameters::Trainable(n)`. `output_shape()` renders `(None, units)`, where `None` is the dynamic batch dimension, mirroring Keras' summary. A `Dense(4, 3, ...)` therefore holds a `4 x 3` weight matrix and a `1 x 3` bias. That is `12 + 3 = 15` trainable scalars.

Input must be a 2-D tensor of shape `(batch, input_dim)`. A 1-D or 3-D input is rejected with `Error::InvalidInput`, rather than being silently reshaped. If you are feeding the output of a convolutional or recurrent stack, flatten it to 2 dimensions first. The following program constructs a layer, inspects it, and injects known weights:

```rust
use ndarray::Array2;
use rustyml::neural_network::layers::layer_weight::LayerWeight;
use rustyml::neural_network::layers::{Activation, Dense, TrainingParameters};
use rustyml::neural_network::traits::Layer;

fn main() {
    // A 4 -> 3 dense layer with a fused ReLU activation.
    let mut dense = Dense::new(4, 3, Activation::ReLU).unwrap();

    // 4*3 weights + 3 bias = 15 trainable scalars.
    assert_eq!(dense.param_count(), TrainingParameters::Trainable(15));
    println!("output shape: {}", dense.output_shape()); // (None, 3)

    // Read the freshly initialized parameters without cloning them.
    match dense.get_weights() {
        LayerWeight::Dense(w) => {
            println!("weight {:?}, bias {:?}", w.weight.shape(), w.bias.shape());
        }
        _ => unreachable!(),
    }

    // Inject known weights. The shapes are validated against the layer's config.
    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();
}
```

`set_weights` checks both shapes and returns `Error::NeuralNetwork(NnError::WeightShape)` on a mismatch. A `(3, 3)` weight for a `4 -> 3` layer, or a `(1, 4)` bias, is rejected rather than truncated.

## 3.2.3. Weight initialization

Weights use **Xavier/Glorot uniform** initialization. Each element is drawn from `Uniform(-limit, +limit)`, where `limit = sqrt(6 / (input_dim + units))`. Biases start at exactly zero. This is the fan-in-plus-fan-out Glorot scheme. It matches Keras' `Dense` default exactly (`glorot_uniform` kernel, `zeros` bias).

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. By default the seed comes from the process-global seed if one is set, otherwise from entropy. Two 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 construct the model. This also fixes dropout masks and the fit-time batch shuffle.

You can instead seed one layer explicitly, with `Dense::new(...)?.with_random_state(seed)`. This re-runs Glorot with that seed and 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 five activations

`Activation` has exactly 5 variants. The crate has no LeakyReLU, ELU, GELU, or Swish. 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.

| `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` | shifted `exp`, row-normalized | `a_i * (g_i - sum_j(a_j * g_j))` (row Jacobian) | simplex, each row sums to `1` |
| `Linear` | `x` | pass `g` through | `(-inf, 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.

**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.

**Softmax** turns a row of logits into a probability distribution over the last axis. The forward pass subtracts each row'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 needs at least a 2-D input. A 1-D tensor returns `Error::InvalidInput`. Its backward pass is the true Jacobian-vector product across the row, not an elementwise multiply, and each gradient row sums to zero.

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 two 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.
model
    .add(Dense::new(8, 3, Activation::Linear).unwrap())    // raw logits, NOT probabilities
    .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. Forward, backward, and the GEMM path

The forward pass is `activation(input * W + b)`. It runs 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 one pass. The bias and activation apply in the kernel's epilogue while the output tile is still in registers.

The fused result matches the unfused product plus scalar activation bit for bit, with one 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_dim * 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 one 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. `ReLU` is a memory-bound "cheap map" whose crossover sits at 4,000,000 elements. At any practical layer size, it runs serial. `Sigmoid`, `Tanh`, and `Softmax` are `exp`-dominated, and go parallel above 131,072 elements. Moving these gates only trades serial for parallel. The results are identical, and every product is run-to-run deterministic on the same machine.

`forward` caches the input and the activated output for the backward pass. `predict` is the eval-mode twin. It computes the same values, but writes no caches. The backward pass first differentiates the activation using the cached output.

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.

Calling `backward` before `forward` returns `Error::NeuralNetwork(NnError::ForwardPassNotRun)`. An upstream gradient whose shape does not match the cached output returns `Error::ShapeMismatch`. Both are errors, never panics.

## 3.2.6. Reading and setting weights

There is no standalone `weights()` / `bias()` getter. The accessor is `get_weights()` from the `Layer` trait. It returns a `LayerWeight` enum. For a dense layer, that is `LayerWeight::Dense(DenseLayerWeight { weight, bias })`, where `weight` and `bias` are `Cow<Array2<f32>>` borrowed from the live layer, with no clone.

`weight` has shape `(input_dim, units)`, and `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. This same `LayerWeight` enum is the on-disk weight format. 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.7. Two 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::layers::{Activation, Dense};
use rustyml::neural_network::losses::MeanSquaredError;
use rustyml::neural_network::optimizers::SGD;
use rustyml::neural_network::sequential::Sequential;

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 = Sequential::new();
    model
        .add(Dense::new(3, 8, Activation::ReLU).unwrap()) // hidden, fused ReLU
        .add(Dense::new(8, 1, Activation::Linear).unwrap()) // regression head: identity
        .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::layers::{Activation, Dense};
use rustyml::neural_network::losses::CategoricalCrossEntropy;
use rustyml::neural_network::optimizers::Adam;
use rustyml::neural_network::sequential::Sequential;

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 = Sequential::new();
    model
        .add(Dense::new(3, 8, Activation::ReLU).unwrap())
        .add(Dense::new(8, 2, Activation::Softmax).unwrap()) // probability head
        .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).