Bidirectional RNNs › PyTorch Implementation
Language AI HandbookAll chapters Bidirectional RNNs › PyTorch ImplementationOne section · Sections in this chapter8 / 16 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
PyTorch makes bidirectional RNNs straightforward with a single constructor argument. Let's build a BiLSTM sequence tagger and apply it to a POS tagging task.
Setup and Data Preparation
We start by importing dependencies and preparing a small toy dataset for demonstration.
import numpy as np
import torch
# Toy corpus: (sentence_tokens, pos_tags)
# Tags: 0=NOUN, 1=VERB, 2=ADP, 3=DET, 4=ADJ
training_data = [
(["banks", "near", "the", "river"], [0, 2, 3, 0]),
(["she", "runs", "every", "morning"], [0, 1, 3, 0]),
(["the", "old", "cat", "sleeps"], [3, 4, 0, 1]),
(["dogs", "chase", "the", "ball"], [0, 1, 3, 0]),
(["he", "reads", "ancient", "texts"], [0, 1, 4, 0]),
(["birds", "fly", "over", "lakes"], [0, 1, 2, 0]),
]
# Build vocabulary from all words
all_words = [w for sent, _ in training_data for w in sent]
vocab = {word: idx + 1 for idx, word in enumerate(sorted(set(all_words)))}
vocab["<PAD>"] = 0
vocab_size = len(vocab)
num_tags = 5 # NOUN, VERB, ADP, DET, ADJVocabulary size: 23 words Number of POS tags: 5 Training sentences: 6
We have a small vocabulary built from 6 training sentences. In a real system you would use hundreds of thousands of sentences, but this is sufficient to demonstrate the architecture.
Defining the BiLSTM Tagger
The bidirectional=True parameter in nn.LSTM is all that changes from a standard LSTM. PyTorch internally creates two LSTM layers and concatenates their outputs.
import torch.nn as nn
class BiLSTMTagger(nn.Module):
def __init__(
self, vocab_size, embedding_dim, hidden_dim, num_tags, padding_idx=0
):
super().__init__()
self.embedding = nn.Embedding(
vocab_size, embedding_dim, padding_idx=padding_idx
)
# bidirectional=True creates forward + backward LSTM
self.lstm = nn.LSTM(
input_size=embedding_dim,
hidden_size=hidden_dim,
num_layers=1,
batch_first=True,
bidirectional=True, # The key parameter
)
# Output dimension is hidden_dim * 2 because forward + backward are concatenated
self.classifier = nn.Linear(hidden_dim * 2, num_tags)
def forward(self, x):
# x: (batch, seq_len)
embeddings = self.embedding(x) # (batch, seq_len, embed_dim)
lstm_out, _ = self.lstm(embeddings) # (batch, seq_len, hidden_dim * 2)
logits = self.classifier(lstm_out) # (batch, seq_len, num_tags)
return logitsInput shape: [1, 4] (batch=1, seq_len=4) Output shape: [1, 4, 5] (batch=1, seq_len=4, num_tags=5) Model parameters: Total parameters: 13,493
Notice the output shape: (batch=1, seq_len=4, num_tags=5). The model produces a tag distribution for every position in the sequence simultaneously, which is what sequence labeling means: one prediction per input token, not one prediction per sequence.
Understanding the Output Tensor
The output tensor from a bidirectional LSTM in PyTorch deserves careful examination. When bidirectional=True, the lstm_out tensor at each position is the concatenation of the forward hidden state and the backward hidden state. Specifically, lstm_out[b, t, :hidden_dim] is the forward hidden state at position t in batch element b, and lstm_out[b, t, hidden_dim:] is the backward hidden state at position t.
This means the first half of the output vector at each position carries information about all tokens to the left (and the current token), and the second half carries information about all tokens to the right (and the current token). The classifier layer, a simple linear projection, takes this concatenated vector and maps it to tag scores. In a real system, you might use a more sophisticated output layer, but even a linear classifier applied independently to each position's concatenated representation is often highly effective because the BiLSTM has already done the hard work of context integration.
The final hidden states, which PyTorch returns as the second output of lstm(), work slightly differently under bidirectionality. PyTorch returns them with shape (num_directions * num_layers, batch, hidden_dim). For a single-layer bidirectional LSTM, this is shape (2, batch, hidden_dim), where index 0 contains the forward LSTM's final hidden state (at position ) and index 1 contains the backward LSTM's final hidden state (at position ). If you want a single fixed-length representation for the whole sequence, a common approach is to concatenate these two terminal states.
Comparing Unidirectional and Bidirectional Models
The key shape difference from a unidirectional LSTM is that the output dimension is hidden_dim * 2.
class UnidirectionalLSTMTagger(nn.Module):
def __init__(self, vocab_size, embedding_dim, hidden_dim, num_tags):
super().__init__()
self.embedding = nn.Embedding(vocab_size, embedding_dim)
self.lstm = nn.LSTM(
embedding_dim, hidden_dim, batch_first=True, bidirectional=False
)
self.classifier = nn.Linear(hidden_dim, num_tags)
def forward(self, x):
emb = self.embedding(x)
out, _ = self.lstm(emb)
return self.classifier(out)
# Build both models with equivalent total capacity
uni_model = UnidirectionalLSTMTagger(
vocab_size, 16, 64, num_tags
) # larger hidden
bi_model = BiLSTMTagger(
vocab_size, 16, 32, num_tags
) # smaller hidden, 2x output
uni_params = sum(p.numel() for p in uni_model.parameters())
bi_params = sum(p.numel() for p in bi_model.parameters())Parameter comparison (matched total capacity): Unidirectional LSTM (hidden=64): 21,685 parameters Bidirectional LSTM (hidden=32): 13,493 parameters Note: BiLSTM hidden=32 gives 2 x 32 = 64 effective units per position
A bidirectional LSTM with hidden dimension is not twice as expensive as a unidirectional LSTM with the same . If you want to match the unidirectional model's capacity, use a bidirectional model with hidden dimension for each direction, achieving comparable parameter counts while still gaining full-context representations.
Training the Model
def prepare_sequence(words, vocab_map, tags):
"""Convert words and tags to tensors."""
word_ids = torch.tensor(
[vocab_map.get(w, 0) for w in words], dtype=torch.long
)
tag_ids = torch.tensor(tags, dtype=torch.long)
return word_ids.unsqueeze(0), tag_ids # Add batch dimension
# Training setup
model = BiLSTMTagger(
vocab_size, embedding_dim=16, hidden_dim=32, num_tags=num_tags
)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()
# Training loop
n_epochs = 100
loss_history = []
model.train()
for epoch in range(n_epochs):
epoch_loss = 0.0
for words, tags in training_data:
x, y = prepare_sequence(words, vocab, tags)
optimizer.zero_grad()
logits = model(x) # (1, seq_len, num_tags)
logits = logits.squeeze(0) # (seq_len, num_tags)
loss = criterion(logits, y)
loss.backward()
nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
epoch_loss += loss.item()
loss_history.append(epoch_loss / len(training_data))Initial loss: 1.5750 Final loss: 0.0001 Loss reduction: 100.0%
Evaluating on Training Sequences
tag_names = ["NOUN", "VERB", "ADP", "DET", "ADJ"]
def predict_tags(model, words, vocab_map):
"""Predict POS tags for a list of words."""
model.train(False)
with torch.no_grad():
x = torch.tensor([[vocab_map.get(w, 0) for w in words]])
logits = model(x).squeeze(0) # (seq_len, num_tags)
predictions = logits.argmax(dim=-1) # (seq_len,)
return [tag_names[p.item()] for p in predictions]
# Test on training sentences
results = []
for words, true_tags in training_data:
predicted = predict_tags(model, words, vocab)
true_names = [tag_names[t] for t in true_tags]
correct = sum(p == t for p, t in zip(predicted, true_names))
results.append((words, true_names, predicted, correct))
total_correct = sum(r[3] for r in results)
total_tokens = sum(len(r[0]) for r in results)Token accuracy on training set: 24/24 = 100.0% Sentence: banks near the river True: ['NOUN', 'ADP', 'DET', 'NOUN'] Predicted: ['NOUN', 'ADP', 'DET', 'NOUN'] Correct: 4/4 Sentence: she runs every morning True: ['NOUN', 'VERB', 'DET', 'NOUN'] Predicted: ['NOUN', 'VERB', 'DET', 'NOUN'] Correct: 4/4 Sentence: the old cat sleeps True: ['DET', 'ADJ', 'NOUN', 'VERB'] Predicted: ['DET', 'ADJ', 'NOUN', 'VERB'] Correct: 4/4
Key Parameters
The key parameters for nn.LSTM and nn.GRU when using bidirectionality are:
- bidirectional: Set to
Trueto run both directions. WhenTrue, the module internally creates two recurrent networks with independent weight matrices. - hidden_size: The dimension of each directional hidden state. The output tensor's last dimension will be
hidden_size * 2because the two directions are concatenated. - num_layers: Number of stacked BiRNN layers. Each layer takes the concatenated output of the previous layer as input to both its forward and backward networks.
- batch_first: Set to
Trueso tensors have shape(batch, seq_len, features)rather than(seq_len, batch, features).