rustyml 0.15.0

A high-performance machine learning & deep learning library in pure Rust, offering ML algorithms and neural network support
Documentation
# 3.7. Recurrent Layers

RustyML ships 3 recurrent layers: `SimpleRNN`, `LSTM`, and `GRU`. They live under `rustyml::neural_network::layers::recurrent`, and the prelude re-exports them. All 3 share one contract. Each layer consumes a 3D sequence tensor, runs a left-to-right recurrence over the time axis, and returns the final hidden state.

If you know Keras, think of these layers as `SimpleRNN`, `LSTM`, and `GRU` with `return_sequences=False` fixed on. That fixed setting affects how you stack the layers. Read [section 3.7.7](#377-stacking-and-building-a-model) before you build a deep recurrent network.

## 3.7.1. The input/output contract

Every recurrent layer expects a 3D input tensor with shape `(batch_size, timesteps, features)`. It produces a 2D output `(batch_size, units)`. The `features` axis must equal the `input_dim` you passed to the constructor. `units` is the hidden width you set. The recurrence consumes the timestep axis, so it does not appear in the output. Only the last hidden state, `h_T`, survives.

```rust
use ndarray::Array;
use rustyml::neural_network::layers::activation::Tanh;
use rustyml::neural_network::layers::recurrent::SimpleRNN;
use rustyml::neural_network::traits::Layer;

fn main() {
    // input_dim = 4 features per timestep, units = 3 hidden neurons
    let rnn = SimpleRNN::new(4, 3, Tanh::new()).unwrap();

    // (batch = 2, timesteps = 5, features = 4)
    let x = Array::zeros((2, 5, 4)).into_dyn();

    // predict runs the recurrence without recording backward caches
    let out = rnn.predict(&x).unwrap();

    // The timestep axis is gone: only the last hidden state survives.
    println!("output shape: {:?}", out.shape()); // [2, 3] == (batch, units)
}
```

```text
output shape: [2, 3]
```

The 3 layers share one constructor signature: `new(input_dim, units, activation) -> Result<Self, Error>`. The activation argument takes `impl Into<Activation>`. You can pass an `Activation` enum variant, for example `Activation::Tanh`, or a thin layer wrapper such as `Tanh::new()`, `ReLU::new()`, `Sigmoid::new()`, `Linear::new()`, or `Softmax::new()`. Every wrapper converts to the same enum. See [3.2. Dense Layers and Activations](./3.2._Dense_Layers_and_Activations.md) for the full activation catalog.

`new` returns `Error::InvalidParameter` when `input_dim` or `units` is 0. RustyML has no `return_sequences`, `return_state`, `bidirectional`, or in-cell dropout option. Each layer always returns the last state, always runs forward in time, and always uses one dense implementation.

A 2D or 4D input is a hard error, not a silent reshape. `forward` and `predict` return `Error::InvalidInput` for any input that is not 3D. The layer caches its forward activations for the backward pass. Calling `backward` before `forward` returns `Error::NeuralNetwork(NnError::ForwardPassNotRun("SimpleRNN"))`, or `"LSTM"` or `"GRU"` for those layers. See [1.6. Error Handling](../Chapter-01/1.6._Error_Handling.md) for how these error variants work together.

RustyML differs from Keras in one way. The `activation` argument controls only the candidate or output nonlinearity, not the gates. In `SimpleRNN`, this activation applies to every hidden state. In `LSTM` and `GRU`, it applies to the candidate, and in LSTM, it also applies to the cell state before the output gate.

The gates always use sigmoid. RustyML has no separate `recurrent_activation` option like Keras has. The gate nonlinearity is fixed. Tanh is the default activation, and almost every published architecture uses it.

## 3.7.2. SimpleRNN and why gated cells exist

`SimpleRNN` is the textbook Elman recurrence. It starts from a zero hidden state. Each timestep mixes the current input with the previous hidden state, using 2 weight matrices and 1 bias:

```text
h_0 = 0
h_t = activation( x_t @ W + h_{t-1} @ U + b )     for t = 1..T
output = h_T
```

Here, `W` is the input kernel `(input_dim, units)`. `U` is the recurrent kernel `(units, units)`. `b` is the bias `(1, units)`. `@` is a matmul that runs over the batch dimension.

RustyML initializes `W` with Xavier/Glorot uniform, and `U` with an orthogonal matrix (Gram-Schmidt). The orthogonal recurrent kernel is deliberate. It keeps the state transition norm-preserving at initialization. This is a cheap way to delay the vanishing-gradient problem described next.

That problem is vanishing gradients, and its counterpart, exploding gradients. Backpropagating from `h_T` to `h_1` multiplies the upstream gradient by a fresh Jacobian at every step. That step is roughly `grad_{t-1} = (activation'(h_t) * grad_t) @ U^T`. Chaining `T` of those steps raises a matrix to the `T`-th power. If the matrix's effective magnitude is below 1, the gradient decays exponentially, and the network cannot learn dependencies more than a few steps back. If the magnitude is above 1, the gradient explodes instead.

The orthogonal `U` keeps the `U^T` factor norm-preserving, and `tanh' <= 1` keeps the product bounded. Even so, the product still tends toward zero over long sequences. This is why `SimpleRNN` works only for short sequences, up to a few dozen steps, and fails at long-range dependencies. LSTM and GRU exist to solve this problem.

## 3.7.3. LSTM: an additive memory highway

LSTM adds a second state, the cell state `c_t`. Its update is additive. LSTM uses 3 sigmoid gates to decide what to write, what to keep, and what to read. RustyML stores the 4 weight blocks fused side by side. The column order matches Keras, `[input | forget | cell | output]`, written `[i | f | g | o]`:

```text
i_t = sigmoid( x_t @ W_i + h_{t-1} @ U_i + b_i )     input gate  (how much candidate to write)
f_t = sigmoid( x_t @ W_f + h_{t-1} @ U_f + b_f )     forget gate (how much old cell to keep)
g_t = act(     x_t @ W_g + h_{t-1} @ U_g + b_g )     candidate   ("cell gate")
o_t = sigmoid( x_t @ W_o + h_{t-1} @ U_o + b_o )     output gate (how much cell to expose)
c_t = f_t * c_{t-1} + i_t * g_t                      cell state  (additive update)
h_t = o_t * act(c_t)                                 hidden state
```

Here `*` is elementwise multiplication. `act` is the configurable activation, Tanh by default, applied to both the candidate and the cell state. The key line is `c_t = f_t * c_{t-1} + i_t * g_t`. Its gradient with respect to `c_{t-1}` is just `f_t`, an elementwise multiply with no repeated matmul.

When the forget gate is open, meaning `f` is close to 1, the cell state carries the gradient backward with almost no loss. This is a near-identity highway, and the additive term feeds into it. Memory persists, and gradient flows, because the dominant path is addition, not repeated matrix multiplication.

RustyML initializes the **forget-gate bias to 1.0**, and every other bias to zero. This gives the memory highway a head start. Before training shapes the gates, `f` starts biased open, so the cell state, and its gradient, survive from the first epoch. The integration test `lstm_forget_bias_is_one_not_zero` checks this behavior. LSTM has 4 gates' worth of parameters: `param_count = 4 * (input_dim * units + units * units + units)`.

## 3.7.4. GRU: the same idea with merged gates

GRU keeps the additive-blend idea from LSTM. It folds the input and forget gates into a single **update** gate, and it drops the separate cell state. GRU needs only 3 weight blocks instead of 4. RustyML stores them fused, in the order `[update | reset | candidate]`, written `[z | r | h]`. This order matches Keras:

```text
z_t = sigmoid( x_t @ W_z + h_{t-1} @ U_z + b_z )          update gate
r_t = sigmoid( x_t @ W_r + h_{t-1} @ U_r + b_r )          reset gate
n_t = act( x_t @ W_h + (r_t * h_{t-1}) @ U_h + b_h )      candidate
h_t = z_t * h_{t-1} + (1 - z_t) * n_t                     hidden state
```

The update gate `z_t` runs a convex blend. When `z` is close to 1, the layer copies the previous hidden state through unchanged. This acts as a gradient highway, the same way a closed LSTM forget gate does. When `z` is close to 0, the layer replaces the previous state with the fresh candidate.

This is Keras' convention. Some write-ups use the complement, so a `z` value ported from one of those needs flipping. The tests `gru_update_gate_one_keeps_previous_hidden` and `gru_update_gate_zero_takes_the_candidate` check these 2 extremes. The test `gru_fused_kernel_first_block_is_the_update_gate` checks the column order.

Note precisely where the reset gate acts. RustyML computes `r_t * h_{t-1}` before the candidate's recurrent matmul, as `(r_t * h_{t-1}) @ U_h`. This matches the original Cho et al. formulation, which is Keras' `reset_after=False`. It differs from the CuDNN `reset_after=True` variant, which applies the reset after the matmul and needs 2 biases per gate. RustyML uses a single bias per gate here.

GRU's `param_count = 3 * (input_dim * units + units * units + units)`. This is 3 quarters of an LSTM of the same width. In practice, GRU trains a little faster and matches LSTM on many tasks. LSTM sometimes performs better when a task needs a long, precisely controlled memory. Both layers initialize their input kernels with Xavier/Glorot, using the per-gate fan `input_dim + units`, not the fused width. Both layers initialize each gate's recurrent block as an independent orthogonal matrix.

## 3.7.5. Weights, shapes, and setting them by hand

The trainable tensors and their shapes:

| Layer | kernel | recurrent_kernel | bias | fused column blocks |
|---|---|---|---|---|
| SimpleRNN | `(input_dim, units)` | `(units, units)` | `(1, units)` | none |
| LSTM | `(input_dim, 4 * units)` | `(units, 4 * units)` | `(1, 4 * units)` | `[i \| f \| g \| o]` |
| GRU | `(input_dim, 3 * units)` | `(units, 3 * units)` | `(1, 3 * units)` | `[z \| r \| h]` |

Fusing every gate into one matrix is not just cosmetic. It lets the input projection and the recurrent projection run as one large GEMM per timestep, instead of one GEMM per gate. This gives a real cache and SIMD benefit. Use `get_weights()` to inspect the live arrays. It returns a `LayerWeight::{SimpleRNN,LSTM,GRU}` value, carrying borrowed `kernel`, `recurrent_kernel`, and `bias` fields:

```rust
use rustyml::neural_network::layers::activation::Tanh;
use rustyml::neural_network::layers::layer_weight::LayerWeight;
use rustyml::neural_network::layers::recurrent::LSTM;
use rustyml::neural_network::traits::Layer;

fn main() {
    // input_dim = 4, units = 8. with_random_state makes the init reproducible.
    let lstm = LSTM::new(4, 8, Tanh::new()).unwrap().with_random_state(42);

    match lstm.get_weights() {
        LayerWeight::LSTM(w) => {
            // All 4 gates are fused side by side: width == 4 * units.
            println!("kernel           {:?}", w.kernel.shape()); // [4, 32]
            println!("recurrent_kernel {:?}", w.recurrent_kernel.shape()); // [8, 32]
            println!("bias             {:?}", w.bias.shape()); // [1, 32]
        }
        _ => unreachable!(),
    }
}
```

```text
kernel           [4, 32]
recurrent_kernel [8, 32]
bias             [1, 32]
```

Call `with_random_state(seed)` for reproducible initialization. It re-runs the kernel and recurrent-kernel draws deterministically, and the forget-bias-1.0 rule still applies. Without a seed, RustyML seeds the weights from the global seed or from entropy. See [7.1. Reproducibility and Random Seeds](../Chapter-07/7.1._Reproducibility_and_Random_Seeds.md).

To install weights by hand, for example when you port from another framework or unit-test an exact recurrence, every layer has `set_weights(kernel, recurrent_kernel, bias)`. This method takes the fused matrices directly. LSTM and GRU also have `set_gate_weights(...)`. It accepts one `(kernel, recurrent_kernel, bias)` triple per gate, and concatenates the triples into the fused layout for you.

The per-gate argument order is `(input, forget, cell, output)` for LSTM, 12 arrays in total. For GRU it is `(reset, update, candidate)`, 9 arrays in total. Any shape mismatch returns `Error::NeuralNetwork(NnError::WeightShape { .. })`. When you save or load a whole model, these arrays serialize as-is. See [3.9. Saving and Loading Weights](./3.9._Saving_and_Loading_Weights.md).

## 3.7.6. BPTT and the cost model

Training uses backpropagation through time with a **full unroll**. There is no truncation window. On the forward pass, the layer caches everything the backward pass needs. `SimpleRNN` stores every hidden state, with `h_0 = 0` prepended. `LSTM` also stores the cell states, `activation(c_t)`, and all 4 gate activations per timestep. `GRU` stores the reset and update gates, the candidate, and the `r_t * h_{t-1}` product.

`predict` passes `None` for these caches and skips the recording and its clones. This is why inference costs less than a training `forward` call. Memory scales linearly with `timesteps`. Long sequences cost RAM, not just time.

`backward` walks the timesteps in reverse. It expects a 2D upstream gradient `(batch, units)`, the gradient of the loss with respect to the final hidden state. This is the only gradient the layer needs, because the final state is the only thing the layer emitted. At each step, `backward` computes the per-timestep pre-activation gradient. It threads `grad_h`, and `grad_c` for LSTM, back to the previous step. This part is inherently **sequential over time**.

`backward` then batches the weight-gradient reductions. It collapses the per-timestep `dz` values over `(batch, timesteps)`, so the kernel and bias gradients each fall out of a single large GEMM. The recurrent-kernel gradient does too, except for GRU, which needs 2 GEMMs for its 2 distinct recurrent inputs. RustyML stores gradients with replace semantics, and does **not** clip them inside the layer. If you need gradient clipping, use an optimizer that offers it. See [3.4. Optimizers](./3.4._Optimizers.md).

The performance shape is **sequential over the time axis, and parallel over the batch and fused gate axis**. The input projection `x @ W` does not depend on the recurrence. RustyML computes it once, as a single batched GEMM across all timesteps, up front. Only the `h_{t-1} @ U` term must run step by step. Each of those steps is itself a batch-parallel GEMM.

GRU saves a little more work here. It fuses the reset and update recurrent projections into one GEMM, because both gates read `h_{t-1}`. Only the candidate's recurrent projection stays separate, because its input `r_t * h_{t-1}` depends on the freshly computed reset gate.

Every matmul goes to the gemmkit backend (see [6.2. Matrix Multiplication](../Chapter-06/6.2._Matrix_Multiplication.md)). gemmkit sizes its own parallelism from the amount of work. Wide layers and large batches get thread parallelism automatically. Small layers stay serial, to avoid overhead. A timestep's fused gate projection always applies its bias in the same pass as the product. Activation fusion into that pass happens only for `SimpleRNN` with `ReLU`, not for `LSTM`, `GRU`, or any other activation.

The practical result: more sequences, meaning a larger batch, parallelize well. More timesteps do not, because that axis is a serial dependency chain. [7.3. Performance Tuning and Parallelism](../Chapter-07/7.3._Performance_Tuning_and_Parallelism.md) covers the thread-pool settings.

## 3.7.7. Stacking and building a model

A recurrent layer returns only the last hidden state, a 2D `(batch, units)` tensor. You **cannot** feed that output straight into another recurrent layer. A recurrent layer needs a 3D `(batch, timesteps, features)` input, and it rejects a 2D tensor with `Error::InvalidInput`. RustyML has no `return_sequences` option, so a layer cannot emit a per-timestep sequence.

`RepeatVector` bridges that gap. It turns a `(batch, units)` state into a `(batch, n, units)` sequence, so a second recurrent layer can read it. Every one of the `n` steps holds the same vector, so the second layer sees a constant sequence, not the per-timestep states of the first. A deep stack in that sense stays out of reach. That is still the encoder-decoder pattern, in which a fixed context seeds the decoder. See [3.5. Convolutional Layers](./3.5._Convolutional_Layers.md), section 3.5.8, for the layer itself.

The standard pattern uses a recurrent layer as a sequence encoder, followed by a dense head that maps the final state to your targets. This composes cleanly. The recurrent layer turns `(batch, timesteps, features)` into `(batch, units)`, and `Dense` consumes exactly that 2D shape.

The example below learns a real sequence task: predict the sum of a length-3 scalar sequence. It uses an `LSTM` encoder, a `Dense` readout, `Adam`, and mean-squared error. It converges in a few hundred full-batch epochs. Recall from [3.1. The Sequential Model](./3.1._The_Sequential_Model.md) that `fit` runs one full-batch gradient step per epoch:

```rust
use ndarray::Array;
use rustyml::neural_network::sequential::Sequential;
use rustyml::prelude::*;

fn main() {
    // 4 sequences, 3 timesteps, 1 feature. Target = sum of the 3 scalars.
    let x = Array::from_shape_vec(
        (4, 3, 1),
        vec![
            0.1, 0.2, 0.1, // sum 0.4
            0.3, 0.1, 0.2, // sum 0.6
            0.0, 0.2, 0.2, // sum 0.4
            0.2, 0.2, 0.1, // sum 0.5
        ],
    )
    .unwrap()
    .into_dyn();
    let y = Array::from_shape_vec((4, 1), vec![0.4, 0.6, 0.4, 0.5])
        .unwrap()
        .into_dyn();

    let mut model = Sequential::new();
    model
        // Recurrent feature extractor: (batch, 3, 1) -> (batch, 16)
        .add(LSTM::new(1, 16, Tanh::new()).unwrap().with_random_state(42))
        // Read the last hidden state out to a single scalar.
        .add(Dense::new(16, 1, Linear::new()).unwrap())
        .compile(
            Adam::new(0.01, 0.9, 0.999, 1e-8, 0.0).unwrap(),
            MeanSquaredError::new(),
        );

    model.fit(&x, &y, 300).unwrap();

    let pred = model.predict(&x).unwrap();
    println!("target      : {:?}", y.as_slice().unwrap());
    println!("prediction  : {:?}", pred.as_slice().unwrap());
}
```

After 300 epochs, the 4 predictions land within a few percent of `[0.4, 0.6, 0.4, 0.5]`. The LSTM has learned to accumulate the sequence, and the dense head reads the accumulator out. Swap `LSTM` for `GRU` or `SimpleRNN`, and the same code still compiles and trains. On this short sequence, all 3 layer types converge. This is because the vanishing-gradient advantage of the gated cells shows up only on long sequences.

`model.summary()` prints each recurrent layer's output shape as `(None, units)`. The `None` stands for the dynamic batch dimension. For anything larger than a toy dataset, use `fit_with_batches` instead of `fit`, so each epoch takes several mini-batch steps and reshuffles the data.

## 3.7.8. Embedding: from indices to vectors

Every layer on this page reads a float tensor. Text does not arrive that way. A tokenizer gives you word indices such as `[7, 42, 3]`, and an index is a name, not a quantity you can add. `Embedding` is the bridge. It holds a trainable table with 1 row per vocabulary entry, and it replaces each index with that row.

`Embedding::new(input_dim, output_dim)` sets the table shape. `input_dim` is the vocabulary size, so the largest usable index is `input_dim - 1`. `output_dim` is the width of 1 vector. The layer appends that width as a new trailing axis, so the output rank is always 1 more than the input rank.

| Input shape | Output shape | Typical use |
| --- | --- | --- |
| `[N]` | `[N, output_dim]` | 1 index per sample |
| `[N, T]` | `[N, T, output_dim]` | a batch of `T`-token sequences |
| `[N, T, K]` | `[N, T, K, output_dim]` | a batch of token grids |

The middle row is the one that matters here. It is exactly the `(batch_size, timesteps, features)` contract from [3.7.1](#371-the-inputoutput-contract), so an `Embedding` feeds a recurrent layer with no reshape between them.

A `Tensor` holds `f32`, so the indices arrive as floating-point values. The layer truncates each value toward zero and then checks it. `2.0` and `2.9` both select row 2, and `-0.5` selects row 0. That is the rule the Keras cast to a whole number uses. An index outside `0..input_dim`, and any non-finite value, gives an `Error::InvalidInput` that names the offending value.

Keras leaves those 2 cases to its backend. An out-of-range index gives `NaN`, and a negative index wraps to the end of the table.

The table is the only parameter, so `param_count()` reports `input_dim * output_dim`. It starts from a uniform draw over `[-0.05, 0.05]`, which is the Keras default for this layer, not the Xavier/Glorot rule that `Dense` uses. Call `.with_random_state(seed)` to make the draw reproducible. Read the table with `get_weights()`, which gives the `LayerWeight::Embedding` variant, and write it with `set_weights(table)`. Both use the `(input_dim, output_dim)` shape.

The backward pass adds the upstream gradient into the row that each index selected. A row that 2 positions selected gets the sum of both, and a row that no index selected keeps a gradient of exactly 0. The gradient is still stored densely over the whole table, so 1 optimizer step costs `input_dim * output_dim` whatever the batch holds. A momentum term or a weight decay also moves a row whose gradient is 0, because both rules act on the parameter itself. Keras behaves the same way on its dense update path.

The example below is the standard text pipeline in miniature. It classifies a 4-token sequence, and it stacks `Embedding`, `LSTM`, and `Dense`:

```rust
use ndarray::Array;
use rustyml::neural_network::sequential::Sequential;
use rustyml::prelude::*;

fn main() {
    // 4 sentences of 4 token indices each, drawn from a 10-word vocabulary
    let x = Array::from_shape_vec(
        (4, 4),
        vec![
            1.0, 2.0, 3.0, 4.0, // opens with token 1
            5.0, 6.0, 7.0, 8.0, // opens with token 5
            1.0, 2.0, 7.0, 8.0, // opens with token 1
            5.0, 6.0, 3.0, 4.0, // opens with token 5
        ],
    )
    .unwrap()
    .into_dyn();
    // Label 1 when the sentence opens with token 1, and 0 when it opens with token 5
    let y = Array::from_shape_vec((4, 1), vec![1.0, 0.0, 1.0, 0.0])
        .unwrap()
        .into_dyn();

    let mut model = Sequential::new();
    model
        // Indices to vectors: (batch, 4) -> (batch, 4, 8)
        .add(Embedding::new(10, 8).unwrap().with_random_state(42))
        // Sequence encoder: (batch, 4, 8) -> (batch, 16)
        .add(LSTM::new(8, 16, Tanh::new()).unwrap().with_random_state(7))
        // Read the last hidden state out to 1 probability
        .add(Dense::new(16, 1, Sigmoid::new()).unwrap().with_random_state(3))
        .compile(
            Adam::new(0.01, 0.9, 0.999, 1e-8, 0.0).unwrap(),
            BinaryCrossEntropy::new(),
        );

    model.summary();
    model.fit(&x, &y, 200).unwrap();

    let pred = model.predict(&x).unwrap();
    println!("target     : {:?}", y.as_slice().unwrap());
    println!("prediction : {:?}", pred.as_slice().unwrap());
}
```

After 200 epochs the 4 predictions sit within a thousandth of `[1.0, 0.0, 1.0, 0.0]`. `model.summary()` prints the embedding output shape as `(None, 4, 8)` once a forward pass has run, and `Unknown` before that. The layer learns its sequence length from the data rather than from the constructor.

One Keras argument has no counterpart here. `mask_zero` marks index 0 as padding and builds a mask for the layers after it. RustyML propagates no mask, so the flag would change nothing. It is omitted rather than accepted and ignored. Pad your sequences to a fixed length, and let the model learn that the padding index carries no information.