Course contents 05 / 33

05 / 33 · Foundations

Batch normalization, gradients and convolutions

The training machinery behind every neural generative model in the course — batch normalization, the vanishing and exploding gradient problem, and the shapes and parameters of a convolution layer.

On this page 7

The previous note fitted a model with a handful of parameters. From here on the models are neural networks. The autoencoders of the next note are built from blocks of convolution → batch normalization → activation, and every later model is trained by backpropagating gradients through many such layers. The evaluation metrics later in the course also run images through pretrained convolutional networks. This note collects that background: what batch normalization computes and how instance normalization differs, why gradients can vanish or explode, how a convolution layer is shaped, and the routine around a training run.

Batch normalization#

Batch normalization (Ioffe and Szegedy, 2015) standardizes the activations of a layer using statistics of the current mini-batch. Take one feature (in a convolutional network, one channel) and a mini-batch of nn values x1,…,xnx_1, \dots, x_n of it. The layer first computes the batch mean and batch variance

μB=1n∑i=1nxi,σB2=1n∑i=1n(xi−μB)2,\mu_B = \frac1n\sum_{i=1}^{n} x_i, \qquad \sigma_B^2 = \frac1n\sum_{i=1}^{n} \left(x_i - \mu_B\right)^2,

and then standardizes every value:

x^i=xi−μBσB2+ε.\hat x_i = \frac{x_i - \mu_B}{\sqrt{\sigma_B^2 + \varepsilon}} .

The small constant ε>0\varepsilon > 0 (PyTorch uses 10−510^{-5}) keeps the division safe when a batch has almost no spread. After this step the batch has mean 00 and variance σB2/(σB2+ε)≈1\sigma_B^2/(\sigma_B^2 + \varepsilon) \approx 1.

Figure

Batch normalization of one feature

Histogram of a batch before and after normalizationThree histograms on a shared horizontal axis. The top one is a right-skewed batch centred near 6.32. An arrow labelled subtract mu B moves it to the middle row, now centred at zero with the same width. Two arrows labelled divide by the square root of sigma B squared plus epsilon squeeze the plus and minus one standard deviation marks to plus and minus one, giving the bottom histogram: narrower, taller, still skewed, with mean 0 and variance 1. A dashed standard normal curve is drawn on the bottom row for comparison.batch xᵢμB = 6.32subtract μBxᵢ − μB−σBσBdivide by √(σB2 + ε)x̂imean 0.00, variance 1.00𝒩(0, 1)051015−11Histogram of a batch before and after normalizationThree histograms on a shared horizontal axis. The top one is a right-skewed batch centred near 6.32. An arrow labelled subtract mu B moves it to the middle row, now centred at zero with the same width. Two arrows labelled divide by the square root of sigma B squared plus epsilon squeeze the plus and minus one standard deviation marks to plus and minus one, giving the bottom histogram: narrower, taller, still skewed, with mean 0 and variance 1. A dashed standard normal curve is drawn on the bottom row for comparison.batch xᵢμB = 6.32subtract μBxᵢ − μB−σBσBdivide by √(σB2 + ε)x̂imean 0.00, variance 1.00𝒩(0, 1)051015−11
A skewed batch of n = 256 activations, drawn as x = 2 + 4·exp(0.45 z) with z ~ 𝒩(0, 1) and binned in steps of 0.6, with batch mean 6.32 and batch standard deviation 2.06. Subtracting the mean slides the histogram to zero; dividing by √(σ_B² + ε) squeezes it to unit variance. Areas are preserved, so the narrower histogram is taller. The shape does not change: the batch stays skewed and is not turned into the dashed standard normal. Here ε = 10⁻⁵.

A small worked example. For the batch x=(1,2,3,6)x = (1, 2, 3, 6) with ε=10−5\varepsilon = 10^{-5}:

firstsecondthirdfourth
xix_i11223366
xi−μBx_i - \mu_B−2-2−1-10033
x^i\hat x_i−1.069-1.069−0.535-0.535001.6041.604

Here μB=12/4=3\mu_B = 12/4 = 3 and σB2=(4+1+0+9)/4=3.5\sigma_B^2 = (4 + 1 + 0 + 9)/4 = 3.5, so every centred value is divided by 3.5+10−5≈1.871\sqrt{3.5 + 10^{-5}} \approx 1.871. The variance divides by nn, not n−1n - 1: it is the plain variance of the batch.

