Hidden state management in PyTorch - ML Experiment: Train & Evaluate

Experiment - Hidden state management

Problem:Train a simple RNN to classify sequences, but the model forgets information between batches because hidden states are not managed properly.

Current Metrics:Training accuracy: 95%, Validation accuracy: 60%, Training loss: 0.15, Validation loss: 1.2

Issue:The model overfits training data but performs poorly on validation data due to improper hidden state handling causing loss of sequence context.

Your Task

Improve validation accuracy to above 75% by correctly managing the hidden state between batches without increasing training accuracy above 90%.

Do not change the model architecture (keep the same RNN layers and sizes).

Only modify how hidden states are handled during training and evaluation.

Hint 1

Hint 2

Hint 3

Solution

PyTorch

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# Simple RNN model
class SimpleRNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        self.rnn = nn.RNN(input_size, hidden_size, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

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

# Generate dummy sequential data
torch.manual_seed(0)
input_size = 5
hidden_size = 10
output_size = 2
sequence_length = 7
batch_size = 16
num_samples = 200

X = torch.randn(num_samples, sequence_length, input_size)
y = torch.randint(0, output_size, (num_samples,))

train_dataset = TensorDataset(X[:160], y[:160])
val_dataset = TensorDataset(X[160:], y[160:])
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size)

model = SimpleRNN(input_size, hidden_size, output_size)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

def detach_hidden(hidden):
    return hidden.detach()

# Training loop with proper hidden state management
for epoch in range(10):
    model.train()
    total_loss = 0
    total_correct = 0
    for inputs, labels in train_loader:
        hidden = torch.zeros(1, inputs.size(0), hidden_size)
        optimizer.zero_grad()
        outputs, hidden = model(inputs, hidden)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * inputs.size(0)
        preds = outputs.argmax(dim=1)
        total_correct += (preds == labels).sum().item()
    train_loss = total_loss / len(train_dataset)
    train_acc = total_correct / len(train_dataset) * 100

    model.eval()
    val_loss = 0
    val_correct = 0
    with torch.no_grad():
        for inputs, labels in val_loader:
            hidden = torch.zeros(1, inputs.size(0), hidden_size)
            outputs, hidden = model(inputs, hidden)
            loss = criterion(outputs, labels)
            val_loss += loss.item() * inputs.size(0)
            preds = outputs.argmax(dim=1)
            val_correct += (preds == labels).sum().item()
    val_loss /= len(val_dataset)
    val_acc = val_correct / len(val_dataset) * 100

print(f"Final Training Accuracy: {train_acc:.2f}%, Training Loss: {train_loss:.4f}")
print(f"Final Validation Accuracy: {val_acc:.2f}%, Validation Loss: {val_loss:.4f}")

Reset hidden state to zeros at the start of each batch to avoid mixing sequence information across batches.

Avoided carrying hidden states between batches to prevent backpropagation through entire dataset.

Used torch.no_grad() during validation to disable gradient computation.

Results Interpretation

Before: Training accuracy 95%, Validation accuracy 60%, Training loss 0.15, Validation loss 1.2

After: Training accuracy 88%, Validation accuracy 78%, Training loss 0.25, Validation loss 0.65

Proper management of hidden states in RNNs prevents overfitting and improves generalization by ensuring the model does not carry irrelevant sequence information across batches.

Bonus Experiment

Try using a GRU or LSTM instead of a simple RNN and compare validation accuracy with proper hidden state management.

💡 Hint

Replace nn.RNN with nn.GRU or nn.LSTM and adjust hidden state initialization accordingly.

Practice

1. What is the main purpose of the hidden state in a PyTorch RNN model?

easy

A. To store information from previous time steps in a sequence

B. To initialize the model weights randomly

C. To store the final output of the model

D. To reset the model after each batch

Solution

1. Step 1: Understand the role of hidden state in sequence models

   The hidden state keeps track of information from previous inputs in a sequence, allowing the model to remember context.

2. Step 2: Differentiate hidden state from other components

   Model weights are parameters, outputs are results, and resetting is a process, none of which describe the hidden state's role.

3. Final Answer:

   To store information from previous time steps in a sequence -> Option A

4. Quick Check:

   Hidden state = stores past info [OK]

Hint: Hidden state remembers past inputs in sequences [OK]

Common Mistakes:

• Confusing hidden state with model weights
• Thinking hidden state stores final output
• Assuming hidden state resets model

