Learning Compact Features via In-Training Representation Alignment
This work addresses training instability in deep learning for researchers and practitioners, but it is incremental as it builds on existing SGD methods with a novel alignment technique.
The paper tackles the problem of noisy and jumpy updates in deep neural networks due to variance in mini-batch gradients during SGD training by proposing In-Training Representation Alignment (ITRA), which aligns feature distributions of different mini-batches to stabilize training and extract compact features, resulting in superior performance on image and text classification tasks compared to strong baselines.
Deep neural networks (DNNs) for supervised learning can be viewed as a pipeline of the feature extractor (i.e., last hidden layer) and a linear classifier (i.e., output layer) that are trained jointly with stochastic gradient descent (SGD) on the loss function (e.g., cross-entropy). In each epoch, the true gradient of the loss function is estimated using a mini-batch sampled from the training set and model parameters are then updated with the mini-batch gradients. Although the latter provides an unbiased estimation of the former, they are subject to substantial variances derived from the size and number of sampled mini-batches, leading to noisy and jumpy updates. To stabilize such undesirable variance in estimating the true gradients, we propose In-Training Representation Alignment (ITRA) that explicitly aligns feature distributions of two different mini-batches with a matching loss in the SGD training process. We also provide a rigorous analysis of the desirable effects of the matching loss on feature representation learning: (1) extracting compact feature representation; (2) reducing over-adaption on mini-batches via an adaptive weighting mechanism; and (3) accommodating to multi-modalities. Finally, we conduct large-scale experiments on both image and text classifications to demonstrate its superior performance to the strong baselines.