39. Dropout From Scratch (Train vs. Eval Mode)

Problem

Dropout randomly zeroes a fraction p of activations during training, forcing the network to not rely on any single unit, then does nothing at all at inference time. The subtlety is the scaling: "inverted dropout" (what every modern framework implements) rescales the surviving activations by 1 / (1 - p) during training, so that the expected value of the output matches the un-dropped input. That's what lets eval mode be a pure no-op copy, with zero extra scaling work needed at inference.

Implement dropout's train/eval branching, including the inverted scaling.

Source: src/39_dropout_from_scratch.py

def dropout(x: torch.Tensor, p: float, training: bool) -> torch.Tensor:
    ...

Examples:
>>> x = torch.ones(4)
>>> dropout(x, p=0.5, training=False)
tensor([1., 1., 1., 1.])

Step-by-Step Approach

1. Handle eval mode first: if training is False, return x completely unchanged — no masking, no scaling, no randomness at all.
2. Handle the degenerate p == 0 case in training mode the same way (a no-op), since there's nothing to drop.
3. In training mode, sample a keep-mask: draw a uniform random value per element and keep it if the draw is >= p (equivalently, drop with probability p). torch.rand_like(x) >= p is one clean way to express this.
4. Apply the mask by elementwise multiplication, zeroing the dropped positions.
5. Rescale the survivors by 1 / (1 - p) — this is "inverted" dropout, and it's what makes the output's expected value match the input's, regardless of p.
6. Validate statistically rather than value-for-value: on a large all-ones tensor, the fraction of exact zeros in the output should land close to p, and every surviving (nonzero) value should be exactly 1 / (1 - p), not some other scale.

The key insight is that the scaling has to happen at train time, not eval time — that's the "inverted" part, and it's the detail that lets eval-mode dropout collapse to literally returning the input, with no mode-dependent scaling logic needed downstream.

Reference solution

def dropout(x: torch.Tensor, p: float, training: bool) -> torch.Tensor:
    if not training or p == 0.0:
        # Eval mode (or p=0): dropout is a pure no-op. This is only correct
        # because training already applied the 1/(1-p) rescale -- "inverted"
        # dropout pushes all the scaling work into the training pass so
        # inference stays a free, unmodified forward.
        return x
    if p >= 1.0:
        return torch.zeros_like(x)

    # Sample a keep-mask: True (1.0) where the uniform draw clears the drop
    # threshold p. Using a plain uniform-vs-threshold comparison instead of
    # torch.bernoulli is equivalent and slightly cheaper.
    keep_mask = (torch.rand_like(x) >= p).float()
    # Inverted dropout: rescale surviving activations by 1/(1-p) so that
    # E[output] == input in expectation, matching the un-dropped signal
    # magnitude -- this is what lets eval mode skip scaling entirely.
    return x * keep_mask / (1.0 - p)

Key Functions & Tricks

• torch.rand_like(x) >= p — vectorized Bernoulli-equivalent mask via uniform threshold, cheaper than torch.bernoulli.
• keep_mask.float() — converts the boolean mask to a multipliable float tensor.
• x * keep_mask / (1.0 - p) — the inverted-dropout rescale, applied at train time only.
• torch.zeros_like(x) — the degenerate p >= 1 edge case, guarded separately to avoid a division by zero.
• Statistical assertion (abs(zero_frac - p) < tol) instead of an exact value match — the right test strategy for anything involving genuine randomness.

How to Recognize This Pattern