LSTM Gradient Flow › Code Implementation

import numpy as np
import torch

torch.manual_seed(42)
np.random.seed(42)

# Task: detect whether the first token equals 0 in a sequence of random symbols.
# This forces the model to carry information from step 0 across T steps.


def make_dataset(n_samples, T, n_symbols=10):
    X = torch.randint(0, n_symbols, (n_samples, T))
    y = (X[:, 0] == 0).float()
    return X, y


INPUT_SIZE = 10  # one-hot dimension
HIDDEN_SIZE = 32
NUM_CLASSES = 1
BATCH_SIZE = 64
N_TRAIN = 512

class VanillaRNN(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super().__init__()
        self.hidden_size = hidden_size
        self.rnn = nn.RNN(input_size, hidden_size, batch_first=True)
        self.fc = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        out, h_n = self.rnn(x)
        return self.fc(h_n.squeeze(0))


class LSTMModel(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super().__init__()
        self.hidden_size = hidden_size
        self.lstm = nn.LSTM(input_size, hidden_size, batch_first=True)
        self.fc = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        out, (h_n, c_n) = self.lstm(x)
        return self.fc(h_n.squeeze(0))

def measure_gradient_norms(model_class, sequence_lengths, n_train=N_TRAIN):
    """
    For each sequence length, train the model briefly and measure
    the gradient norm with respect to the first-position input.
    """
    grad_norms_by_length = []

    for T in sequence_lengths:
        model = model_class(INPUT_SIZE, HIDDEN_SIZE, NUM_CLASSES)
        optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
        criterion = nn.BCEWithLogitsLoss()

        X, y = make_dataset(n_train, T)
        X_one_hot = torch.zeros(n_train, T, INPUT_SIZE)
        X_one_hot.scatter_(2, X.unsqueeze(2), 1.0)

        # Brief training to move weights away from initialization
        model.train()
        for step in range(50):
            idx = torch.randint(0, n_train, (BATCH_SIZE,))
            xb = X_one_hot[idx]
            yb = y[idx]
            logits = model(xb)
            loss = criterion(logits.squeeze(), yb)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        # Measure gradient norm w.r.t. first-position input
        model.eval()
        X_test, y_test = make_dataset(BATCH_SIZE, T)
        X_test_oh = torch.zeros(BATCH_SIZE, T, INPUT_SIZE)
        X_test_oh.scatter_(2, X_test.unsqueeze(2), 1.0)
        X_test_oh.requires_grad_(True)

        logits = model(X_test_oh)
        loss = criterion(logits.squeeze(), y_test)
        loss.backward()

        # Gradient with respect to the first token's embedding
        grad = X_test_oh.grad[:, 0, :]  # (batch, input_size)
        grad_norm = grad.norm(dim=1).mean().item()
        grad_norms_by_length.append(grad_norm)

    return grad_norms_by_length


sequence_lengths = [5, 10, 20, 30, 50]
rnn_grad_norms = measure_gradient_norms(VanillaRNN, sequence_lengths)
lstm_grad_norms = measure_gradient_norms(LSTMModel, sequence_lengths)

T_long = 30

lstm_final = LSTMModel(INPUT_SIZE, HIDDEN_SIZE, NUM_CLASSES)
optimizer = torch.optim.Adam(lstm_final.parameters(), lr=1e-3)
criterion = nn.BCEWithLogitsLoss()

X_long, y_long = make_dataset(N_TRAIN, T_long)
X_long_oh = torch.zeros(N_TRAIN, T_long, INPUT_SIZE)
X_long_oh.scatter_(2, X_long.unsqueeze(2), 1.0)

for step in range(300):
    idx = torch.randint(0, N_TRAIN, (BATCH_SIZE,))
    xb = X_long_oh[idx]
    yb = y_long[idx]
    logits = lstm_final(xb)
    loss = criterion(logits.squeeze(), yb)
    optimizer.zero_grad()
    loss.backward()
    nn.utils.clip_grad_norm_(lstm_final.parameters(), 1.0)
    optimizer.step()

# In PyTorch LSTM, bias_hh_l0 stores [b_ii, b_if, b_ig, b_io] each of size HIDDEN_SIZE.
# Index HIDDEN_SIZE:2*HIDDEN_SIZE corresponds to the forget gate bias.
forget_bias_hh = lstm_final.lstm.bias_hh_l0.data[HIDDEN_SIZE : 2 * HIDDEN_SIZE]
forget_bias_ih = lstm_final.lstm.bias_ih_l0.data[HIDDEN_SIZE : 2 * HIDDEN_SIZE]
total_forget_bias = forget_bias_hh + forget_bias_ih
forget_gate_init = torch.sigmoid(total_forget_bias)