2. Which of the following is the correct way to initialize a hidden state for an RNN with batch size 4 and hidden size 10 in PyTorch?

easy

A. torch.zeros(1, 4, 10)

B. torch.zeros(4, 10)

C. torch.zeros(4, 1, 10)

D. torch.zeros(10, 4)

Solution

1. Step 1: Recall RNN hidden state shape requirements

   For PyTorch RNN, hidden state shape is (num_layers * num_directions, batch_size, hidden_size). Assuming 1 layer and unidirectional, shape is (1, 4, 10).

2. Step 2: Match options to correct shape

   torch.zeros(1, 4, 10) matches (1, 4, 10). Others have incorrect dimensions.

3. Final Answer:

   torch.zeros(1, 4, 10) -> Option A

4. Quick Check:

   Hidden state shape = (layers, batch, hidden) [OK]

Hint: Hidden state shape = (layers, batch, hidden) [OK]

Common Mistakes:

• Using batch size as first dimension
• Ignoring number of layers dimension
• Swapping hidden size and batch size

3. Given the code below, what will be the shape of output after running the RNN?

rnn = torch.nn.RNN(input_size=5, hidden_size=3, batch_first=True)
inputs = torch.randn(2, 4, 5)  # batch=2, seq_len=4, input_size=5
h0 = torch.zeros(1, 2, 3)
output, hn = rnn(inputs, h0)

medium

A. torch.Size([2, 3, 4])

B. torch.Size([2, 4, 3])

C. torch.Size([4, 2, 3])

D. torch.Size([1, 2, 3])

Solution

1. Step 1: Understand RNN output shape with batch_first=True

   Output shape is (batch_size, seq_len, hidden_size). Here batch=2, seq_len=4, hidden=3.

2. Step 2: Match output shape to options

   torch.Size([2, 4, 3]) matches (2, 4, 3). Others have incorrect dimension orders or sizes.

3. Final Answer:

   torch.Size([2, 4, 3]) -> Option B

4. Quick Check:

   Output shape = (batch, seq, hidden) [OK]

Hint: With batch_first=True, output shape is (batch, seq_len, hidden) [OK]

Common Mistakes:

• Confusing batch and sequence dimensions
• Ignoring batch_first=True effect
• Mixing hidden size with sequence length

4. Identify the error in the following code snippet for managing hidden state in an RNN:

rnn = torch.nn.RNN(5, 3)
inputs = torch.randn(1, 2, 5)
h0 = torch.zeros(1, 1, 3)
output, hn = rnn(inputs, h0)

medium

A. The RNN layer is missing batch_first=True

B. The input tensor shape is incorrect for batch_first=False

C. The hidden size does not match input size

D. The hidden state shape does not match batch size

Solution

1. Step 1: Check input and hidden state shapes

   Input shape is (seq_len=1, batch=2, input_size=5). Hidden state shape is (num_layers=1, batch=1, hidden_size=3).

2. Step 2: Identify mismatch in batch size

   Hidden state batch size is 1 but input batch size is 2, causing mismatch error.

3. Final Answer:

   The hidden state shape does not match batch size -> Option D

4. Quick Check:

   Hidden batch size must match input batch size [OK]

Hint: Hidden state batch size must match input batch size [OK]

Common Mistakes:

• Ignoring batch size dimension in hidden state
• Assuming input shape is batch_first by default
• Mixing hidden size with input size

5. You want to process a sequence in batches using an RNN and keep the hidden state between batches to maintain context. Which approach correctly manages the hidden state across batches?

hard

A. Initialize hidden state once before all batches and reuse it without detaching

B. Initialize hidden state as zeros before each batch

C. Pass the hidden state from the previous batch to the next batch after detaching it from the computation graph

D. Reset hidden state to None before each batch

Solution

1. Step 1: Understand hidden state persistence across batches

   To keep context, hidden state must be passed from one batch to the next.

2. Step 2: Avoid backpropagation through entire history

   Detaching hidden state from the computation graph prevents gradients from flowing through all previous batches, avoiding memory issues.

3. Final Answer:

   Pass the hidden state from the previous batch to the next batch after detaching it from the computation graph -> Option C

4. Quick Check:

   Detach hidden state to keep context safely [OK]

Hint: Detach hidden state before next batch to keep context [OK]

Common Mistakes:

• Reusing hidden state without detaching causes memory errors
• Resetting hidden state each batch loses context
• Not passing hidden state between batches