Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Stacked RNNs › Variational Dropout for Recurrent Nets

Michael BrenndoerferMay 13, 20253 min read

Language AI HandbookAll chapters

Part 11 of 67Recurrent Neural Networks

Chapter 82Stacked RNNs › Variational Dropout for Recurrent Nets

Reading section 8 of 16 · Variational Dropout for Recurrent Nets · 3 minSection 8 of 16 · 3 min

Section 8 of 16 · 3 min read

Stacked RNNs › Variational Dropout for Recurrent NetsOne section · Sections in this chapter8 / 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

Let's look more closely at implementing variational dropout, which is the recommended approach for regularizing stacked RNNs. The key requirement is that the same mask is used across all timesteps in a single forward pass.

In[8]:

Code

import torch
import torch.nn as nn


class VariationalDropout(nn.Module):
    """Applies the same dropout mask at every timestep within a forward pass."""

    def __init__(self, dropout_p):
        super().__init__()
        self.dropout_p = dropout_p

    def forward(self, x):
        # x shape: (batch, seq_len, features)
        if not self.training or self.dropout_p == 0.0:
            return x
        # Sample one mask for the whole sequence (shape: batch x 1 x features)
        mask = x.new_ones(x.size(0), 1, x.size(2))
        mask = torch.bernoulli(mask * (1 - self.dropout_p)) / (
            1 - self.dropout_p
        )
        # Broadcast mask over the seq_len dimension
        return x * mask


class VariationalLSTM(nn.Module):
    """Single-layer LSTM with variational dropout on inputs and between layers."""

    def __init__(self, input_size, hidden_size, dropout_p=0.3):
        super().__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, batch_first=True)
        self.var_drop = VariationalDropout(dropout_p)

    def forward(self, x, hidden=None):
        x = self.var_drop(x)
        out, (h, c) = self.lstm(x, hidden)
        return out, (h, c)


class DeepVariationalLSTM(nn.Module):
    """Stacked LSTM with variational dropout between layers."""

    def __init__(self, input_size, hidden_size, num_layers, dropout_p=0.3):
        super().__init__()
        sizes = [input_size] + [hidden_size] * num_layers
        self.layers = nn.ModuleList(
            [
                VariationalLSTM(sizes[i], hidden_size, dropout_p)
                for i in range(num_layers)
            ]
        )

    def forward(self, x):
        hidden_states = []
        for layer in self.layers:
            x, (h, c) = layer(x)
            hidden_states.append((h, c))
        return x, hidden_states

Out[9]:

Console

Input shape:         (2, 10, 64)
Output shape:        (2, 10, 128)
Number of layers:    3
Each hidden state:   (1, 2, 128)

Within-pass variation (should be ~0): 0.0308
Cross-pass variation (should be > 0): 0.0090

The within-pass variation is near zero because the same mask applies at every timestep, confirming that information can flow consistently through the recurrence. The cross-pass variation shows that different forward passes do get different masks. This provides the stochastic regularization.