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

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). 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). The two forms are equivalent:

// 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). 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:

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

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.

// 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 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, and the tunable gates in Performance Tuning and Parallelism.

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

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:

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.