Explained: Characterizing Signal Propagation to Close the Performance Gap in Unnormalized ResNets

Introduction - High level introduction to the paper's key ideas and contributions
Code Implementation - In depth understanding of the paper with code implementation in PyTorch (code referenced from timm)

💭: Throughout this report, you will find sidebars like this. These are my own personal comments that might help you along the way. These sidebars could also include some funny bits or failed experimentation stories.

💭: This report is long, longer than an average blog post, but it is complete. I have made a conscious effort to not leave any concept out of this report. IMHO, this report is best read in parts and might require re-reading for users who are being introduced to NF-ResNets for the first time.

💭: Note that this paper is built on past research, and therefore some prior knowledge about various concepts mentioned in the Prerequisites section below will really help. If you feel like this report on its own wasn't sufficient to help explain this research paper, then I am also open to hosting a paper reading session where we could discuss past research to help fill the gap. Let me know and I'll make it happen :)

💭: My personal aim is to create a report that distills the research paper into a simple consumable format. If some parts of the report don't make sense or are confusing, please feel free to provide constructive feedback towards the end of this report.

Prerequisites

Introduction to ResNets, skip connections & bottleneck layers by Jeremy Howard & Sylvain Gugger.
The Drawbacks of Batch Normalization and introduction to other Normalization techniques sections in Group Normalization blog post by Aman Arora
Weight Standardization (Paper Explained) by Yannic Kilcher

💭: I am bit biased on the second resource mentioned above. ;)

Introduction

BatchNorm has become a core computational primitive in deep learning, and it is used in almost all state-of-the-art image classifiers. A number of different benefits of BatchNorm have been identified. It smoothens the loss landscape, which allows training with larger learning rates, and the noise arising from the mini-batch estimates of the batch statistics introduces implicit regularization. However, BatchNorm also has many disadvantages. Its behavior is strongly dependent on the batchsize, performing poorly when the per device batch size is too small or too large, and it introduces a discrepancy between the behaviour of the model during training and at inference time. A number of alternative normalization layers have been proposed, but typically these alternatives generalize poorly or introduce their own drawbacks, such as added compute costs at inference. Another line of work has sought to eliminate layers which normalize hidden activations entirely.

❓: Why do we want to remove BatchNorm? This has been answered in Why do we need Normalizer-Free networks? What's wrong with BatchNorm? section of this report.

Signal Propagation Plots: The authors propose a simple set of visualizations which helps practitioners inspect signal propagation at initialization on the forward pass in deep residual networks.
Scaled Weight Standardization: The authors identify a key failure mode in past unnormalized ResNets with ReLU or Swish activations and Gaussian weights. Because the mean output of these non linearities is positive, the squared mean of the hidden activations on each channel grows rapidly as the network depth increases. To resolve this, the authors propose Scaled Weight Standardization which is an extension of Weight Standardization. Essentially, weight standardization normalizes weights in convolution layers, i.e., making the weights have zero mean and unit variance.
Comparable performance with BatchNorm Counterparts: The authors apply the normalization-free network structure in conjunction with Scaled Weight Standardization to ResNets on ImageNet, where they, for the first time attain performance which is comparable or better than batch-normalized ResNets on networks as deep as 288 layers.

❓: What does "mean output of these non linearities is positive" in the second bullet point above mean? Remember ReLU is nothing but a $max(0, x)$ operation given we have some input $x$. So the output of ReLU activation is always positive, thus, there is a mean shift towards $mean(x)>0$. To counter this, authors introduced Scaled Weight Standardization which is discussed later in Scaled Weight Standardization section of the report.

💭: Networks without BatchNorm until this point weren't able to achieve comparable SOTA performance, so you can imagine why this is pretty exciting! In fact, in their follow up paper, the authors even manage to attain a new SOTA! But, let's not jump ahead of ourselves for now..

Why do we need Normalizer-Free networks? What's wrong with BatchNorm?

A number of benefits of BatchNorm have been identified. It smoothens the loss landscape (Santurkar et al, 2018), which allows training with larger learning rates (Bjorck et al, 2018) and the noise arising from the minibatch estimates of the batch statistics introduces implicit regularization (Luo et al, 2019). It also mantains good signal propogation at initialization in deep residual networks with identity skip connections (De & Smith, 2020).

Its behavior is strongly dependent on the batchsize, performing poorly when the per device batch size is too small or too large (Hoffer et al., 2017),and it introduces a discrepancy between the behaviour of the model during training and at inference time. BatchNorm also adds memory overhead (Rota Bulo et al., 2018), and is a common source of implementation errors (Pham et al., 2019). In addition, it is often difficult to replicate batch normalized models trained on different hardware.

