Stochastic Depth — Training a Deep Network That's Sometimes Shallow
Most regularization techniques you can name work at the level of activations or weights. Dropout zeroes out random neurons. Weight decay shrinks individual parameters. Label smoothing softens output distributions. These all share an assumption: the structure of the network is fixed, and we just perturb what flows through it.
Stochastic depth makes a stranger move. It perturbs the network itself. Each training step, entire layers vanish. The 50-layer ResNet is sometimes a 48-layer network, sometimes a 45-layer one — the depth is a random variable. By the time training finishes, the optimizer hasn't really trained one network. It has trained an ensemble of subnetworks that share the same weights but differ in which blocks they actually use. This turns out to be a remarkably effective regularizer, and it is the main thing keeping our ResNet-50 backbone from memorizing our small dataset.
1. The Setup — Why Regularize the Backbone at All?
Our backbone is a ResNet-50 with about 23 million parameters. Our dataset is roughly 5,000 annotated pages — handwriting, scanned documents, and form fields. The capacity-to-data ratio is wildly unfavorable. Without regularization, the backbone will quietly learn to identify training pages by their pixel-level quirks rather than the underlying letter shapes.
Dropout, the classical answer, doesn't translate cleanly to convolutional feature maps. Dropping random activations from a feature map disrupts spatial structure — the very thing convolutions are trying to preserve. Worse, the BatchNorm layers everywhere in ResNet already do half of dropout's job (they decorrelate features and add noise) and the two interact poorly. Empirically, dropout in CNN feature maps either does nothing or makes training slower.
2. The Mechanism — Skipping Entire Blocks
A standard residual bottleneck block computes:
The block adds whatever the convolutions inside \( F_l \) produce to its input. Stochastic depth turns this into a random choice:
When the block is "skipped," its convolutions don't run, no gradient is computed for them, and the input passes through untouched. This is not a no-op: the optimizer sees a different effective network on this batch than on the previous one, and it has to make the surviving blocks robust to the possibility that some of their downstream consumers will randomly disappear. The block can no longer rely on any specific neighbor always being there.
At inference time, all blocks run. The result is something like an ensemble: each individual block has been trained to be useful in many different subnetworks, and using all of them at once tends to give a strictly better final prediction than any of the random configurations seen during training.
3. The Linear Keep-Probability Schedule
The natural question is: how often should each block be skipped? We don't apply the same probability to every block. Instead, the keep probability decreases linearly with depth, so shallow blocks are almost always kept and deep blocks are skipped more often:
with \( p_{\min} = 0.8 \) and \( L = 16 \) (the total number of bottleneck blocks across ResNet-50's layer2, layer3, and layer4). Working out the schedule for every block:
| Block index l | Keep probability \( p_l^{\text{keep}} \) | Skip probability |
|---|---|---|
| 0 | 1.000 | 0.0% |
| 3 | 0.960 | 4.0% |
| 5 | 0.933 | 6.7% |
| 8 | 0.893 | 10.7% |
| 11 | 0.853 | 14.7% |
| 13 | 0.827 | 17.3% |
| 15 | 0.800 | 20.0% |
The first block is never skipped; the last block is skipped one batch in five. Why this asymmetry? Because shallow blocks and deep blocks do fundamentally different work. Shallow blocks compute low-level features — edges, strokes, oriented gradients. These are the foundation everything else is built on, and if you knock them out, the rest of the network has nothing to work with. Deep blocks compute higher-level, more specialized features — semantic features about what a word looks like, what shape a paragraph has. Those features are more redundant across blocks, more "swappable," and so the network can afford to lose them occasionally.
The schedule encodes a hard-won engineering principle: regularize where the model has capacity to spare, protect the foundations. Dropping a shallow block is structural damage; dropping a deep block is just noise. The linear schedule is a clean compromise between the two regimes.
4. One Subtle Implementation Detail — Downsample Blocks Always Run
The mechanism above quietly assumed that "skipping a block" means passing the input through unchanged. That only works if the input and output shapes match. In ResNet, however, the first block of each layer is a downsample block — it halves the spatial resolution and changes the channel count. The identity path through such a block is not actually identity; it's a 1×1 convolution that adjusts shape. Skipping the convolutional path of a downsample block would leave the spatial dimensions and channel counts mismatched between layers, which crashes the forward pass with a shape error.
The fix is simple but worth being explicit about: downsample blocks always run. Only non-downsample blocks are eligible for stochastic skipping. Our implementation detects this at module-construction time by checking whether the block has a non-trivial downsample submodule, and if so, marks it as never-skip:
This costs us nothing — downsample blocks are a minority, and the surrounding non-downsample blocks still get plenty of randomization.
5. The Picture — Two Forward Passes, Same Weights
The clearest way to internalize what stochastic depth does is to draw two random forward passes side by side. Same model, same weights, two different training steps:
Two crucial properties are visible in this picture. First, the two sub-networks share their weights — every block that is kept in both passes uses the exact same parameters. So the optimizer is not training two separate models; it is training one set of weights that has to work across many random sub-networks. Second, the path through skipped blocks is the identity shortcut. Information still flows; gradients still propagate. There is no broken network, just a temporarily shallower one.
6. Stochastic Depth vs Dropout — A Sharper Comparison
It is worth comparing this directly to dropout to see why stochastic depth is the right tool for ResNet-50.
| Aspect | Dropout | Stochastic Depth |
|---|---|---|
| Granularity | individual neurons | entire residual blocks |
| Effect on spatial structure | shreds it (random pixels disappear) | preserves it (whole block bypassed) |
| Where it lives best | fully connected and recurrent layers | residual CNNs and transformers |
| Compute during training | same as without dropout | faster — skipped blocks compute nothing |
| Effect on gradient flow | can hurt for very deep networks | improves it (shorter effective paths) |
| Interaction with BatchNorm | often counterproductive | orthogonal — both work fine together |
| Test-time behavior | scaled activations | all blocks active, no scaling needed |
The key row is the second: effect on spatial structure. A convolutional feature map encodes spatial relationships — pixel \( (i, j) \) is meaningfully near pixel \( (i+1, j) \). Dropout punches random holes through this structure. The convolutions in the next layer then have to learn around the holes, which mostly amounts to learning to be robust to noise rather than learning anything useful. Stochastic depth never touches the feature map values at all. When a block is skipped, the spatial structure is passed through perfectly intact — just unmodified by that particular block's transformation.
7. Why It Matters Here
Our model lives in exactly the regime stochastic depth was designed for. ResNet-50 with full residual structure, a small dataset that overfitting will eat alive, and a downstream pipeline (FPN, segmentation head, recognition head) that depends on the backbone producing useful features at all three pyramid levels — not just memorizing pages.
Stochastic depth doesn't fight overfitting alone. It works in concert with the rest of our regularization stack — RandomErase wipes out spatial patches at the input, geometric augmentation perturbs the pixel grid, label smoothing softens the target distribution, EMA smooths the weight trajectory, and dropout still lives inside the transformer decoder where it actually belongs. Each tool takes out a different failure mode. Stochastic depth's particular job is making sure no single block in the backbone becomes load-bearing — making sure the model has learned features, not a specific chain of computations.
Comments
Post a Comment