Squared Earth Mover's Distance-based Loss for Training Deep Neural Networks

Deep neural networks (DNNs) especially convolutional neural networks (CNNs) have achieved state-of-the-art performances in many applications, and they have been shown to be especially powerful in multi-class classification tasks. Existing DNNs and CNNs are typically trained with a soft-max cross-entropy loss which only considers the ground-truth class by maximizing the predicted probability of the correct label. This cross-entropy loss ignores the intricate inter-class relationships that exist in the data. In this work, we propose to use the exact squared Earth Mover's Distance (EMD) as a loss function for multi-class classification tasks. The squared EMD loss uses the predicted probabilities of all classes and penalizes the miss-predictions accordingly. In experiments, we evaluate our squared EMD loss in ordered-classes datasets such as age estimation and image aesthetic judgment. We also generalize the squared EMD loss to classification datasets with orderless-classes such as the ImageNet. Our results show that the squared EMD loss allows networks to achieve lower errors than the standard cross-entropy loss, and result in state-of-the-art performances on two age estimation datasets and one image aesthetic judgment dataset.