Signal Propagation Plots

💭: This is one of my favorite sections of the paper.Basically, signal propagation plots are plots that help measure "signal propagation" inside the network. How? We calculate some statistics at different points inside the network (during a single forward pass) and plot them.

Although papers have recently theoretically analyzed signal propagation in ResNets, practitioners rarely empirically evaluate the scales of the hidden activations at different depths inside a specific deep network when designing new models or proposing modifications to existing architectures. By contrast, we have found that plotting the statistics of the hidden activations at different points inside a network, when conditioned on a batch of either random Gaussian inputs or real training examples, can be extremely beneficial.

💭: We ideally want hidden activations to have zero mean and unit variance throughout the network. This is a good measure of "good" signal propagation.

😕: Do equations confuse you? This equation represents the Pre-Activation ResNet Block that is shown in Figure - 4(a) of the research paper. Here $f(.)$ represents the function on the Residual Branch which consists on BatchNorm, ReLU, Conv operations as in Figure - 4(a). This is also referred to as Residual Block in the coming sections.

💭: In simple words, create a network of your choice, initialize it with a suitable initialization scheme and make a forward pass on Guassian inputs with mean zero and unit variance.

Average Channel Squared Mean, computed as the square of the mean across the NHW axes, and then averaged across the C axis.
Average Channel Variance, computed by taking the channel variance across theNHWaxes, and then averaging across theCaxis.
Average Channel Variance on the end of the residual branch, before merging with the skip path.

💭: This might be a good time to pause, take a second, maybe even close your eyes and get a feel for SPPs? How do you think they look like? What's your expectation?

❓: What's the difference BN-ReLU-Conv layer and ReLU-BN-Conv ordering? Just swap the positions of BN and ReLU activation layer in Figure - 4(a) of the research paper.

🤔: Do the plots match your expectation? If not, why not?

Figure 1: Signal Propagation Plot for a ResNetV2-600 at initialization with BatchNorm, ReLU activations and He init, in response to an Gaussian input (with zero mean and unit variance) at 512px resolution. Black dots indicate the end of a stage. Blue plots use the BN-ReLU-Conv ordering while red plots use ReLU-BN-Conv.

💭: This model has been defined in the code provided by the authors here.

💭: Figure 1 really confused me for quite some time. I somehow thought that ResNet-V2 600 has 600 residual blocks and therefore was confused by the X-axis of figure 1 as the maximum value on the X-axis is 200. But, now having looked at the code, this makes sense. I hope this is something that's not as confusing to you as it was for me.

Key Observations

⚠️: Pay extra attention here.

From Figure 1(b), Average Channel Variance grows linearly with the depth at a given stage, and resets at each Transition Blocks.
From Figure 1(a), for BN-ReLU-Conv ordering, the Average Square Channel Means display similar behavior growing linearly with depth.

💭: I direct the curious reader to the SkipInit research paper by the same authors where they replace BatchNorm with a scalar that downscales the signal at the residual branch at initialization.

💭: If we can mimic these patterns in Normalizer-Free ResNets, then we can make sure that the new networks without normalization also train well and are competitive to their normalized counterparts.

Plotting Signal Propagation Plots for custom models in PyTorch

import torchvision
from timm.utils.model import extract_spp_stats, avg_ch_var, avg_ch_var_residual, avg_sq_ch_mean

model = torchvision.models.resnet50()
spp_stats = extract_spp_stats(m, 
			hook_fn_locs=['layer?.?', 'layer?.?', 'layer?.?.bn3'], 
			hook_fns=[avg_sq_ch_mean, avg_ch_var, avg_ch_var_residual])

# plot stats
fig, ax = plt.subplots(1, 3, figsize=(18,3), sharey=True)
ax[0].plot(stats['avg_sq_ch_mean'], label='avg_sq_ch_mean');
ax[0].legend(); ax[0].grid();
ax[1].plot(stats['avg_ch_var'], label='avg_ch_var');
ax[1].legend(); ax[1].grid();
ax[2].plot(stats['avg_ch_var_residual'], label='avg_ch_var_residual');
ax[2].legend(); ax[2].grid();

Figure 2: Signal Propagation Plots for ResNet-50 BN-ReLU-Conv ordering using timm

💭: I was lucky enough to have contributed SPPs to timm. For a complete notebook to replicate Figure-1 for a ResNet V2 600, refer here.

Normalizer-Free ResNets

💭: You might find this section below a little more complicated than the ones above but it is also the most important as this is where Normalizer-Free ResNets are introduced. Feel free to re-read this section or reach out to me should you have any questions.

