Training Tricks Miscellany: Reparameterization, Label Smoothing, and Dropout

This post gathers a few scattered but handy training tricks: reparameterization, weight EMA, Shake-Shake, label smoothing, and Dropout.

Reparameterization

Reparameterization means: at training time you add extra layers or paths to the network to aid optimization (making training smoother, or preventing overfitting), and at test time you fuse these parts back into the original structure. This way the extra cost is incurred only during training; once fused at inference, it is no different from the original structure, yet you gain a bit of accuracy for free. (The parameters of a BN layer can also be fused into the preceding convolution at test time, which belongs to the same line of thinking.) Three representative works are ACNet, RepVGG, and RepMLP.

ACNet splits one convolutional layer into three paths at training time (adding 1×31\times3 and 3×13\times1 convolutions, along with a BN for each path):

Reparameterization in ACNet
ACNet: at training time, one convolution is split into three paths with BN—3×3, 1×3, and 3×1.

At training time each path has its own parameters and ordinary backpropagation suffices; at test time the three paths are fused back into the original 3×33\times3 convolution and its BN—first fuse each path’s BN into its convolution, then add the three paths’ convolution parameters together, at which point the forward pass of the three paths is equivalent to a single convolutional layer.

Parameter fusion after ACNet training is complete
ACNet’s fusion: first fuse BN into the convolution, then add the three paths’ convolutions into one.

RepVGG uses reparameterization to train a VGG-style simple structure to very high accuracy—at training time it uses residual connections, which are easier to train than the original VGG and can exceed the accuracy of a ResNet with the same compute; at inference time the residual connections are not needed, so it gets to enjoy the high inference efficiency that comes with a simple VGG-style structure.

RepVGG fuses three branches into a single 3×3 convolution
RepVGG: multi-branch + residual at training time, fused into a single 3×3 convolution at inference.

RepMLP takes this further, simplifying the structure from VGG all the way to an MLP.

Thinking deeper: ACNet and RepVGG split one branch into multiple branches at training time, and these branches contain only linear transformations (only convolutions and BN, no activation functions), so they can be fused together directly. In theory, any number of branches containing only linear transformations can be fused into one—it’s just that performance doesn’t necessarily improve as a result. For branches that are all convolutions, as long as none of the kernels exceed k×kk\times k, they can all be fused into one k×kk\times k convolution. Generalizing further, a convolution can be implemented with img2col + matrix multiplication, so any linear transformations across all branches can be fused into a single matrix multiply-add on one branch—which means reparameterization may have promising applications in structures like the Transformer as well.

Weight EMA

Weight EMA (Exponential Moving Average) maintains an additional exponential moving average of the weights during training, and at test time uses these smoothed weights rather than the latest ones. It adds almost no training cost, yet often yields more stable and better generalization.

Shake-Shake

Shake-Shake randomly splits a branch into two parts for computation at training time (regenerating the random number αi\alpha_i on every forward pass). The figure below shows its application on a residual connection: left is the forward pass, middle is the backward pass, and right is the test phase (where α=0.5\alpha=0.5). As you can see, Shake-Shake brings a considerable amount of extra computational overhead.

Shake-shake applied to a residual connection
Shake-Shake: different random weights for the forward and backward passes, with α=0.5 at test time.

Label Smoothing

Label smoothing, as the name suggests, “smooths out” the labels a bit, and experiments show this improves generalization. The cross-entropy loss commonly used for classification is:

H(p,q)=i=1Kqilog(pi)\mathcal{H}(p,q)=-\sum_{i=1}^{K}q_i\log(p_i)

An ordinary one-hot label is:

qi={1,i=y0,otherwiseq_i=\begin{cases}1,& i=y\\ 0,& \text{otherwise}\end{cases}

Label smoothing changes it to (where KK is the number of classes and ϵ\epsilon is a small smoothing coefficient):

qi={1ϵ,i=yϵ/(K1),otherwiseq_i=\begin{cases}1-\epsilon,& i=y\\ \epsilon/(K-1),& \text{otherwise}\end{cases}

Dropout

Dropout is a common means of preventing overfitting, and the approach is simple: during training, for each element entering the dropout layer, it is dropped (set to zero) with probability 1p1-p.

The effect of Dropout
Dropout: at training time, a fraction of elements is randomly dropped (set to zero) by probability.

At test time all elements are kept, so training amounts to approximating a single network as an ensemble of multiple sub-networks. In addition, at test time the dropout layer scales the input from ww to pwpw, the purpose being to make each element’s output through the dropout layer have the same expected value at training and test time.

The difference in dropout between training and inference
Drop at training time; keep all and scale by p at test time, to keep the output expectation consistent.

References

  • Ding, Xiaohan, et al. ACNet: Strengthening the Kernel Skeletons for Powerful CNN via Asymmetric Convolution Blocks. ICCV, 2019.
  • Ding, Xiaohan, et al. RepVGG: Making VGG-style ConvNets Great Again. arXiv:2101.03697, 2021.
  • Ding, Xiaohan, et al. RepMLP: Re-parameterizing Convolutions into Fully-connected Layers. arXiv:2105.01883, 2021.
  • Gastaldi, Xavier. Shake-Shake Regularization. arXiv:1705.07485, 2017.
  • Szegedy, Christian, et al. Rethinking the Inception Architecture for Computer Vision. CVPR, 2016.
  • Srivastava, Nitish, et al. Dropout: A Simple Way to Prevent Neural Networks from Overfitting. JMLR, 2014.