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

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)
}

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

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

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:

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:

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:

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!(),
    }
}

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.

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

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). 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 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 that fit runs one full-batch gradient step per epoch:

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.