Dropout Masks, Train-Time Scaling, and Inference Behaviour

Dropout is a single elementwise multiplication per step, and the interesting content is entirely in the distribution of the multiplier and in the fact that the layer behaves differently at training and at inference time. The multiplier is a Bernoulli random variable drawn independently for each element of the activation tensor, including along the batch axis, so two samples in the same batch see different masks. Under a fixed mask the activation is linear in the incoming value, which means the backward pass through it is the same multiplication, and a unit whose mask entry is zero receives exactly zero gradient for that step.

What the mask does to the expectation

Write the mask as $m \in {0,1}$ with $\Pr[m = 1] = 1 - p$, where $p$ is the drop probability, and the masked activation as $\hat{a} = m \odot a$ without compensation. Then $\mathbb{E}[\hat{a}] = (1-p),a$, since the mask is independent of the activation. The expected magnitude of every downstream pre-activation is therefore multiplied by $(1-p)$, and this bias has to be corrected somewhere, because the layers below were not trained for it and a test-time prediction would be systematically smaller than what the parameters expect.

There are two places to put the correction. Multiplying activations, or equivalently weights, by $(1-p)$ at inference time restores the expectation exactly, and is the convention in the original formulation. Dividing by $(1-p)$ during training instead gives $\mathbb{E}[\hat{a}/(1-p)] = a$, and then inference needs no correction at all: the evaluation path is the identity elementwise function, and the drop probability does not affect the shape of the inference graph. Implementations favor the second form because it keeps the inference path free of a scale factor whose value would otherwise have to be carried along with the parameters. The factor is large when $p$ is large: $p = 0.5$ doubles the surviving activations, $p = 0.9$ multiplies them by ten, so a unit whose mask survives contributes a pre-activation an order of magnitude above what the same unit produces at inference time if the scaling is applied in the wrong place or twice.

The backward pass through the same mask

Because $\hat{a} = m \odot a / (1-p)$ is affine in $a$ with coefficient $m/(1-p)$, the reverse sweep multiplies the incoming cotangent elementwise by the same mask scaled by the same factor. Dropped elements contribute exactly zero, and surviving elements contribute a gradient amplified by $1/(1-p)$ relative to the unmasked case. The amplification is what makes the expected update to a given weight the same as it would be without dropout, up to the correlation between the mask and the activation, while the variance of the update is larger.

Because the layer below is affine, dividing the surviving activations by $1-p$ is the same operation as multiplying the weights of that layer by $1-p$ and leaving the activations alone: the scale factor commutes with the matrix product. Folding the factor into the parameters applies one multiply per stored weight instead of one per activation element, and the activation tensor is usually the larger of the two, so the folded form saves work when the layer is wide. It also makes the check on the forward pass identical in both modes, since the factor disappears into the weights rather than into the data flowing through them. Either formulation leaves the inference path free of the factor, which is the property the convention exists to provide.

A unit that is dropped in a given step receives no update to its incoming weights in that step; with $p = 0.5$ it receives an update on about half the steps. That is the mechanism behind the ensemble argument made for dropout: each step trains a different subnetwork of the $2^{k}$ subnetworks that a layer of $k$ units admits, and the shared weights end up trained by an implicit collection of those subnetworks. The argument is an approximation, since the subnetworks are not independent, but it explains why the method changes the effective capacity rather than merely adding noise to the loss.

The mask varies by tensor element, which makes the meaning of “unit” depend on the tensor rank. An elementwise mask over a feature map drops individual spatial positions of individual channels for individual samples. Dimensional variants draw one mask entry per channel and apply it across the spatial extent and the batch, so an entire feature channel is removed for the whole batch; that is the version that produces a genuine subnetwork of channels and is the one that makes sense for convolutional feature maps, where a single spatial position carries little information on its own.

import numpy as np

def dropout_train(a, p, rng):
    keep = 1.0 - p
    mask = (rng.random(a.shape) < keep).astype(a.dtype)
    return a * mask / keep, mask      # scale here so inference is the identity

def dropout_backward(dA, mask, p):
    return dA * mask / (1.0 - p)

def dropout_eval(a, p):
    return a                          # no mask, no scaling

rng = np.random.default_rng(0)
a = np.ones((4, 6))
out, m = dropout_train(a, 0.5, rng)
assert np.allclose(out, 2.0 * m)      # survivors are doubled

Train and evaluation mode are different functions

A model that contains dropout computes two different functions depending on a mode flag, exactly as a normalization layer with batch statistics does. Because the masked training path is stochastic and the evaluation path is deterministic, an evaluation performed while the flag still says training produces different numbers on every call, and those numbers are drawn from the training distribution rather than from the distribution the parameters were finally fit to. The failure is usually noticed as an accuracy figure that jumps around between identical runs.

Keeping the training behaviour at inference time on purpose is a different matter, and it has a use. Repeating the forward pass many times with independent masks and averaging the resulting distributions is an approximation to a posterior predictive distribution, and the spread of the predictions across the passes is used as an uncertainty estimate. The cost is one forward pass per sample of the mask, and the interpretation rests on assumptions about the prior induced by the mask distribution that do not hold for every architecture. This is also the reason a single stochastic evaluation should not be trusted as a measurement of anything.

Where the mask interacts with the layer below

Three interactions account for most of the surprises. The first is with any layer that stores statistics of its input. Masking with the $1/(1-p)$ scaling changes the second moment of the activation entering the next layer by a factor of $1/(1-p)$ relative to the unmasked case, so an operation that accumulates running estimates of its input’s mean and variance accumulates them for the masked distribution while inference presents the unmasked one. The correction applied to the mean happens to be exact in expectation; the correction to the variance is not, which is why the combination of dropout above and batch statistics below has to be reasoned about layer by layer.

The second is with finite-difference gradient checks. The mask is resampled on every forward call, so an analytic gradient and a numerical gradient computed from two calls use two different masks and their difference reflects the masks, not the parameters. The check is only meaningful with every stochastic layer disabled or with a fixed seed shared by both calls; the general problem is the subject of gradient checking.

The third is with the shape of the activations it feeds. The scaling factor multiplies the output of the preceding dense layer by a constant up to ten, so a downstream saturating activation can be pushed into its tail by dropout alone, and a downstream unbounded activation propagates the increase. The magnitudes seen during training are then not the magnitudes the parameters were trained for.

Degenerate settings

A drop probability of 0 makes the layer the identity in both modes, and the mask draw becomes pure overhead. A probability of 1 removes every activation, so the output is zero, the reciprocal scaling factor is undefined, and the gradient is zero for all inputs. Placing dropout on a skip connection rather than inside the branch attached to it removes the block input on some samples, which destroys the unconditional path through the network that makes the block safe to stack. Placing it inside the branch leaves the skip path intact and keeps the block’s output defined for every sample, which is the arrangement that behaves as intended when residual blocks and dropout are combined.