The rest of the layer#

A full batch-norm layer has two learnable parameters per feature, a scale γ\gamma and a shift β\beta, and outputs

yi=γ x^i+β.y_i = \gamma\,\hat x_i + \beta .

The parameters γ\gamma and β\beta are fixed numbers, while μB\mu_B and σB2\sigma_B^2 change from batch to batch, so the identity cannot be matched batch by batch. It can be matched against the population statistics of the feature: with γ=Var⁡[x]+ε\gamma = \sqrt{\operatorname{Var}[x] + \varepsilon} and β=E[x]\beta = \mathbb E[x] the layer reproduces its input at inference, where running averages that estimate these population statistics (described below) replace μB\mu_B and σB2\sigma_B^2. In training it does so only approximately, up to the batch-to-batch fluctuation of μB\mu_B and σB2\sigma_B^2 (Ioffe and Szegedy, 2015, section 3). Standardization therefore does not restrict what the layer can express. It changes the parametrization, so that the scale of each feature is controlled by one parameter.

Training and inference behave differently. During training the layer uses the batch statistics μB,σB2\mu_B, \sigma_B^2 and also keeps running averages of them. In PyTorch, with the default momentum 0.10.1, these are μrun←0.9 μrun+0.1 μB\mu_{\text{run}} \leftarrow 0.9\,\mu_{\text{run}} + 0.1\,\mu_B and σrun2←0.9 σrun2+0.1 nn−1σB2\sigma^2_{\text{run}} \leftarrow 0.9\,\sigma^2_{\text{run}} + 0.1\,\tfrac{n}{n-1}\sigma_B^2; the running variance stores the unbiased estimate. At inference, when a “batch” may be a single image, the running averages are used instead, so the output of one example no longer depends on the other examples in the batch. In PyTorch the switch is model.eval(). Turning off gradients with torch.no_grad() does not switch batch norm to its running statistics.

In a convolutional network, BatchNorm2d normalizes each channel over all values of that channel in the batch: batch size × height × width of them. A convolution followed by batch norm is usually created with bias=False, because subtracting μB\mu_B cancels any constant bias, and β\beta takes over its role.

For a similar reason, batch norm does not see the scale of what comes before it. If a bias-free convolution with weights WW computes the channel values xix_i from an input uu, then scaling uu or WW by a factor r>0r > 0 turns every xix_i into r xir\,x_i, because the convolution is linear. Then μB\mu_B becomes r μBr\,\mu_B and σB2\sigma_B^2 becomes r2σB2r^2\sigma_B^2, and

rxi−rμBr2σB2+ε=xi−μBσB2+ε/r2.\frac{r x_i - r\mu_B}{\sqrt{r^2\sigma_B^2 + \varepsilon}} = \frac{x_i - \mu_B}{\sqrt{\sigma_B^2 + \varepsilon/r^2}} .

Up to the small change of ε\varepsilon to ε/r2\varepsilon/r^2, the output is unchanged (Ioffe and Szegedy, 2015, section 3.3). A bias-free convolution followed by batch norm is therefore insensitive to the overall scale of its input and of its weights.

Why does it help? The original paper motivated batch norm by reducing “internal covariate shift”, the drift of each layer’s input distribution during training. Santurkar et al. (2018), sections 2–3, argue that the gain does not come from reducing internal covariate shift. In their analysis batch norm makes the loss and its gradients smoother, which tolerates larger learning rates; why batch norm works is still debated. Either way, it keeps the scale of activations under control, which is exactly the issue of the next section.

Instance normalization#

Batch norm pools each channel over the whole batch, so the output for one image depends on the other images in it. Instance normalization (Ulyanov, Vedaldi and Lempitsky, 2016, Eq. 3) standardizes each image on its own. For image nn and channel cc of an H×WH \times W feature map xx,

μnc=1HW∑h,wxnchw,σnc2=1HW∑h,w(xnchw−μnc)2,ynchw=γc xnchw−μncσnc2+ε+βc,\begin{gathered} \mu_{nc} = \frac{1}{HW}\sum_{h,w} x_{nchw}, \qquad \sigma_{nc}^2 = \frac{1}{HW}\sum_{h,w} \left(x_{nchw} - \mu_{nc}\right)^2, \\ y_{nchw} = \gamma_c\,\frac{x_{nchw} - \mu_{nc}}{\sqrt{\sigma_{nc}^2 + \varepsilon}} + \beta_c , \end{gathered}

