Stochastic Depth from scratch: randomly drop whole residual layers to train deeper nets faster

#machinelearning #ai #deeplearning #neuralnetworks

A 100-layer ResNet is slow to train and its gradients still fade on the way to the earliest layers. Stochastic Depth (Huang et al., 2016) attacks both at once with a trick that sounds reckless: randomly drop whole residual blocks on each training step. A dropped block becomes a plain identity bypass, so the network is shorter that step — cheaper to compute, with fewer nonlinearities for the gradient to cross. It's dropout's drop-a-whole-layer cousin. I built a from-scratch demo of the block tower; here's the whole thing, one line at a time.

The identity skip is what makes dropping a layer legal

You can't delete an arbitrary layer from a plain net — the path breaks. Stochastic depth works because it drops residual blocks, which already carry an identity shortcut: y = x + F(x). Set F to nothing and you're left with y = x — a perfectly valid, still-connected, shape-preserving transformation. A dropped block isn't a hole; it's a well-defined pass-through. That is why this is specifically a ResNet-family technique.

function block(x, F) { return add(x, F(x)); }   // y = x + F(x)
// drop F entirely and you still have a valid net: y = x  (identity)

Survival probabilities decay linearly with depth

Not every block should be dropped equally. Early blocks learn low-level, widely reused features — dropping them is destructive. Deep blocks are more specialised and redundant, so they can go freely. So each block l gets a keep-chance p_l that decays linearly from ≈1 at the first block down to p_L (typically 0.5) at the last. One hyperparameter, p_L, controls the whole schedule.

function survival(l, L, pL) {
  return 1 - (l / L) * (1 - pL);   // l = 1..L
}
// pL = 0.5, L = 54  ->  p_1 = 0.99 ... p_54 = 0.50

Training forward — a Bernoulli coin per block

On each training step, draw an independent coin for every block. Heads (prob p_l): run it normally, x + F(x). Tails: skip the branch entirely and pass x straight through. The dropped branch does no compute and receives no gradient that step. In the original paper the surviving branch is not scaled during training — it's a hard 0/1 gate.

function trainBlock(x, F, p) {
  if (Math.random() < p) return add(x, F(x));  // survives: x + F(x)
  else                    return x;             // dropped: identity
}

Test time — keep every block, scale by survival

Randomness is unacceptable at inference; you want one deterministic answer. So at test time nothing is dropped: every block runs, but each branch is weighted by its survival probability, y = x + p_l·F(x). That makes each block's expected contribution match what training saw — exactly dropout's train-drop / test-scale bargain, moved up from neurons to whole layers.

function testBlock(x, F, p) {
  return add(x, scale(F(x), p));    // y = x + p * F(x)  (no dropping)
}

Expected depth — how deep the net really is

Because each block survives independently, the expected number of active blocks is just the sum of the survival probabilities, E[depth] = Σ p_l. With linear decay to 0.5 that's about 0.75·L — so a 54-block ResNet trains as a ~40-block one on average, and the famous 1202-layer version behaves like a ~900-layer one per step. That's why stochastic depth actually trains faster in wall-clock time, not only regularises.

function expectedDepth(L, pL) {
  let e = 0;
  for (let l = 1; l <= L; l++) e += survival(l, L, pL);
  return e;                         // ~0.75 * L active blocks on average
}

Why it works: an implicit ensemble + a gradient shortcut

Two effects give the payoff. (1) With L blocks there are 2^L possible keep/drop sub-networks, and each step optimises a different one; the final model behaves like an average over this huge family of shallower nets — strong regularisation, like an ensemble, which is what the demo's "sample ×25" histogram shows. (2) A shorter active path means backprop crosses fewer nonlinear branches, so gradients to early layers stay strong and training converges faster. Together they let you train nets far deeper than a plain residual net can — the paper trained a 1202-layer ResNet on CIFAR-10 that improved with stochastic depth where the plain version overfit.

You never hand-roll the coin. Torchvision ships StochasticDepth and timm calls it DropPath; you drop it inside each residual block with a per-block drop rate (1 − p_l, decayed with depth). It's a no-op at eval(), needs no architecture change, and is standard equipment in ResNets, EfficientNet, ConvNeXt, and essentially every Vision Transformer.

from torchvision.ops import StochasticDepth   # timm: from timm.models.layers import DropPath

class ResBlock(nn.Module):
    def __init__(self, ch, drop_prob):
        super().__init__()
        self.F = nn.Sequential(conv(ch), norm(ch), relu(), conv(ch))
        self.drop = StochasticDepth(drop_prob, mode="row")  # drop_prob = 1 - p_l
    def forward(self, x):
        return x + self.drop(self.F(x))       # train: drops branch; eval: scales

drop_rates = [i / (L - 1) * 0.5 for i in range(L)]   # 0.0 ... 0.5, deeper = higher