  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.
