Normalizer-Free (Res)Nets

Assumptions

• You know how deep-learning (and back-prop) works
• You understand the concept of ResNets
• You know (basics of) PyTorch
• You heard about Batch Normalization

Agenda

• Introduction
• The problem with Batch Normalization
• How to eliminate Batch Normalization
• Gradient clipping - novel idea
• View from an ML Engineer perspective
• Discussion

Quick summary

I'm not into CV, why should I care?

Ideas propagate from one domain to another

What works for CV is likely to work for other fields

Batch Normalization (BN)

Quick recap

Why is BN useful?

1. BN downscales the residual branch

2. BN eliminates mean-shift

3. BN allows efficient large-batch training

BN downscales the residual branch

After BN model is "biased" towards skip connection =>
deeper networks can be trained

BN eliminates mean-shift

BN allows efficient large-batch training

BN smoothens the loss so bigger batches can be used

The bigger the batch size the larger stable lr

Bigger batch size => less "update steps" required

If it is so useful, why is BN a problem?

BN is ridiculously slow to compute on a GPU

BN breaks the independence between training examples in the mini-batch

Proposed solution

Instead of doing costly BN

Let's "predict" the variance shift

Then just scale the result by a right scalar

Variance shift correction

In inference mode, independent from other examples in a batch

Intuition

Prevention of a mean shift
in the hidden activations

NF ResNet

Batch size scaling problem

Adaptive Gradient Clipping

Clip only too big gradients

Adaptive Gradient Clipping in code

def adaptive_clip_grad(parameters, clip_factor=0.01, eps=1e-3, norm_type=2.0):
    for p in parameters:
        p_data = p.detach()
        g_data = p.grad.detach()
        max_norm = unitwise_norm(p_data, norm_type=norm_type).clamp_(min=eps).mul_(clip_factor)
        grad_norm = unitwise_norm(g_data, norm_type=norm_type)
        clipped_grad = g_data * (max_norm / grad_norm.clamp(min=1e-6))
        new_grads = torch.where(grad_norm < max_norm, g_data, clipped_grad)
        p.grad.detach().copy_(new_grads)
            

How is Adaptive Gradient Clipping useful?

Gradient clipping prevents optimizer from too big jumps

Adaptive makes additional usege of gradient to parameter proportion

Less dependent on hyper-parameter $\lambda$?

Training is more smooth (no jumps due to a noise in data)

Last, but not least

Proposed architecture

Start with: SE-ResNeXt-D

Add few tweaks

Overpriced TPU go brr

New SOTA on ImageNet

Do they provide code?

How to use it?

import timm
from utils import example_batch_of_images

model = timm.create_model('dm_nfnet_f0', pretrained=True)
model.eval()

prediction = model(example_batch_of_images)
prediction.size()
>>> torch.Size(128, 1000)
            

How to train it?

Exactly same approach as with Resnet/Efficientnet/Whatevernet

class NFNetBasedfCustomClassifier(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.criterion = nn.CrossEntropyLoss()
        self.model = timm.create_model('dm_nfnet_f0', pretrained=True)

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self.forward(x)
        loss = self.criterion(logits, y)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)

        return [optimizer], [scheduler]
            

So far so good
Where is the catch?

NFNet paper vs Efficiennet paper
What is the difference?

Number of parameters increased

Despite far larger number of parameters, comparable times of a single full training step

Number of FLOPS increased

Better results, more FLOPs, but similar times (is that comparison fair?)

Discussion

Speculations

If BN is not a limiting factor more parameters can be added making
future models even less useful for "non-google" labs

Predictions

10.III.2021

• Google/Deepmind will do "EfficientNet-like" grid search for optimal architecture


• The NF Net (or its efficient descendant) will be used as a backbone of a new SOTA object detection model


• The idea of Adaptive Gradient Clipping will be successfully applied in transformer models, especially for CV (as they require far more resources)

TL;DR

• NF Net = New ImageNet SOTA
• Batch Normalization = slow operations
• No BN layer = faster forward/backward pass
• BN replaced with scaling and Adaptive Gradient Clipping
• (Probably) A new paradigm for
  designing architectures

Thanks

Feel free to ask ANY question

Piotr Mazurek
Presentation avalibe at: https://tugot17.github.io/NF-Nets-Presentation