RNN vs LSTM Memory Benchmarks

gistfile1.txt

# ==============================================================================
# PROGRAM: RNN vs LSTM Memory Benchmark (100% Deterministic)
# ==============================================================================

import torch
import torch.nn as nn
import torch.optim as optim

torch.manual_seed(42)

# Vocabulary Setup:
# 0 = 'apple', 1 = 'banana', 2 = 'filler_noise', 3 = 'recall_trigger'
vocab_size = 4
hidden_dim = 16

# --- TOGGLE SEQUENCE LENGTH HERE ---
seq_len = 20   # <-- Vanilla RNN FAILS (50%), LSTM SUCCEEDS (100%)
# seq_len = 3  # <-- Vanilla RNN SUCCEEDS (100%), LSTM SUCCEEDS (100%)

id_to_word = {0: "apple", 1: "banana", 2: "filler_noise", 3: "recall_trigger"}

# Prepare Sequences
seq_A = [0] + [2] * seq_len + [3]  # Starts with 'apple'
seq_B = [1] + [2] * seq_len + [3]  # Starts with 'banana'

# One-hot encode inputs directly so training is instant and stable
def to_one_hot(seqs):
    batch = []
    for seq in seqs:
        one_hot_seq = []
        for token in seq:
            vec = [0.0] * vocab_size
            vec[token] = 1.0
            one_hot_seq.append(vec)
        batch.append(one_hot_seq)
    return torch.tensor(batch)

inputs_batch = to_one_hot([seq_A, seq_B])  # Shape: [2, seq_len + 2, 4]
targets_batch = torch.tensor([0, 1])       # Targets: 0 ('apple'), 1 ('banana')


# 1. DEFINE VANILLA RNN MODEL
class VanillaRNNModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.rnn = nn.RNN(vocab_size, hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, vocab_size)

    def forward(self, x):
        out, _ = self.rnn(x)
        return self.fc(out[:, -1, :])


# 2. DEFINE LSTM MODEL
class LSTMModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.lstm = nn.LSTM(vocab_size, hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, vocab_size)
        
        # Set forget gate bias to 1.0 to keep Cell State highway wide open
        for name, param in self.lstm.named_parameters():
            if 'bias' in name:
                n = param.size(0)
                param.data[n//4:n//2].fill_(1.0)

    def forward(self, x):
        out, _ = self.lstm(x)
        return self.fc(out[:, -1, :])


# 3. RUN EXPERIMENT
def run_experiment(model_class, model_name):
    torch.manual_seed(42)
    model = model_class()
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.03)

    # Train for 300 epochs
    for epoch in range(300):
        optimizer.zero_grad()
        output = model(inputs_batch)
        loss = criterion(output, targets_batch)
        loss.backward()
        optimizer.step()

    print(f"==================================================")
    print(f"          BENCHMARK: {model_name}")
    print(f"==================================================")
    
    total_correct = 0
    with torch.no_grad():
        output = model(inputs_batch)
        probabilities = torch.softmax(output, dim=1)
        predicted_ids = torch.argmax(probabilities, dim=1)

        for idx in range(2):
            target_id = targets_batch[idx].item()
            predicted_id = predicted_ids[idx].item()
            is_correct = predicted_id == target_id
            
            if is_correct:
                total_correct += 1

            key_word = id_to_word[idx]
            pred_word = id_to_word[predicted_id]
            target_word = id_to_word[target_id]
            confidence = probabilities[idx, target_id].item() * 100

            print(f"Sample {idx+1} | Input Start: '{key_word}' ──► Target: '{target_word}'")
            print(f"         Predicted  : '{pred_word}' (Confidence: {confidence:.2f}%)")
            print(f"         Status     : {'✓ SUCCESS' if is_correct else '✗ FAILED'}\n")

    accuracy = (total_correct / 2) * 100
    print(f"OVERALL ACCURACY: {accuracy:.1f}%\n")


# 4. EXECUTE
run_experiment(VanillaRNNModel, "Vanilla RNN (Gradients Vanish)")
run_experiment(LSTMModel, "LSTM Network (Cell State Highway C_t)")