with a learnable scale γc\gamma_c and shift βc\beta_c per channel, as in batch norm (PyTorch’s InstanceNorm2d includes them only with affine=True). The layers differ only in which values share a mean and a variance:

LayerStatistics computed overOne mean and variance per
batch normthe NN images and all H×WH \times W positionschannel
instance normthe H×WH \times W positions of one imageimage and channel
layer norm, image (CNN)all channels and positions of one imageimage
layer norm, Transformerthe dd features of one tokentoken

Layer norm normalizes over whatever trailing axes it is configured with (normalized_shape in PyTorch). The CNN row is the configuration that pools a whole image; in a Transformer the same layer normalizes the feature vector of each token separately.

If the four values (1,2,3,6)(1, 2, 3, 6) of the worked example above are the 2×22 \times 2 map of one channel of one image, instance norm (with γc=1\gamma_c = 1, βc=0\beta_c = 0) turns them into (−1.069,−0.535,0,1.604)(-1.069, -0.535, 0, 1.604), whatever the other images in the batch contain. This has three consequences:

  • No dependence on the batch. Training and inference compute the same function, no running averages are needed, and a batch of a single image works.
  • Per-image contrast is removed. The mean and the spread of every feature map of every image are discarded; γc\gamma_c and βc\beta_c are the same for all images and cannot restore them. This is what fast style transfer needs: the stylized image should not depend on the contrast of the content image (Ulyanov et al., section 2). Where the overall brightness or contrast of an image carries information, the layer throws it away.
  • It needs spatial extent. On a 1×11 \times 1 map each channel has a single value, its variance is zero, and the output is βc\beta_c whatever the input.

Three uses of instance normalization:

  • Style transfer. Replacing batch norm by instance norm in a feed-forward stylization network, and keeping it at test time, improved the stylized images markedly (Ulyanov et al., section 3).
  • CycleGAN. Both the generators and the discriminators use it, and training runs with batches of one image (Zhu et al., 2017, section 4 and appendix 7.2; see the networks of CycleGAN).
  • The critic of WGAN-GP. The gradient penalty constrains the gradient of the critic’s output for each input separately, which is meaningful only if each output depends on its own input alone. Batch norm breaks this, so WGAN-GP omits it from the critic. Any normalization without interactions between examples is allowed; Gulrajani et al. (2017), section 4, recommend layer normalization, and instance normalization qualifies as well.

Vanishing and exploding gradients#

We train by computing the gradient of the loss with respect to every weight (backpropagation) and stepping against it. By the chain rule, the gradient that reaches an early layer is a product of one factor per layer, or per time step, that it passed through. With one unit per layer, hi=φ(ai)h_i = \varphi(a_i) with pre-activation ai=wihi−1+bia_i = w_i h_{i-1} + b_i and activation function φ\varphi, each factor is

∂hi∂hi−1=wi φ′(ai),\frac{\partial h_i}{\partial h_{i-1}} = w_i\,\varphi'(a_i),

a weight times the slope of the activation. For layers with many units the factors are Jacobian matrices, and what matters is their size: the product of many Jacobians tends to grow or shrink geometrically with depth, at a rate set by their largest singular values (Pascanu, Mikolov and Bengio, 2013, section 2.1; Goodfellow, Bengio and Courville, 2016, sections 8.2.5 and 10.7). So if the per-layer factors are mostly

  • larger than 11 in size, the product grows geometrically: exploding gradients;
  • smaller than 11 in size, the product shrinks geometrically: vanishing gradients.

The further the factors are from 11, the faster this happens, but the departure need not be large (see the table below). The problem was first analysed for recurrent networks (Hochreiter, 1991; Bengio, Simard and Frasconi, 1994), where it is easiest to see.

Example: an RNN#

A recurrent network applies the same weight at every time step. Take a single recurrent unit with a ReLU, ht=ReLU⁡(w ht−1+u xt+b)h_t = \operatorname{ReLU}(w\,h_{t-1} + u\,x_t + b), and recurrent weight w=2w = 2.

Figure

Unrolled RNN with recurrent weight w = 2

