Medium Research

Stochastic Depth: Inference

A residual connection adds a branch's contribution to an unchanged input. Stochastic depth randomly skips that branch during training. This exercise implements only the deterministic inference rule, with the branch output already supplied.

Your task

Implement stochastic_depth(x, block_output, drop_prob, training).

x and block_output are finite floating-point tensors with shape (batch, d) and the same dtype and device.
block_output is the residual branch's contribution before adding x. It is not the result of the complete residual block.
drop_prob is a scalar in [0, 1]. The survival probability is 1 - drop_prob.
training is always False for this exercise. Random training behavior is outside the task; do not draw a random mask during inference.

Return a tensor with the same shape and dtype as x:

output = x + (1 - drop_prob) * block_output

Scale only the branch contribution, not x. Leave both inputs unchanged. Do not apply ReLU or another activation after the addition.

For x=[[1, -2]], block_output=[[8, -4]] and drop_prob=0.25, return [[7, -5]]. With drop_prob=0, return x + block_output. With drop_prob=1, return x, without division by zero.

Which scaling convention is this?

This uses the inference-time survival scaling described by Huang et al. (2016). For fixed x and branch output, if a training gate is 1 with probability 1 - drop_prob and otherwise 0, the expected pre-activation sum is x + (1 - drop_prob) * block_output.

The exercise stops at that sum; the paper's full block also applies ReLU. The expectation statement concerns the sum before an activation, not the output of a whole nonlinear network. Do not substitute inverted-dropout scaling, which uses a different training/evaluation convention.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Choose initial weights that keep signals at a useful scale

Initialization helps keep activations and gradients from shrinking or growing rapidly across layers. Under suitable independence and variance assumptions, summing fan_in weighted inputs motivates weight variance proportional to 1 / fan_in.

Xavier initialization considers both input and output widths. He initialization also accounts for how ReLU removes the negative part of its input. Follow the specified formula and distribution rather than using these names as complete instructions.

A standard deviation is not the same as a uniform distribution's bound. Uniform samples from [-a, a] have variance a²/3. To get standard deviation s, the bound must be sqrt(3) * s.

Drop activations during training

With keep probability q, inverted dropout multiplies each activation x by m/q, where m is a zero-or-one mask sampled with probability q of being one. For q > 0, the expected output is x. A particular draw can still be very different. At evaluation, inverted dropout leaves the input unchanged.

For x = 4 and q = 0.5, the training output is either 0 or 8, each with probability one half. Their average is 4. The endpoint q = 0 needs explicit handling because the formula divides by zero.

Drop residual branches with stochastic depth

Stochastic depth drops an entire residual branch, often sharing a mask across its feature dimensions. The shortcut remains present. Some conventions scale surviving branches during training; others scale the branch at evaluation. Do not combine the two conventions.

Check mask shape, training/evaluation mode, and the required random seed or key behavior. Reusing a JAX key for the same random operation repeats its draw. PyTorch and JAX need not produce the same samples from the same integer seed.

Detecting runtime…

Stochastic Depth: Inference

Your task

Which scaling convention is this?

Choose initial weights that keep signals at a useful scale

Drop activations during training

Drop residual branches with stochastic depth

Your code, saved here

Stuck?

PyTorch reference solution