Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Dropout › Implementation: A Complete Dropout Layer

Michael BrenndoerferApril 27, 20252 min read

Language AI HandbookAll chapters

Part 10 of 67Neural Network Foundations

Chapter 72Dropout › Implementation: A Complete Dropout Layer

Reading section 15 of 19 · Implementation: A Complete Dropout Layer · 2 minSection 15 of 19 · 2 min

Section 15 of 19 · 2 min read

Dropout › Implementation: A Complete Dropout LayerOne section · Sections in this chapter15 / 19 ↓

Sections in this chapter19

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 implement a complete dropout layer that supports standard mode, spatial dropout mode, and MC Dropout mode:

In[24]:

Code

import torch
import torch.nn as nn


class FlexibleDropout(nn.Module):
    """
    Flexible dropout implementation supporting:
    - Standard (element-wise) dropout
    - Spatial dropout (channel-consistent across sequence length)
    - MC Dropout (dropout at inference time for uncertainty)
    """

    def __init__(self, p=0.5, spatial=False, mc_dropout=False):
        super().__init__()
        if not 0.0 <= p < 1.0:
            raise ValueError(f"Dropout probability must be in [0, 1), got {p}")
        self.p = p
        self.spatial = spatial
        self.mc_dropout = mc_dropout  # If True, dropout is always active

    def forward(self, x):
        active = self.training or self.mc_dropout
        if not active or self.p == 0.0:
            return x

        keep_prob = 1.0 - self.p

        if self.spatial and x.dim() == 3:
            # Spatial dropout: (batch, time, features) -> mask over (batch, 1, features)
            batch, time, features = x.shape
            mask = torch.bernoulli(
                torch.full(
                    (batch, 1, features),
                    keep_prob,
                    device=x.device,
                    dtype=x.dtype,
                )
            )
            mask = mask.expand(batch, time, features)
        else:
            # Standard dropout: independent mask per element
            mask = torch.bernoulli(torch.full_like(x, keep_prob))

        return (x * mask) / keep_prob

    def extra_repr(self):
        return (
            f"p={self.p}, spatial={self.spatial}, mc_dropout={self.mc_dropout}"
        )

Out[25]:

Console

Standard dropout (batch 0, first 3 features across 4 time steps):
[[0. 0. 2.]
 [2. 0. 0.]
 [0. 0. 0.]
 [2. 0. 2.]]

Spatial dropout (batch 0, first 3 features across 4 time steps):
[[0. 2. 2.]
 [0. 2. 2.]
 [0. 2. 2.]
 [0. 2. 2.]]

Note: spatial dropout drops identical features across all time steps.
Standard dropout has a unique pattern at each time step.