Unrolled recurrent network with exploding factorsThree time steps stacked vertically. In each, an input box feeds a recurrent cell drawn with a ReLU icon; the first two cells feed dashed output boxes and the third a solid output box. Curved red arrows carry the hidden state from each cell to the next, each labelled times 2; a dashed red squiggle continues below the third cell, also labelled times 2 and so on, exploding. A note explains that after tau steps the factor is w to the power tau.inputw = 2outputx1x2x3ReLU cell×2×2×2 …explodinghidden state hpasses throughthe same wafter τ steps:×wτUnrolled recurrent network with exploding factorsThree time steps stacked vertically. In each, an input box feeds a recurrent cell drawn with a ReLU icon; the first two cells feed dashed output boxes and the third a solid output box. Curved red arrows carry the hidden state from each cell to the next, each labelled times 2; a dashed red squiggle continues below the third cell, also labelled times 2 and so on, exploding. A note explains that after tau steps the factor is w to the power tau.inputw = 2outputx1x2x3ReLU cell×2×2×2 …explodingRed arrows: the hidden state h passesthrough the same w at every step.After τ steps the factor is wτ.
The same weight w = 2 is applied at every time step, so every red recurrent arrow multiplies by 2. The sensitivity of a later hidden state to an earlier one, and with it the backpropagated gradient, is therefore 2 after one step, 4 after two, 8 after three, and w raised to the power τ after τ steps. Intermediate outputs are dashed because only the last one is used here.

