Logit Attenuating Weight Normalization

Aman Gupta, Rohan Ramanath, Jun Shi, Anika Ramachandran, SIROU ZHU, Mingzhou Zhou, Sathiya Keerthi

Keywords: deep learning, gradient methods, stochastic optimization, generalization gap, imagenet, adam, large batch training

Abstract:

Over-parameterized deep networks trained using gradient-based optimizers is a popular way of solving classification and ranking problems. Without appropriately tuned regularization, such networks have the tendency to make output scores (logits) and network weights large, causing training loss to become too small and the network to lose its adaptivity (ability to move around and escape regions of poor generalization) in the weight space. Adaptive optimizers like Adam, being aggressive at optimizing the train loss, are particularly affected by this. It is well known that, even with weight decay (WD) and normal hyper-parameter tuning, adaptive optimizers lag behind SGD a lot in terms of generalization performance, mainly in the image classification domain.

An alternative to WD for improving a network's adaptivity is to directly control the magnitude of the weights and hence the logits. We propose a method called Logit Attenuating Weight Normalization (LAWN), that can be stacked onto any gradient-based optimizer. LAWN initially starts off training in a free (unregularized) mode and, after some initial epochs, it constrains the weight norms of layers, thereby controlling the logits and improving adaptivity. This is a new regularization approach that does not use WD anywhere; instead, the number of initial free epochs becomes the new hyper-parameter. The resulting LAWN variant of adaptive optimizers gives a solid lift to generalization performance, making their performance equal or even exceed SGD's performance on benchmark image classification and recommender datasets. Another important feature is that LAWN also greatly improves the adaptive optimizers when used with large batch sizes.

One-sentence Summary: An optimizer for deep learning called Logit Attenuating Weight Normalization (LAWN) for superior generalization performance and scaling to large batches

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/arxiv:2108.05839/code)