🎓 A1Learner
🔍

PyTorch LSTM Cells

An LSTM is a recurrent network whose gates protect its memory over long stretches. A gate is a learned filter that scales values between none and all.

Four gates in one matrix

One weight matrix holds all four gate computations. Hidden size 16 gives 64 rows, four per hidden unit.

lstm = nn.LSTM(input_size=1, hidden_size=16, batch_first=True)
print(tuple(lstm.weight_ih_l0.shape))

Two states come back

The call returns the output and a pair. The pair holds the hidden state h and the cell state c.

x = torch.zeros(2, 5, 1)
output, (h, c) = lstm(x)
print(tuple(output.shape), tuple(h.shape), tuple(c.shape))

Same shell, new cell

One model shell accepts any cell class. Swapping nn.RNN for nn.LSTM changes only the cell line.

class SineModel(nn.Module):
    def __init__(self, cell, hidden):
        super().__init__()
        self.cell = cell(1, hidden, batch_first=True)
        self.head = nn.Linear(hidden, 1)

Example

import torch
from torch import nn


def make_windows(wave, width):
    inputs = []
    targets = []
    for start in range(len(wave) - width):
        inputs.append(wave[start:start + width])
        targets.append(wave[start + width])
    return torch.stack(inputs).unsqueeze(-1), torch.stack(targets).unsqueeze(-1)


class SineModel(nn.Module):
    def __init__(self, cell, hidden):
        super().__init__()
        self.cell = cell(1, hidden, batch_first=True)
        self.head = nn.Linear(hidden, 1)

    def forward(self, x):
        output, _ = self.cell(x)
        return self.head(output[:, -1])


def train_and_score(cell, x_train, y_train, x_test, y_test):
    torch.manual_seed(0)
    model = SineModel(cell, 16)
    loss_fn = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    for epoch in range(50):
        loss = loss_fn(model(x_train), y_train)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    with torch.no_grad():
        return loss_fn(model(x_test), y_test).item()


wave = torch.sin(torch.arange(300) * 0.1)
x, y = make_windows(wave, 12)
x_train, y_train = x[:180], y[:180]
x_test, y_test = x[180:], y[180:]

rnn_error = train_and_score(nn.RNN, x_train, y_train, x_test, y_test)
lstm_error = train_and_score(nn.LSTM, x_train, y_train, x_test, y_test)

print("Train windows:", len(x_train))
print("Test windows:", len(x_test))
print("RNN test MSE:", round(rnn_error, 4))
print("LSTM test MSE:", round(lstm_error, 4))

Key points

  • Cell state: a second memory line that runs straight through the steps
  • Forget gate: decides what to drop from the cell state
  • Input gate: decides what new material to write into the cell state
  • Output gate: decides what part of the cell state to show as the hidden state
  • Use nn.LSTM as a drop-in cell with the same call shape as nn.RNN
  • An LSTM returns the output plus a pair of states, hidden and cell
  • The gates protect the memory, restated plainly: the cell state can carry a value across many steps nearly unchanged

In this example, one training function trains the same model shell twice on the same sine windows. Seed 0 resets before each build, so both models start from the same discipline. The train windows number 180 and the test windows 108. The plain RNN reaches a test MSE of 0.0019. The LSTM reaches 0.002, a shade higher on this short task. The honest reading is a tie at this scale. Twelve-step windows give the plain cell little to forget. The LSTM earns its keep on longer stretches, where the forget gate can protect old material. The next lesson adds a third cell to the same comparison.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does the forget gate decide?