Every step along the recurrent path multiplies by the same w=2w = 2, so the sensitivity of a later state to an earlier one is 22 after one step, 44 after two, 88 after three, and so on. Formally, let the loss be computed only from the last state hk+τh_{k+\tau}, τ\tau steps after hkh_k (as in the unrolled RNN above, where only the last output is used). While the ReLU is in its active region (ReLU⁡′=1\operatorname{ReLU}' = 1), each per-step factor w φ′(at)w\,\varphi'(a_t) equals ww, and the chain rule gives

∂ht∂ht−1=wfor every t,∂hk+τ∂hk=∏s=k+1k+τ∂hs∂hs−1=wτ,∂L∂hk=∂L∂hk+τ wτ.\begin{aligned} \frac{\partial h_t}{\partial h_{t-1}} &= w \quad \text{for every } t, \\ \frac{\partial h_{k+\tau}}{\partial h_k} &= \prod_{s=k+1}^{k+\tau} \frac{\partial h_s}{\partial h_{s-1}} = w^{\tau}, \\ \frac{\partial L}{\partial h_k} &= \frac{\partial L}{\partial h_{k+\tau}}\, w^{\tau} . \end{aligned}

The gradient sent back τ\tau steps carries the factor wτw^\tau:

weightafter 3 stepsafter 10 stepsafter 50 steps
w=2w = 28810241024≈1.1⋅1015\approx 1.1\cdot 10^{15}
w=1.1w = 1.11.331.332.592.59117117
w=1w = 1111111
w=0.9w = 0.90.730.730.350.350.00520.0052
w=0.5w = 0.50.1250.125≈9.8⋅10−4\approx 9.8\cdot 10^{-4}≈8.9⋅10−16\approx 8.9\cdot 10^{-16}

Even a weight only 10 % away from 11 changes the gradient by two orders of magnitude over 50 steps. Nothing in the table is specific to recurrence: a stack of τ\tau layers with the same factor ww each gives the same wτw^\tau, and that depth-wise reading is all the autoencoders of the next notes need. Recurrent units and their gates are developed in the NLP course (recurrent networks, LSTM and GRU).

Figure

Repeated multiplication by w

Gradient factor w to the power τ on a log scaleThree straight lines on a logarithmic vertical axis from 1/1024 to 1024, against time steps 0 to 10. The line for w equals 2 rises to 1024 through the shaded exploding region; the line for w equals 1 stays at 1; the line for w equals 0.5 falls to 1/1024 through the shaded vanishing region. Points at τ equal 1, 2 and 3 on the rising line are labelled 2, 4 and 8.10243211/321/10240246810time steps τfactor wτ (log scale)exploding: wτ → ∞vanishing: wτ → 0w = 2w = 1w = 0.5248Gradient factor w to the power τ on a log scaleThree straight lines on a logarithmic vertical axis from 1/1024 to 1024, against time steps 0 to 10. The line for w equals 2 rises to 1024 through the shaded exploding region; the line for w equals 1 stays at 1; the line for w equals 0.5 falls to 1/1024 through the shaded vanishing region. Points at τ equal 1, 2 and 3 on the rising line are labelled 2, 4 and 8.10243211/321/10240246810time steps τfactor wτ (log scale)exploding: wτ → ∞vanishing: wτ → 0w = 2w = 1w = 0.5248
The factor w raised to the power τ after τ steps, on a logarithmic axis where powers of 2 are evenly spaced. With w = 2 it reaches 1024 after ten steps (exploding); with w = 0.5 it falls to 1/1024 (vanishing); only w = 1 keeps the scale. The first three red points, 2, 4 and 8, are the factors after one, two and three steps of the unrolled RNN above.

What this does to training#

A gradient-descent step is θ←θ−η ∇θL\theta \leftarrow \theta - \eta\,\nabla_\theta L, so the size of the step is proportional to the gradient.

  • Exploding: the gradient arrives multiplied by a huge factor, so the update takes a very large step in the loss landscape. Each update overshoots, and the loss jumps around or diverges.
  • Vanishing: the gradient that reaches the early layers of a deep stack carries one factor per layer. If the factors are small, the product goes to zero and the steps become tiny: the early weights barely move, and the training budget runs out long before the minimum is reached.
Figure

What a huge or tiny gradient does to the training steps

Overshooting steps versus steps that are too smallTwo copies of a parabola-shaped loss. In the first, labelled exploding, gradient-descent iterates start near the minimum and jump to alternating sides, each farther out, with red arcs between them. In the second, labelled vanishing, eleven blue points creep slowly down the right branch and stop at max steps, far from the minimum at zero.w ≫ 1: steps too largeminimumstartovershoots,divergesw → 0: steps too smallminimumstartstopped atmax_steps = 10Overshooting steps versus steps that are too smallTwo copies of a parabola-shaped loss. In the first, labelled exploding, gradient-descent iterates start near the minimum and jump to alternating sides, each farther out, with red arcs between them. In the second, labelled vanishing, eleven blue points creep slowly down the right branch and stop at max steps, far from the minimum at zero.w ≫ 1: steps too largeminimumstartovershoots,divergesw → 0: steps too smallminimumstartstopped atmax_steps = 10
Gradient descent on the toy loss L(θ) = θ²/2, whose gradient is θ. If the gradient arrives multiplied by a factor g picked up through the network, the update with learning rate η is θ ← θ − η g θ = (1 − η g) θ. The iterates overshoot the minimum when η g > 1 and move farther out when η g > 2. Red: η g = 2.3 from θ₀ = 1, so every step multiplies θ by −1.3 and the iterates bounce outwards. Blue: η g = 0.03 from θ₀ = 2.5, so ten steps move θ only to 1.84, and training ends at max_steps far from the minimum.

Common remedies#

The rest of the course relies on the standard fixes:

  • Initialization that keeps the scale of signals near 1: Xavier/Glorot initialization for saturating activations (Glorot and Bengio, 2010) and He initialization for ReLUs (He et al., 2015), described in the next section.
  • Normalization layers such as batch norm above.
  • Non-saturating activations instead of sigmoids. The sigmoid’s derivative is at most 14\tfrac14 (reached at 00), so the factor wi φ′(ai)w_i\,\varphi'(a_i) of a sigmoid layer is at most ∣wi∣/4|w_i|/4 in size: each sigmoid layer shrinks the gradient by at least a factor of 4 even when ∣wi∣=1|w_i| = 1. ReLU, ReLU⁡(x)=max⁡(0,x)\operatorname{ReLU}(x) = \max(0, x) (Nair and Hinton, 2010), has slope exactly 11 wherever the unit is active. LeakyReLU, max⁡(αx, x)\max(\alpha x,\ x) with a small slope α\alpha such as 0.20.2 (Maas, Hannun and Ng, 2013), also keeps slope α\alpha instead of 00 for negative inputs, so an inactive unit still passes some gradient.
  • Gradient clipping against explosions: rescale the gradient when its norm exceeds a threshold (Pascanu, Mikolov and Bengio, 2013, section 3.2).
  • Gated recurrent cells. In the LSTM (Hochreiter and Schmidhuber, 1997) the cell state is updated additively. In the original cell its factor per step is exactly 11. The cell used today adds a forget gate (Gers, Schmidhuber and Cummins, 2000), and the factor equals that gate: a number in (0,1)(0, 1) that the network can keep near 11. The GRU (Cho et al., 2014) works the same way: it interpolates between the old state and a candidate, so the direct path carries the factor of its update gate.
  • Residual connections (He et al., 2016a), xl+1=xl+F(xl)x_{l+1} = x_l + F(x_l). With a pure identity skip, as in pre-activation ResNets (He et al., 2016b, eq. 5), the Jacobian of a block is I+∂F/∂xlI + \partial F/\partial x_l, so the gradient always has a path whose factor is exactly 11.

A convolution layer#

How a 2-D convolution layer (nn.Conv2d) is parametrized and trained (Goodfellow, Bengio and Courville, 2016, chapter 9; Dumoulin and Visin, 2016):

  • The kernels are trained. All entries of the weight tensor and the bias are learned by gradient descent:

    W∈Routc×inc×kH×kW,b∈Routc.W \in \mathbb R^{\text{out}_c \times \text{in}_c \times k_H \times k_W}, \qquad b \in \mathbb R^{\text{out}_c} .
  • Notation. in[c]\text{in}[c] is input channel number cc, an Hin×WinH_{\text{in}} \times W_{\text{in}} plane (subscripted Win,WoutW_{\text{in}}, W_{\text{out}} are widths; WW, WjW_j, Wj,cW_{j,c} are always weights). Wj,cW_{j,c} is the kH×kWk_H \times k_W kernel that output channel jj applies to input channel cc. Collected over cc, Wj∈Rinc×kH×kWW_j \in \mathbb R^{\text{in}_c \times k_H \times k_W} is the filter for output channel jj.

  • Output channel jj, for j=1,…,outcj = 1, \dots, \text{out}_c, sums the per-channel results and adds one scalar bias:

    out[j]=∑c=1incWj,c⋆in[c]+bj.\text{out}[j] = \sum_{c=1}^{\text{in}_c} W_{j,c} \star \text{in}[c] + b_j .

    Each out[j]\text{out}[j] is an Hout×WoutH_{\text{out}} \times W_{\text{out}} plane. There is one bias per output channel, so the bias vector has outc\text{out}_c entries.

Figure

Shapes in a convolution layer: one output channel

Computing one output channel of a convolutionRows for input channels 1, 2 up to in_c. In each row a 6 by 6 input plane in[c] of size H_in by W_in is cross-correlated with a 3 by 3 kernel slice W j c, giving a 4 by 4 partial map. The three kernel slices are boxed together as filter W j of shape in_c by k by k. Arrows bring all partial maps into a summation node, a bias b j is added from above, and the result is output plane out[j] of size H_out by W_out, drawn in front of a stack of out_c planes, one per filter.in[c]Hin × WinWj,ck × kWj,cin[c]c = 1c = 2c = inc⋮⋮filter Wjinc × k × kΣover cbj+ biasout[j]Hout × Woutoutc planesrepeat with each of the outc filters W1, …, Woutc → outc planesComputing one output channel of a convolutionRows for input channels 1, 2 up to in_c. In each row a 6 by 6 input plane in[c] of size H_in by W_in is cross-correlated with a 3 by 3 kernel slice W j c, giving a 4 by 4 partial map. The three kernel slices are boxed together as filter W j of shape in_c by k by k. Arrows bring all partial maps into a summation node, a bias b j is added from above, and the result is output plane out[j] of size H_out by W_out, drawn in front of a stack of out_c planes, one per filter.in[c]Hin × WinWj,ck × kWj,cin[c]c = 1c = 2c = inc⋮⋮filter Wjinc × k × kΣbj+ biasout[j]Hout × Woutrepeat with each of the outc filters W1, …, Woutc→ outc output planes
Each output channel j has its own filter: one k × k kernel slice per input channel. The slice for channel c is slid over that channel, the results for all in_c input channels are summed, and the scalar bias of channel j is added. The amber window and cell show which inputs produce one output value. Repeating this with all out_c filters gives one output plane per filter, so the weight tensor has shape out_c × in_c × k × k. Drawn with a 6 × 6 input, k = 3, stride 1 and no padding, so the output planes are 4 × 4.

Written out entry by entry for stride ss and dilation dd, the operation ⋆\star is a sliding dot product. Deep-learning libraries compute a cross-correlation and do not flip the kernel, unlike the convolution of signal processing; since the kernel is learned, the difference does not matter.

out[j]y,x=bj+∑c=1inc∑m=0kH−1∑l=0kW−1Wj,c,m,l⋅ in[c] sy+dm,  sx+dl,\begin{aligned} \text{out}[j]_{y,x} = b_j &+ \sum_{c=1}^{\text{in}_c}\sum_{m=0}^{k_H-1}\sum_{l=0}^{k_W-1} W_{j,c,m,l} \\ &\quad\cdot\ \text{in}[c]_{\,s y + d m,\; s x + d l} , \end{aligned}

where y=0,…,Hout−1y = 0, \dots, H_{\text{out}} - 1, x=0,…,Wout−1x = 0, \dots, W_{\text{out}} - 1, and in[c]\text{in}[c] denotes the input after zero-padding by pp on each side, indexed from 00. The formula assumes one group, so every output channel sees every input channel; grouped convolutions restrict these connections.

  • Hyperparameters, not trained: stride, padding, kernel size and dilation are fixed when the layer is built. They determine the output size (PyTorch documentation):

    Hout=⌊NHs⌋+1,NH=Hin+2p−d (kH−1)−1\begin{aligned} H_{\text{out}} &= \left\lfloor \frac{N_H}{s} \right\rfloor + 1, \\ N_H &= H_{\text{in}} + 2p - d\,(k_H - 1) - 1 \end{aligned}

    for the height, where NHN_H is the padded height minus the extent d (kH−1)+1d\,(k_H - 1) + 1 of the dilated kernel. The same holds for WoutW_{\text{out}} with WinW_{\text{in}} and kWk_W. In the autoencoder of the next note each of the four encoder blocks has k=3, p=1, s=2, d=1k = 3,\ p = 1,\ s = 2,\ d = 1, so Hout=⌊(Hin−1)/2⌋+1H_{\text{out}} = \lfloor (H_{\text{in}} - 1)/2 \rfloor + 1 halves even sizes: 64→32→16→8→464 \to 32 \to 16 \to 8 \to 4.

  • Initialization. The kernels start random. With fanin=inc kHkW\text{fan}_{\text{in}} = \text{in}_c\, k_H k_W and fanout=outc kHkW\text{fan}_{\text{out}} = \text{out}_c\, k_H k_W, Xavier/Glorot initialization draws weights with variance 2/(fanin+fanout)2/(\text{fan}_{\text{in}} + \text{fan}_{\text{out}}) (Glorot and Bengio, 2010), and He initialization with variance 2/fanin2/\text{fan}_{\text{in}} for layers followed by a ReLU (He et al., 2015, section 2.2). Both are designed for the goal of the previous section: keep signals and gradients at a constant scale from layer to layer. PyTorch’s default for Conv2d is a simpler fan-in rule, a uniform distribution on ±1/fanin\pm 1/\sqrt{\text{fan}_{\text{in}}} (implemented as kaiming_uniform_ with a=5a = \sqrt5). A uniform distribution on ±b\pm b has variance b2/3b^2/3, so this default has variance 1/(3 fanin)1/(3\,\text{fan}_{\text{in}}), six times smaller than the He variance. Without normalization, the signal variance of a deep ReLU stack initialized this way shrinks by about a factor of 6 per layer. In the networks of the next notes, batch norm after the convolutions restores the scale.

  • What the weights learn. Kernels in the first layers typically become small detectors for edges, colour blobs and simple textures. Deeper layers combine them into parts and objects over larger receptive fields (Zeiler and Fergus, 2014, section 4 and figure 2). This is why pretrained convolutional networks are useful as feature extractors for the evaluation metrics of the course.

For scale: a layer with inc=16\text{in}_c = 16, outc=32\text{out}_c = 32 and 3×33 \times 3 kernels has 32⋅16⋅3⋅3=460832 \cdot 16 \cdot 3 \cdot 3 = 4608 weights plus 3232 biases.

A training run in practice#

Around the model and its loss, every training run in the course follows the same routine.

  • Optimizer. Most models use Adam (Kingma and Ba, 2015), which divides each coordinate’s step by a running estimate of its gradient’s magnitude (see the fitting note). A common variant is AdamW (Loshchilov and Hutter, 2019, section 2). An L2 penalty λ2∥θ∥2\tfrac\lambda2\|\theta\|^2 added to the loss enters Adam through the gradient and is divided by the same per-coordinate scale, so weights with historically large gradients are regularized less than intended. AdamW instead shrinks the weights directly, outside the adaptive step:

    θ←θ−ηλ θ−η m^v^+ϵ,\theta \leftarrow \theta - \eta\lambda\,\theta - \eta\,\frac{\hat m}{\sqrt{\hat v} + \epsilon},

    where m^\hat m and v^\hat v are Adam’s running estimates of the gradient’s mean and squared magnitude. For plain SGD, where m^/(v^+ϵ)\hat m/(\sqrt{\hat v} + \epsilon) is replaced by the gradient itself, the two forms give the same update; for Adam they differ. PyTorch’s AdamW defaults to η=10−3\eta = 10^{-3} and λ=10−2\lambda = 10^{-2}.

  • Validation. After each epoch the loss is also computed on held-out data, with the model in evaluation mode (model.eval(), which switches batch norm to its running averages and turns dropout off) and without gradient tracking (torch.no_grad(), which saves memory and time and changes nothing else). Plotting both losses per epoch shows overfitting as a validation curve that levels off or rises while the training curve keeps falling.

  • Checkpoints. The trained parameters (the model’s state_dict) are saved to disk, so that the model can be reloaded and evaluated later without training it again.

  • Reproducibility. Every random number generator involved is seeded: Python’s random, NumPy, and PyTorch, whose torch.manual_seed covers the CPU and all GPUs. On GPUs two more sources of variation remain. cuDNN can time several convolution algorithms and keep the fastest, which may differ between runs, and some algorithms are not deterministic; torch.backends.cudnn.benchmark = False and torch.backends.cudnn.deterministic = True turn both off, at some cost in speed. Even then, results are reproducible only on the same hardware and software versions (PyTorch documentation, “Reproducibility”).

Notebook#

Several notebooks of the HSE-LAMBDA Deep Generative Models repository (pinned commit; cell numbers count every cell, markdown included, from 1) use these pieces.

  • The autoencoder notebook, cell 11, defines Block: Conv2d(..., bias=False) → BatchNorm2d → LeakyReLU(0.2), kernel 3, padding (kernel - 1) // 2 and stride 2 in the encoder. Four blocks map a 1×64×641 \times 64 \times 64 MNIST image to 32×4×432 \times 4 \times 4, and a final stride-1 3×33\times 3 convolution keeps that size, giving a 512-dimensional latent. Cells 17 and 22 switch explicitly between ae.train() and ae.eval(), so batch norm uses batch statistics while training and running averages while reconstructing test images.
  • The VAE notebook, cells 10 and 12, orders its blocks as convolution → LeakyReLU → BatchNorm. Its generation code in cells 18 and 23 wraps sampling in torch.no_grad() without calling eval(), so batch norm still normalizes with the statistics of the generated batch and keeps updating its running buffers. This is the train/inference distinction from above in practice. The VAE notebook has no saved outputs, so it documents procedure only.
  • The WGAN notebook, cell 12, normalizes the critic’s second to fourth convolutions with InstanceNorm2d(..., affine=True) (on 16×1616 \times 16, 8×88 \times 8 and 4×44 \times 4 maps for 64×6464 \times 64 images), and cell 23 does the same in the blocks of its spectrally normalized critic. The generators keep batch norm.
  • Homework 1 follows the routine of the last section: set_seed in cell 3 sets all the seeds and both cuDNN flags, cells 23–31 train an autoencoder with AdamW at its defaults for 30 epochs and plot training and validation MSE per epoch, with the validation pass in eval() mode under torch.no_grad(), and the training function is to load saved weights when they exist. The notebook has no saved outputs.

Next#

With these pieces we can build the first neural model of the course. The next note stacks convolution → batch norm → LeakyReLU blocks into an encoder and a decoder and trains them to reconstruct MNIST digits through a 512-dimensional code. The questions it raises are what such a code captures, and whether new images can be generated from it.

References#