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 values of it. The layer first computes the batch mean and batch variance
and then standardizes every value:
The small constant (PyTorch uses ) keeps the division safe when a batch has almost no spread. After this step the batch has mean and variance .
Batch normalization of one feature
A small worked example. For the batch with :
| first | second | third | fourth | |
|---|---|---|---|---|
Here and , so every centred value is divided by . The variance divides by , not : 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 and a shift , and outputs
The parameters and are fixed numbers, while and 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 and the layer reproduces its input at inference, where running averages that estimate these population statistics (described below) replace and . In training it does so only approximately, up to the batch-to-batch fluctuation of and (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 and also keeps running averages of them. In PyTorch, with the default momentum , these are and ; 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 cancels any constant bias, and 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 computes the channel values from an input , then scaling or by a factor turns every into , because the convolution is linear. Then becomes and becomes , and
Up to the small change of to , 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 and channel of an feature map ,
with a learnable scale and shift 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:
| Layer | Statistics computed over | One mean and variance per |
|---|---|---|
| batch norm | the images and all positions | channel |
| instance norm | the positions of one image | image and channel |
| layer norm, image (CNN) | all channels and positions of one image | image |
| layer norm, Transformer | the features of one token | token |
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 of the worked example above are the map of one channel of one image, instance norm (with , ) turns them into , 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; and 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 map each channel has a single value, its variance is zero, and the output is 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, with pre-activation and activation function , each factor is
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 in size, the product grows geometrically: exploding gradients;
- smaller than in size, the product shrinks geometrically: vanishing gradients.
The further the factors are from , 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, , and recurrent weight .
Unrolled RNN with recurrent weight w = 2
Every step along the recurrent path multiplies by the same , so the sensitivity of a later state to an earlier one is after one step, after two, after three, and so on. Formally, let the loss be computed only from the last state , steps after (as in the unrolled RNN above, where only the last output is used). While the ReLU is in its active region (), each per-step factor equals , and the chain rule gives
The gradient sent back steps carries the factor :
| weight | after 3 steps | after 10 steps | after 50 steps |
|---|---|---|---|
Even a weight only 10 % away from changes the gradient by two orders of magnitude over 50 steps. Nothing in the table is specific to recurrence: a stack of layers with the same factor each gives the same , 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).
Repeated multiplication by w
What this does to training#
A gradient-descent step is , 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.
What a huge or tiny gradient does to the training steps
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 (reached at ), so the factor of a sigmoid layer is at most in size: each sigmoid layer shrinks the gradient by at least a factor of 4 even when . ReLU, (Nair and Hinton, 2010), has slope exactly wherever the unit is active. LeakyReLU, with a small slope such as (Maas, Hannun and Ng, 2013), also keeps slope instead of 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 . The cell used today adds a forget gate (Gers, Schmidhuber and Cummins, 2000), and the factor equals that gate: a number in that the network can keep near . 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), . With a pure identity skip, as in pre-activation ResNets (He et al., 2016b, eq. 5), the Jacobian of a block is , so the gradient always has a path whose factor is exactly .
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:
-
Notation. is input channel number , an plane (subscripted are widths; , , are always weights). is the kernel that output channel applies to input channel . Collected over , is the filter for output channel .
-
Output channel , for , sums the per-channel results and adds one scalar bias:
Each is an plane. There is one bias per output channel, so the bias vector has entries.
Shapes in a convolution layer: one output channel
Written out entry by entry for stride and dilation , the operation 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.
where , , and denotes the input after zero-padding by on each side, indexed from . 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):
for the height, where is the padded height minus the extent of the dilated kernel. The same holds for with and . In the autoencoder of the next note each of the four encoder blocks has , so halves even sizes: .
-
Initialization. The kernels start random. With and , Xavier/Glorot initialization draws weights with variance (Glorot and Bengio, 2010), and He initialization with variance 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
Conv2dis a simpler fan-in rule, a uniform distribution on (implemented askaiming_uniform_with ). A uniform distribution on has variance , so this default has variance , 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 , and kernels has weights plus 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 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:
where and are Adam’s running estimates of the gradient’s mean and squared magnitude. For plain SGD, where is replaced by the gradient itself, the two forms give the same update; for Adam they differ. PyTorch’s
AdamWdefaults to and . -
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, whosetorch.manual_seedcovers 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 = Falseandtorch.backends.cudnn.deterministic = Trueturn 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) // 2and stride 2 in the encoder. Four blocks map a MNIST image to , and a final stride-1 convolution keeps that size, giving a 512-dimensional latent. Cells 17 and 22 switch explicitly betweenae.train()andae.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 callingeval(), 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 , and maps for 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_seedin cell 3 sets all the seeds and both cuDNN flags, cells 23–31 train an autoencoder withAdamWat its defaults for 30 epochs and plot training and validation MSE per epoch, with the validation pass ineval()mode undertorch.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#
- S. Ioffe, C. Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. ICML, 2015.
- S. Santurkar, D. Tsipras, A. Ilyas, A. Mądry. How does batch normalization help optimization? NeurIPS, 2018.
- D. Ulyanov, A. Vedaldi, V. Lempitsky. Instance normalization: the missing ingredient for fast stylization. arXiv:1607.08022, 2016. Section 2 and Eq. (3): instance normalization and the contrast argument; section 3: its effect on stylization networks.
- J. L. Ba, J. R. Kiros, G. E. Hinton. Layer normalization. arXiv:1607.06450, 2016.
- J.-Y. Zhu, T. Park, P. Isola, A. A. Efros. Unpaired image-to-image translation using cycle-consistent adversarial networks. ICCV, 2017. Section 4 and appendix 7.2: instance normalization in the generators and discriminators, batch size 1.
- I. Gulrajani, F. Ahmed, M. Arjovsky, V. Dumoulin, A. Courville. Improved training of Wasserstein GANs. NeurIPS, 2017. Section 4, “No critic batch normalization”.
- D. P. Kingma, J. Ba. Adam: A method for stochastic optimization. ICLR, 2015.
- I. Loshchilov, F. Hutter. Decoupled weight decay regularization. ICLR, 2019. Section 2: L2 regularization and weight decay differ for Adam; AdamW.
- S. Hochreiter. Untersuchungen zu dynamischen neuronalen Netzen. Diploma thesis, TU München, 1991.
- Y. Bengio, P. Simard, P. Frasconi. Learning long-term dependencies with gradient descent is difficult. IEEE Transactions on Neural Networks, 5(2), 1994.
- R. Pascanu, T. Mikolov, Y. Bengio. On the difficulty of training recurrent neural networks. ICML, 2013. Section 2.1 (conditions on the largest singular value) and section 3.2 (gradient-norm clipping).
- S. Hochreiter, J. Schmidhuber. Long short-term memory. Neural Computation, 9(8), 1997.
- F. A. Gers, J. Schmidhuber, F. Cummins. Learning to forget: Continual prediction with LSTM. Neural Computation, 12(10), 2000.
- K. Cho, B. van Merriënboer, C. Gulcehre, D. Bahdanau, F. Bougares, H. Schwenk, Y. Bengio. Learning phrase representations using RNN encoder–decoder for statistical machine translation. EMNLP, 2014.
- V. Nair, G. E. Hinton. Rectified linear units improve restricted Boltzmann machines. ICML, 2010.
- A. L. Maas, A. Y. Hannun, A. Y. Ng. Rectifier nonlinearities improve neural network acoustic models. ICML Workshop on Deep Learning for Audio, Speech and Language Processing, 2013.
- X. Glorot, Y. Bengio. Understanding the difficulty of training deep feedforward neural networks. AISTATS, 2010.
- K. He, X. Zhang, S. Ren, J. Sun. Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification. ICCV, 2015.
- K. He, X. Zhang, S. Ren, J. Sun. Deep residual learning for image recognition. CVPR, 2016 (cited as 2016a).
- K. He, X. Zhang, S. Ren, J. Sun. Identity mappings in deep residual networks. ECCV, 2016 (cited as 2016b).
- M. D. Zeiler, R. Fergus. Visualizing and understanding convolutional networks. ECCV, 2014.
- V. Dumoulin, F. Visin. A guide to convolution arithmetic for deep learning. arXiv:1603.07285, 2016.
- I. Goodfellow, Y. Bengio, A. Courville. Deep Learning. MIT Press, 2016. Sections 8.2.5 and 10.7 (long-term dependencies) and chapter 9 (convolutional networks).
- PyTorch documentation:
BatchNorm2d,InstanceNorm2d,Conv2d,AdamWand Reproducibility.