Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Stacked RNNs › Implementation in PyTorch

Michael BrenndoerferMay 13, 20253 min read

Language AI HandbookAll chapters

Part 11 of 67Recurrent Neural Networks

Chapter 82Stacked RNNs › Implementation in PyTorch

Reading section 12 of 16 · Implementation in PyTorch · 3 minSection 12 of 16 · 3 min

Section 12 of 16 · 3 min read

Stacked RNNs › Implementation in PyTorchOne section · Sections in this chapter12 / 16 ↓

Sections in this chapter16

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's nn.LSTM supports stacking natively through the num_layers parameter. Let's walk through a complete implementation covering the key design decisions.

Basic Stacked LSTMLink Copied

In[14]:

Code

import torch.nn as nn


class LanguageModel(nn.Module):
    """Stacked LSTM language model with configurable depth and dropout."""

    def __init__(
        self,
        vocab_size,
        embed_dim,
        hidden_size,
        num_layers,
        dropout=0.3,
        tie_weights=False,
    ):
        super().__init__()
        self.encoder = nn.Embedding(vocab_size, embed_dim)
        self.drop = nn.Dropout(dropout)
        self.rnn = nn.LSTM(
            input_size=embed_dim,
            hidden_size=hidden_size,
            num_layers=num_layers,
            dropout=dropout if num_layers > 1 else 0.0,  # inter-layer dropout
            batch_first=True,
        )
        self.decoder = nn.Linear(hidden_size, vocab_size)

        if tie_weights:
            # Weight tying: share embedding and output projection weights
            # Requires embed_dim == hidden_size
            assert embed_dim == hidden_size, (
                "For weight tying, embed_dim must equal hidden_size"
            )
            self.decoder.weight = self.encoder.weight

        self.init_weights()

    def init_weights(self):
        """Xavier initialization for embeddings and output projection."""
        nn.init.uniform_(self.encoder.weight, -0.1, 0.1)
        nn.init.zeros_(self.decoder.bias)
        nn.init.uniform_(self.decoder.weight, -0.1, 0.1)

    def forward(self, input_ids, hidden=None):
        # input_ids: (batch, seq_len)
        embedded = self.drop(self.encoder(input_ids))
        output, hidden = self.rnn(embedded, hidden)
        output = self.drop(output)
        # Decode to vocab logits
        logits = self.decoder(output)
        return logits, hidden

    def init_hidden(self, batch_size, device):
        """Initialize hidden and cell states to zeros."""
        weight = next(self.parameters())
        h0 = weight.new_zeros(
            self.rnn.num_layers, batch_size, self.rnn.hidden_size
        )
        c0 = weight.new_zeros(
            self.rnn.num_layers, batch_size, self.rnn.hidden_size
        )
        return h0, c0

Out[15]:

Console

Language Model Comparison: Depth vs. Width

Config                               Params    Memory (MB)
----------------------------------------------------------
1 layer,  d=512                  12,351,248           0.8
2 layers, d=512                  14,452,496           1.6

3 layers, d=512                  16,553,744           2.3
4 layers, d=512                  18,654,992           3.1

1 layer,  d=1024                 21,669,648           1.6
2 layers, d=256                   6,182,672           0.8

Key Hyperparameter GuidelinesLink Copied

The key parameters for stacked LSTM architectures are:

num_layers: The depth of the stack. Start at 2 for most tasks. Add layers only if validation performance keeps improving. Values of 2–4 cover the vast majority of practical use cases.
hidden_size: The width of each layer. This is the dominant factor in parameter count because $d^2$ grows quadratically. Common values: 256 (small), 512 (medium), 1024 (large).
dropout: The inter-layer dropout rate. Applied between all consecutive layer pairs when num_layers > 1. Typical values: 0.3–0.5.
embed_dim: The dimensionality of the token embeddings fed to the first layer. Often set equal to hidden_size to avoid an asymmetric first layer.