← All Problems

48. Dropout From Scratch

General Medium General Cross-Lab PyTorch Fundamentals
Grounding: General pattern common across ML-research-lab technical interviews. A second-hand Blind report describes an Anthropic candidate being asked to "implement a typical torch training loop for an MLP (forward pass, loss, backprop, optimizer) in numpy" -- i.e. reimplementing standard nn building blocks from first principles is a reported theme, though that specific report does not name dropout, so this exercise is written as a general from-scratch nn-primitive task in that same spirit.

Problem

nn.Dropout looks trivial from the outside but has two details that are easy to get subtly wrong when reimplementing it: it must be a no-op at eval time (not just "drop less"), and during training it has to rescale the surviving activations by 1/(1-p) -- "inverted dropout" -- so that the expected activation magnitude is the same whether or not dropout is active, meaning the eval-time forward pass needs no separate rescaling.

Implement inverted dropout: if not training (or p == 0), return x unchanged; otherwise draw one uniform sample per element, keep elements whose sample is >= p, zero the rest, and rescale the survivors by 1/(1-p).

Source: src/48_dropout_from_scratch.py

def dropout(
    x: torch.Tensor, p: float, training: bool = True,
    generator: torch.Generator | None = None,
) -> torch.Tensor: ...

>>> x = torch.ones(4)
>>> dropout(x, p=0.0, training=True).tolist()
[1.0, 1.0, 1.0, 1.0]
>>> dropout(x, p=0.5, training=False).tolist()
[1.0, 1.0, 1.0, 1.0]

Step-by-Step Approach

  1. Check the no-op conditions first: if training is False, or p == 0, return x unchanged -- this must be a hard no-op, not "apply a smaller drop rate."
  2. Draw exactly one uniform(0, 1) sample per element with torch.rand(x.shape, generator=generator) -- passing the generator through is what makes results reproducible for a given seed.
  3. Build the keep mask by comparing the samples to p: keep where the sample is >= p, so that in expectation exactly (1 - p) fraction of elements survive.
  4. Cast the boolean mask to x's dtype so it multiplies cleanly.
  5. Multiply x by the mask to zero out dropped elements.
  6. Divide the result by (1 - p) -- this is the "inverted" part of inverted dropout, applied at train time so eval time needs zero special handling.

The key insight is that the 1/(1-p) rescale has to happen during training, not at eval time -- that's what lets model.eval() simply skip dropout entirely rather than needing to know what drop probability was used during training to compensate for it.

Reference solution

def dropout(
    x: torch.Tensor,
    p: float,
    training: bool = True,
    generator: torch.Generator | None = None,
) -> torch.Tensor:
    # Eval mode (or p == 0, meaning "keep everything") is a hard no-op --
    # NOT a smaller drop probability. Getting this branch wrong is the
    # single most common inverted-dropout bug (rescaling at eval time too).
    if not training or p == 0.0:
        return x.clone()

    # One uniform sample per element; keep where sample >= p so that,
    # in expectation, exactly (1 - p) of elements survive.
    sample = torch.rand(x.shape, generator=generator)
    mask = (sample >= p).to(x.dtype)

    # Inverted dropout: rescale surviving activations by 1/(1-p) HERE, at
    # train time, so eval-time forward passes need no compensating scale.
    return x * mask / (1.0 - p)

Key Functions & Tricks

  • torch.rand(shape, generator=generator) — draws uniform(0, 1) samples; threading a Generator through (instead of relying on global RNG state) makes the function's randomness independently reproducible.
  • (sample >= p).to(x.dtype) — converts a boolean mask to a numeric one so it can be multiplied directly against x.
  • x.clone() — returned on the no-op paths so the caller never accidentally gets back a tensor aliased to the input.
  • "Inverted" scaling by 1/(1-p) — the technique that keeps E[dropout(x)] == x and removes the need for any eval-time rescale.

How to Recognize This Pattern

Signal words: "implement dropout from scratch," "inverted dropout," "why does dropout need to be a no-op at eval time." The tell is any "reimplement a standard nn building block manually" prompt, which shows up across labs as a way to check that a candidate understands what a familiar API is actually doing underneath, not just that they can call it. Common variations: DropConnect (dropping individual weights instead of activations); structured/channel dropout (zeroing entire feature-map channels together, common in conv nets); or implementing it as an nn.Module with a training attribute set automatically by model.train()/model.eval() instead of an explicit function argument. A common pitfall is applying the 1/(1-p) rescale at eval time instead of train time (the exact inverse of correct inverted dropout), which silently shrinks activations at inference instead of leaving them untouched.

← Previous Next →