LSTM and GRU networks in Rust, with gradients you can trust

Quick start

Inputs are column vectors (features × batch). The state holds the hidden and cell state of every layer:

use ndarray::Array2;
use rust_lstm::LSTMNetwork;

fn main() {
    let mut network = LSTMNetwork::new(3, 10, 2); // input size, hidden size, layers
    let input = Array2::from_shape_vec((3, 1), vec![0.5, 0.1, -0.3]).unwrap();
    let state = network.zero_state(1);

    // One time step: output of the top layer and the next state
    let (output, next_state) = network.forward(&input, &state);
    println!("output shape: {:?}, layers: {}", output.dim(), next_state.h.len());
}

For whole sequences use forward_sequence, or forward_sequence_with_cache followed by backward_sequence to get gradients through time.

Train a model

A two-layer LSTM learns to predict the next value of a sine wave. The trainer runs backpropagation through the whole sequence, clips gradients and updates with Adam:

use ndarray::Array2;
use rust_lstm::{seed, Adam, LSTMNetwork, LSTMTrainer, MSELoss, TrainingConfig};

fn main() {
    seed(42); // reproducible weights and dropout masks

    let wave: Vec<f64> = (0..200).map(|i| (i as f64 * 0.1).sin()).collect();
    let step = |x: f64| Array2::from_elem((1, 1), x);
    let train_data: Vec<_> = wave
        .windows(21)
        .step_by(5)
        .map(|w| {
            let inputs = w[..20].iter().map(|&x| step(x)).collect();
            let targets = w[1..].iter().map(|&x| step(x)).collect();
            (inputs, targets)
        })
        .collect();

    // input size 1, hidden size 1 (the top layer's output is the prediction), 2 layers
    let network = LSTMNetwork::new(1, 1, 2);
    let mut trainer = LSTMTrainer::new(network, MSELoss, Adam::new(0.01)).with_config(TrainingConfig {
        epochs: 200,
        clip_gradient: Some(1.0),
        ..Default::default()
    });

    trainer.train(&train_data, None);
    println!("MSE after training: {:.5}", trainer.evaluate(&train_data));
}

For wider networks, add a LinearLayer on top of the hidden state; the examples cover stock and weather prediction, text classification with a BiLSTM, character-level text generation, early stopping and batched training.

What is included

• Networks: multi-layer LSTM, GRU and bidirectional LSTM, peephole LSTM cell, linear (dense) layer, character embeddings.
• Training: full backpropagation through time, gradient clipping, early stopping, validation metrics, padded variable-length batches.
• Optimizers and schedules: SGD, Adam, RMSprop and 12 learning-rate schedulers including OneCycle, cosine annealing with warm restarts, warmup and ReduceLROnPlateau.
• Regularisation: input, recurrent, output, cell-update and variational dropout, and zoneout.
• Losses: MSE, MAE and cross-entropy with softmax.
• Persistence: save and load models as JSON or binary.

How correctness is checked

• cargo test --test gradient_check compares every analytic gradient (LSTM with 1–3 layers, GRU, BiLSTM in all combine modes, peephole LSTM, embedding → LSTM → linear → cross-entropy) with central finite differences, also in training mode with every kind of dropout and zoneout.
• A validation suite loads the weights into torch.nn.LSTM and compares outputs and BPTT gradients in float64, and replays Adam and RMSprop steps against torch.optim. Differences stay at float64 rounding, and the check runs in CI.

Version 0.9 replaced truncated one-step gradients with full BPTT and gave each stacked layer its own state; 0.10 fixed Adam's per-parameter bias correction. If you trained models with an older version, retrain them.

Used in research

The gradient-reach study trains LSTM and GRU networks with rust-lstm 0.10.0 to measure how dropout, zoneout and chrono gate initialisation change how far gradients reach back in time. Its code and training runs are on Zenodo and GitHub.

Cite

If rust-lstm helps your work, please cite the software. The concept DOI 10.5281/zenodo.23178447 always resolves to the latest version.

@software{kholodniak_rust_lstm,
  author  = {Kholodniak, Oleksandr},
  title   = {rust-lstm: recurrent neural networks with verified backpropagation through time in Rust},
  year    = {2026},
  version = {0.10.0},
  doi     = {10.5281/zenodo.23178447},
  url     = {https://github.com/SyntaxSpirits/rust-lstm}
}