💭: Really, all we needed is a good understanding of BatchNorm normalized networks and SPPs to develop normalizer free networks. Essentially, the authors create normalizer free networks whose SPPs mimick the SPPs normalized ResNets. Why? This has been answered in the Q&A section below.

BatchNorm downscales the input to residual block by a factor proportional to the standard deviation of the input signal.
Each residual block increases the variance of the signal by an approximately constant factor. (There is linear growth in Average Channel Variance)

💭: Note that this is in contrast to Residual Blocks where $x_{L+1} = f_{L}(x_{L}) + x_{L}$

⚠️: Pay extra attention here.

$f_l(.)$, the function computed by the residual branch is parameterized to be variance preserving at initialization. That is, $Var(f(x_l)) = Var(x_l)$.
$β_l$ is a scalar function chosen as $\sqrt{var(x_l)}$. This ensures that $f_l(.)$ has unit variance.
$α$ is a scalar hyperparameter which controls the rate of variance growth between blocks.

💭: Yes, I know that that's a lot of information and a lot less explanation. Explanation comes next in the form of Q&A.

Q&A on Normalizer Free ResNets

💭: In this section I try to answer some questions that were very confusing to me when I first read the paper.

This constraint enables us to reason about the signal growth in the network, and estimate the variances analytically.

The choice of how signals should propagate in unnormalized networks is largely one of design. In Appendix G.2, we note that we initially explored designing networks to have constant variance, which without prior knowledge one might assume to be a superior choice. We found that such networks were not as performant, and reasoned that mimicking a signal propagation template which we know to work well was a good design choice, as is supported by our experiments.

💭: It took me two weeks to get to this point and understand that $Var(x_{l+1}) = Var(x_l) + α^2$, so it's okay if you don't get it the first time. Let me know if you don't get this part! It's important for a complete understanding of NF-ResNets but not so much from an implementation perspective..

💭: This is based on my understanding of the paper and hasn't been explicitly explained elsewhere.

Downscale the input to residual branch
Increases the variance of the signal by an approximately constant factor for each residual block

💭: I have tried my best to explain all the magic and thinking behind Normalizer Free networks above. If there are parts of this section that didn't make sense, please feel free to let me know towards the end of this report. :)

NF-ResNets in PyTorch/Jax

import timm 
import torch 

m = timm.create_model('nf_resnet50')
x = torch.randn(1, 3, 224, 224)
m(x).shape

>> torch.Size([1, 1000])

💭: I absolutely love timm. It is one of the fastest growing libraries and is kept up to date by Ross. Latest research papers make their way to timm really really quickly! I have also been lucky enough to work on timmdocs project for a more in depth documentation on timm.

💭: If you want to experiment with Nf-ResNets, in PyTorch, using timm is one of the easiest ways to get started with these networks. This is also a great resource by Ayush Thakur that uses PyTorch Lightning and timm for some quick experimentation on batch sizes.

💭: We could have also looked at the official code implementation in Jax by deepmind, but for simplicity, let's stick to PyTorch. We are going to be using the source code from timm.

# copied from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/nfnet.py
import torch
import torch.nn as nn

from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
from timm.models.helpers import build_model_with_cfg
from timm.models.layers import ClassifierHead, DropPath, AvgPool2dSame, ScaledStdConv2d, ScaledStdConv2dSame,\
    get_act_layer, get_act_fn, get_attn, make_divisible

class NormFreeBlock(nn.Module):
    """Normalization-Free pre-activation block."""

    def __init__(
            self, in_chs, out_chs=None, stride=1, dilation=1, first_dilation=None,
            alpha=1.0, beta=1.0, bottle_ratio=0.25, group_size=None, ch_div=1, reg=True, extra_conv=False,
            skipinit=False, attn_layer=None, attn_gain=2.0, act_layer=None, conv_layer=None, drop_path_rate=0.):
        super().__init__()
        out_chs = out_chs or in_chs
        mid_chs = int(in_chs * bottle_ratio)
        self.alpha = alpha
        self.beta = beta

        if in_chs != out_chs or stride != 1 or dilation != first_dilation:
            self.downsample = DownsampleAvg(
                in_chs, out_chs, stride=stride, dilation=dilation, first_dilation=first_dilation, conv_layer=conv_layer)
        else:
            self.downsample = None

        self.act1 = act_layer()
        self.conv1 = conv_layer(in_chs, mid_chs, 1)
        self.act2 = act_layer(inplace=True)
        self.conv2 = conv_layer(mid_chs, mid_chs, 3, stride=stride, dilation=first_dilation, groups=groups)
        self.act3 = act_layer()
        self.conv3 = conv_layer(mid_chs, out_chs, 1, gain_init=1. if skipinit else 0.)

    def forward(self, x):
        out = self.act1(x) * self.beta

        # shortcut branch
        shortcut = x
        if self.downsample is not None:
            shortcut = self.downsample(out)

        # residual branch
        out = self.conv1(out)
        out = self.conv2(self.act2(out))
        out = self.conv3(self.act3(out))
        out = out * self.alpha + shortcut
        return out



