rustyml 0.14.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. This means deep stacks of recurrent layers, in the Keras sense, are not possible here. Plan around this limitation.

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.