Computationally Efficient W asserstein Loss for Structured Labels (2021)

Abstract: The problem of estimating the probability distribution of labels has been widely studied as a label distribution learning (LDL) problem, whose applications include age estimation, emotion analysis, and semantic segmentation. We propose a tree-Wasserstein distance regularized LDL algorithm, focusing on hierarchical text classification tasks. We propose predicting the entire label hierarchy using neural networks, where the similarity between predicted and true labels is measured using the tree-Wasserstein distance. Through experiments using synthetic and real-world datasets, we demonstrate that the proposed method successfully considers the structure of labels during training, and it compares favorably with the Sinkhorn algorithm in terms of computation time and memory usage.

## introduction label distribution learning (ldl), which is a generalized framework for performing single/multilabel classification and estimating the probability distribution over labels, is an important machinelearning problem @xcite . its applications include age estimation @xcite , emotion estimation @xcite , head-pose estimation @xcite , and semantic segmentation @xcite . in particular, multilabel classification is an important problem in many nlp areas, and has several applications including multi-label text classification @xcite . typically, kullback-leibler (kl) divergence is used to measure the similarity between two distributions. however, the kl divergence can tend to infinity if the supports of the two distributions do not overlap, resulting in model failure. to solve this support problem, wasserstein distance is used instead of kl divergence @xcite . wasserstein distance is defined as the cost of optimally transporting one probability distribution to match another @xcite . because it can compare two probability measures while considering the ground metric, it is more powerful than measurements that do not consider geometrical information. an ldl framework with wasserstein distance has been recently proposed @xcite . this framework employs the sinkhorn algorithm @xcite to calculate the wasserstein distance, which requires quadratic computational-time. thus, when we consider extremely large label-sets, for example, 10 5 , the computation cost can be significant. however, the wasserstein distance on a tree (hereinafter called tree-wasserstein distance) can be written in a closed-form and calculated in linear computation time @xcite . in this paper, we propose a tree-regularized ldl algorithm with a tree-wasserstein distance. the key advantage of the tree-wasserstein distance is that it considers the hierarchical label information explicitly, whereas the sinkhorn-based algorithm needs a cost matrix using tree-structured data. moreover, the tree-wasserstein distance has an analytic form that can be computed in linear time using significantly less memory. we experimentally demonstrate that the proposed algorithm compares favorably with the sinkhorn-based ldl algorithm @xcite with considerably lower memory consumption and computational costs. we demonstrate that the calculation is more efficient than that of the existing wasserstein loss. illustration of a tree-structured label with the root "animal". γ("mammal") = {"mammal", "dog", "cat"}, v @xmath0. ## problem setting we observe n input and output samples {(x 1 , y 1 ), • • • , (x n , y n )} from (x , y), where x ⊂ r d . we consider the problem of learning a map from a feature space x into p, which is a set of distributions over a finite set y. for example, multi-class classification is included in this problem, y, which represents the ℓ-th class, and it is expressed as the following onehot vector: when multi-label classification is considered, p denotes binary vectors that indicate existing labels. for example, if the sample x belongs to classes ℓ and ℓ , y is given as we aim to estimate the conditional probability vector p y for x by considering the structure information of y from {(x 1 , p y 1 ), • • • , (x n , p y n )}. ## proposed method in this study, we assume y has a tree-structure. accordingly, we propose ldl with tree-wasserstein distance. ## experiments we applied our proposed method to ldl on trees based on a synthetic dataset and to multi-label text classification of a hierarchical structure based on a real dataset. we implemented all the methods using pytorch @xcite . our models were optimized using a gradient method with the adam (kingma and ba, 2015) optimizer. baselines: we compared our proposed method to the wasserstein-loss-based ldl framework @xcite and a multi-class kl loss mentioned in (3). notably, in the original paper @xcite , they did not include kl loss and used only wasserstein loss, but @xcite used a linear combination of kl divergence and wasserstein distance as the loss. to ensure fair comparison, we also report the combination of wasserstein loss and multi-class kl loss as a strong baseline. therefore, we set the combination parameter λ = {0, foot_0 2 , 1} defined in eq 2 and the weight of all edges to 1. the wasserstein loss was computed using the sinkhorn-knopp algorithm in the log domain @xcite on gpus. for the proposed method, we computed the tree-wasserstein loss on the cpu and then passed it to the gpu to compute the gradient. then, we set the number of iterations of the sinkhorn-knopp algorithm to 10 and the regularization parameter to 50, respectively. ## conclusions this study proposed the use of a tree-wasserstein reguralizer for learning. the experimental results indicate that our proposed method can successfully predict the distributions of structured labels and that it outperforms existing wasserstein loss calculation methods in terms of both computational speed and memory consumption.