💭: I leave it as an exercise to the reader here to map the Figure-3 - (a) & (b) to the code implementation from timm for Normalizer Free Blocks shared above. The snippet of code is capable of implementing both the Pre-Activation ResNet Block & Pre-Activation ResNet Transition Block.

💭: Also, then to create Nf-ResNet is a simple case of repeating the blocks as per the various ResNet configurations mentioned here. This is again left as an exercise but happy to dig deeper into this if needed. I leave it as an exercise because it would really be a good exercise/project to understands Nf-ResNets by recreating the networks using the source code from timm as a guide.

💭: All the pieces required for understanding the complete implementation have been shared as part of this report already.

Scaled Weight Standardization

💭: There is still a mystery that's left to be solved, but, it's the last one.

Figure 4: SPPs for three different variants of the ResNetV2-600 network(with ReLU activations). In red, a batch normalized network with ReLU-BN-Conv ordering. In green, a normalizer-free network with He-init and α = 1. In cyan, the same normalizer-free network but with Scaled Weight Standardization.

For NF-ResNets, the average value of the squared channel mean grows rapidly with depth, achieving large values which exceed the average channel variance. (comparing Figure-5(a) and Figure-5(b))
For NF-ResNets, The scale of the empirical variances on the residual branch are consistently smaller than one. (Figure - 5(c) in green)

💭: As I've mentioned before, it is considered good propagation if the activations have a zero mean and unit variance throughout the network. Thus, having an average value of the squared channel mean that grows rapidly with depth & values of variances that are consistently smaller than one represents instability in the network.

💭: Scaled Weight Standardization is an extension of Weight Standardization. Essentially, weight standardization normalizes weights in convolution layers, i.e., making the weights have zero mean and unit variance. Note that this is different from Batch Normalization (BatchNorm) which standardizes the hidden activations instead. For the curious reader, here is a wonderful video by Yannic Kilcher, that explains Weight Standardization.

💭: Don't worry if this seems confusing, things become clearer with code implementation shared below. Essentially, we standardize the weights of every CNN by subtracting the mean and dividing by standard deviation of the weights.

💭: The only difference between Scaled Weight Standardization and Weight Standardization is the introduction of a fixed constant γ which is non-linearity dependant. I refer the curious reader to section 4.2 of the paper that explains why.

💭: I have an implementation of Weight Standardization in PyTorch here.

# from https://github.com/deepmind/deepmind-research/tree/master/nfnets
_nonlin_gamma = dict(
    identity=1.0,
    celu=1.270926833152771,
    elu=1.2716004848480225,
    gelu=1.7015043497085571,
    leaky_relu=1.70590341091156,
    log_sigmoid=1.9193484783172607,
    log_softmax=1.0002083778381348,
    relu=1.7139588594436646,
    relu6=1.7131484746932983,
    selu=1.0008515119552612,
    sigmoid=4.803835391998291,
    silu=1.7881293296813965,
    softsign=2.338853120803833,
    softplus=1.9203323125839233,
    tanh=1.5939117670059204,
)

Scaled Weight Standardization in PyTorch

 class ScaledStdConv2d(nn.Conv2d):
    """Conv2d layer with Scaled Weight Standardization.

    Paper: `Characterizing signal propagation to close the performance gap in unnormalized ResNets` -
        https://arxiv.org/abs/2101.08692
    """
    def __init__(
            self, in_channels, out_channels, kernel_size, stride=1, padding=None, dilation=1, groups=1,
            bias=True, gamma=1.0, eps=1e-5, gain_init=1.0):
        if padding is None:
            padding = get_padding(kernel_size, stride, dilation)
        super().__init__(
            in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation,
            groups=groups, bias=bias)
        self.gain = nn.Parameter(torch.full((self.out_channels, 1, 1, 1), gain_init))
        self.scale = gamma * self.weight[0].numel() ** -0.5  # gamma * 1 / sqrt(fan-in)
        self.eps = eps

    def get_weight(self):
	std, mean = torch.std_mean(self.weight, dim=[1, 2, 3], keepdim=True, unbiased=False)
        weight = self.scale * (self.weight - mean) / (std + self.eps)
        return self.gain * weight

💭: Again, I leave it as an exercise to the match the implementation above with the formula. But of course, feel free to reach out to me should there be any confusion.

Conclusion