Stacked RNNs › Stacked LSTM for Machine Translation
Language AI HandbookAll chapters Stacked RNNs › Stacked LSTM for Machine TranslationOne section · Sections in this chapter7 / 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
Machine translation was one of the first real-world tasks where stacked LSTMs demonstrated dramatic improvements over single-layer models. The landmark result was Sutskever, Vinyals, and Le's 2014 paper "Sequence to Sequence Learning with Neural Networks," which used a 4-layer stacked LSTM encoder-decoder architecture and achieved state-of-the-art results on English-to-French translation.
The Architecture
The seq2seq architecture for machine translation consists of:
- Encoder: A stacked LSTM reads the source sentence and compresses it into a fixed-size context vector (the final hidden state of the top layer).
- Decoder: A separate stacked LSTM takes the context vector as its initial hidden state and generates the target sentence one token at a time.
The original model used 4 layers of 1000-unit LSTMs in both the encoder and decoder. The input sentence was fed in reverse order to the encoder, a trick that brought the beginning of the source sentence closer to the beginning of the target sentence in the unrolled computation graph, making gradient flow easier.
import torch.nn as nn
class StackedLSTMEncoder(nn.Module):
"""Multi-layer LSTM encoder for seq2seq."""
def __init__(
self, vocab_size, embed_dim, hidden_size, num_layers, dropout=0.3
):
super().__init__()
self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
self.lstm = nn.LSTM(
input_size=embed_dim,
hidden_size=hidden_size,
num_layers=num_layers,
dropout=dropout if num_layers > 1 else 0.0,
batch_first=True,
)
self.dropout = nn.Dropout(dropout)
def forward(self, src_tokens):
# src_tokens: (batch, seq_len)
embedded = self.dropout(self.embedding(src_tokens))
# output: (batch, seq_len, hidden_size)
# hidden: (num_layers, batch, hidden_size), one per layer
# cell: (num_layers, batch, hidden_size), one per layer
output, (hidden, cell) = self.lstm(embedded)
return output, hidden, cell
class StackedLSTMDecoder(nn.Module):
"""Multi-layer LSTM decoder for seq2seq."""
def __init__(
self, vocab_size, embed_dim, hidden_size, num_layers, dropout=0.3
):
super().__init__()
self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
self.lstm = nn.LSTM(
input_size=embed_dim,
hidden_size=hidden_size,
num_layers=num_layers,
dropout=dropout if num_layers > 1 else 0.0,
batch_first=True,
)
self.fc_out = nn.Linear(hidden_size, vocab_size)
self.dropout = nn.Dropout(dropout)
def forward(self, tgt_token, hidden, cell):
# tgt_token: (batch,) -> unsqueeze to (batch, 1)
tgt_token = tgt_token.unsqueeze(1)
embedded = self.dropout(self.embedding(tgt_token))
output, (hidden, cell) = self.lstm(embedded, (hidden, cell))
# output: (batch, 1, hidden_size)
prediction = self.fc_out(output.squeeze(1))
return prediction, hidden, cellEncoder parameters: 6,238,208 Decoder parameters: 12,906,208 Total parameters: 19,144,416 Encoder output shape: (4, 20, 512) (batch, seq_len, hidden) Hidden state shape: (2, 4, 512) (num_layers, batch, hidden) Cell state shape: (2, 4, 512) (num_layers, batch, hidden) Decoder prediction shape: (4, 12000) (batch, tgt_vocab)
The hidden state tensor has shape (num_layers, batch, hidden_size), meaning each layer maintains its own separate hidden state. When the encoder finishes reading the source sentence, you pass all layers' hidden and cell states to the decoder, initializing its entire stack with the compressed source representation.
Why Multiple Layers Mattered
The Sutskever et al. paper explicitly tested 1, 2, 4, and 8 layer models. Their key finding: 4 layers outperformed 1 and 2 layers substantially, while 8 layers degraded (without residual connections). The improvement from 4 layers came from the encoder learning a richer, more hierarchical representation of the source sentence. The final context vector captures both local word-level information (from lower layers) and global sentence structure (from higher layers).