Label Disentanglement in Partition-based Extreme Multilabel Classification

1 Introduction . The eXtreme Multi-label Classification ( XMC ) task is to find relevant labels from an enormous output space of candidate labels , where the size is in the order of millions or more ( [ 1 , 2 , 3 , 4 , 5 ] ; etc . ) This problem is of interest in both academia and industry : for instance , tagging a web page given its contents from tens of thousands of categories [ 6 ] ; finding a few products that a customer will purchase from among an enormous catalog on online retail stores [ 7 ] ; or recommending profitable keywords given an item/product from millions of advertisement keywords [ 8 ] . The XMC problem is challenging not only because of the data scalability ( e.g. , the number of instances , features , and labels are of the scale of millions or more ) , but also due to the label sparsity issue where there is little training signal for long-tailed labels . To tackle these issues , most prevailing XMC algorithms use a partition-based approach . Instead of ranking the entire set of millions of labels , they partition the label space into smaller mutually-exclusive clusters . Each instance is only mapped to one or a few label clusters based on a matching model , and then ranking is only conducted within the smaller subset of labels . Some exemplars are Parabel [ 8 ] , eXtremeText [ 9 ] , AttentionXML [ 10 ] , XR-Linear [ 11 ] and X-Transformer [ 12 ] . ∗Work done while Xuanqing was interning at Amazon Inc. Code will be released at https : //github.com/ amzn/pecos . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Partitioning labels into mutually-exclusive clusters may not be ideal . When labels are semantically complex and multi-modal , it is more natural to assign a label to multiple semantic clusters . In product categorization , for instance , the tag “ belt ” can be related to a vehicle belt ( under “ vehicle accessories ” category ) , or a man ’ s belt ( under “ clothing ” category ) . Assigning “ belt ” to just one of the clusters but not the other is likely to cause a mismatch for certain queries . To solve this problem , we reformulate the label clustering step as an assignment problem , where each label can be assigned to multiple clusters to allow disentanglement of mixed semantics . Further , we formulate learning optimal assignments by maximizing the precision as an optimization problem , and propose efficient solvers that automatically learn a good label assignment based on the current matching models . With this formulation , we apply our algorithm to alternatively refine label assignments and matching model in existing partition-based XMC methods to boost their performance . Similar idea can be found in [ 13 , 14 ] , however , our tree structures are constructed in one pass with hierarchical k-means . Our contributions can be summarized below : • We propose a novel way to obtain label assignments that disentangle multi-modal labels to multiple clusters . • Unlike previous methods that partition the label set before training , we propose an optimizationbased framework that allows optimizing label assignment with the matching and ranking modules inside a partition-based XMC solver . Our method is plug-and-play ; it is orthogonal to the models so most of the partition-based methods can benefit from our method . • Our proposed solution yields consistent improvements over two leading partition-based methods , XR-Linear [ 11 ] and X-Transformer [ 12 ] . Notably , with the concatenation of tfidf features and X-Transformer embeddings , we achieve new SOTA results on four XMC benchmark datasets . 2 Related Work . 2.1 XMC literature . Sparse Linear Models Sparse linear one-versus-rest ( OVR ) methods such as DiSMEC [ 4 ] , ProXML [ 15 ] , PDSparse [ 16 ] , PPDSparse [ 17 ] explore parallelism to speed up the algorithm and reduce the model size by truncating model weights to encourage sparsity . OVR approaches are also building blocks for many other XMC approaches . For example , in Parabel [ 8 ] , SLICE [ 18 ] , X-Transformer [ 12 ] , linear OVR classifiers with negative sampling are used . Partition-based Methods The efficiency and scalability of sparse linear models can be further improved by incorporating different partitioning techniques on the label spaces . For instance , Parabel [ 8 ] partitions the labels through a balanced 2-means label tree using label features constructed from the instances . Other approaches attempt to improve on Parabel , for instance , eXtremeText [ 9 ] , Bonsai [ 19 ] , NAPKINXC [ 20 ] , and XR-Linear [ 11 ] relax two main constraints in Parabel by : 1 ) allowing multi-way instead of binary partitions of the label set at each intermediate node , and 2 ) removing strict balancing constraints on the partitions . More recently , AttentionXML [ 10 ] uses BiLSTMs and label-aware attention to replace the linear functions in Parabel , and warm-up training the models with hierarchical label trees . In addition , AttentionXML considers various negative sampling strategies on the label space to avoid back-propagating the entire bottleneck classifier layer . Graph-based Methods SLICE [ 18 ] uses an approximate nearest neighbor ( ANN ) graph as an indexing structure over the labels . For a given instance , the relevant labels can be found quickly via ANN search . SLICE then trains linear OVR classifiers with negative samples induced from ANN . Graph-based partitions can be viewed as an extension of tree-based partitions , where at each layer of the tree , random edges are allowed to connect two leaf nodes to further improve connectivity . Nevertheless , such construction of overlapping tree structures is fully unsupervised , and is agnostic to the data distribution or training signals of the downstream XMC problem . 2.2 Overlapped Clustering . Finding overlapped clusters has been studied in the unsupervised learning literature [ 21 , 22 , 23 , 24 ] . For example , Cleuziou et al . [ 22 ] extend K-means with multi-assignments based on coverage , but tends to yield imbalanced clusters . This issue was later improved by Lu et al . [ 23 ] with sparsity constraints . Recently , Whang et al . [ 24 ] propose variants of k-mean objectives with flexible clustering assignment constraints to trade-off the degree of overlapping . However , all these unsupervised approach can not be optimized with existing partition-based XMC methods . 3 Background of partition-based XMC . The XMC problem can be characterized as follows : given an input instance x ∈ Rd and a set of labels L = { 1 , 2 , . . . , L } , find a model that retrieves top relevant labels in L efficiently . The model parameters are estimated from the training dataset { ( xi , yi ) : i = 1 , . . . , n } where yi ∈ { 0 , 1 } L denotes the relevant labels for xi from the label space L. We further denote X ∈ Rn×d as the feature matrix and Y ∈ Rn×L as the label matrix . Partition-based XMC methods rely on label space partitioning to screen out irrelevant labels before running the ranking algorithm . They often have the following three components ( depicted in Figure 1 ) : • The cluster assignments , where a set of labels are assigned to each cluster . Assuming there are K clusters , we use Si = { ` i1 , ` i2 , . . . } to denote the labels assigned to cluster i , and ⋃K i=1 Si = L. The clusters are constructed by K-means on label features , such that the clusters { Si } Ki=1 are mutually exclusive and unaware of the matcherM and rankerR , since clustering is typically performed only once before the matcher and ranker are trained . On the contrary , as we will see shortly , our new method provides a way to refine cluster assignments based on the matcher and is able to effectively unveil multimodal labels . • The matcher , which matches the input data xi ∈ Rd to a small set of candidate label clusters M : xi 7→ { S1 , S2 , . . . , Sb } , b ≤ K. ( 1 ) Each label set Sk contains a small fraction of labels ( a few hundreds ) , and b is called the beam size . One recursive implementation of matcher is seen in XR-Linear [ 11 ] , where inside the matcher , there is a tree constructed by recursive K-means clustering . On each level of the tree , a maximum of b nodes is selected according to the scores obtained from linear models . After that , all siblings of b nodes will be expanded at the next level to repeat this process recursively until b leaves are obtained , which results in the match set { S1 , S2 , . . . , Sb } . See the “ Matcher ” block in Figure 1 . • The ranker , which ranks the candidate labels collected from the matcher , with denoting the ranking order : R : M ( xi ) 7→ ` ( 1 ) ` ( 2 ) · · · ` ( w ) , where { ` ( 1 ) , ` ( 2 ) , . . . , ` ( w ) } = b⋃ i=1 Si . ( 2 ) Lastly , top-k labels are returned as the final prediction . See the “ Ranker ” block of Figure 1 . Partition-based XMC typically assumes labels in the same cluster are similar , and thus when training the ranker , they focus on distinguishing the labels ( and the corresponding samples ) within each cluster . Taking the widely used linear ranker as an example , they often assign a weight vector for each label . The weight vector w ` for label ` can be obtained by the following training objective : min w ` ∑ i∈Dpositive L ( x > i w ` ; +1 ) + ∑ i∈Dnegative L ( x > i w ` ; −1 ) , ( 3 ) where L ( · , · ) is the loss function ( e.g. , squared-hinge loss ) ; Dpositive and Dnegative are the positive and negative samples for label ` . In XMC , it is not efficient to collect all non-positive data D \ Dpositive as the negative part to train w ` ; instead , they often sample the “ hard negatives ” which is a tiny subset Dnegative ⊂ D \ Dpositive . Different negative sampling changes the loss function and the final results . For instance , in Parabel [ 8 ] and XR-Linear [ 11 ] , Dnegative are the examples having similar labels but not ` itself , based on the intuition that the learning can be more efficient by separating similar but different labels . In other words , Dpositive = { xi | ` ∈ yi } , Dnegative = { xi | ` 6∈ yi but “ similar ” } . ( 4 ) The label similarity information is hidden in the cluster assignments { S1 , S1 , . . . , SK } , with the assumption that similar labels are clustered . 4 Proposed Method . 4.1 Motivation . The central idea of the partition-based XMC methods is to partition labels into disjoint clusters , such that labels in each cluster share similar semantics . This relies on the unimodal assumption that each label only represents a pure , uniform semantic across all positive samples . However , we observe in practice that this assumption may not hold in general . Given labels that have multi-modal semantics , it is natural to seek a method to disentangle their semantics from each other and treat each semantic as a separate label , and further assign them to different clusters . To achieve this , we weaken the requirement that label clusters under leaf nodes are mutually exclusive Si ∩ Sj = ∅ , to a more suitable , limited overlapping : for any label ` ∈ { 1 , 2 , . . . , L } , it can appear at most λ-times among { S1 , S2 , . . . , SK } . Although allowing label clusters to overlap with each other does not explicitly disentangle the semantics of a label , it paves the way for learning multiple versions of weights for the same label . Take the linear model in Eq . ( 3 ) for simplicity , label ` is seen in both cluster SA and SB , so we will have two weight vectors wA and wB for the same label . wA and wB are trained separately in Eq . ( 3 ) , where they have different negative examples DAnegative and DBnegative . Recall in Eq . ( 4 ) the negative part of the data relies on the label similarity induced from clusters SA and SB . Therefore , wA and wB will eventually converge to two different solutions – this is how the semantics are disentangled when we assign a label to multiple clusters .

Summary: Label partitioning approaches are well-explored in Extreme Classification (XC) literature whose core idea is to cluster semantically similar labels together and then to sample the negative labels only from a small number of relevant clusters. Although this has been shown to significantly speed up the training and prediction routines, most existing partitioners simplistically assume that each label belongs to exactly one cluster. This could degrade prediction accuracies since labels in XC can often be multi-modal with multiple associated semantics. This paper addresses this issue by relaxing the single cluster assumption and thus allowing each label to instead end up in \lambda different clusters. This is achieved though a novel cluster assignment algorithm which is shown to be theoretically well-defined and leads to small but consistent accuracy gains. The proposed algorithm has the added advantage of being plug-and-play, i.e. it can be used with many matching and ranking models based on label partitioning. Enough experiments have been conducted on simulated and real datasets to demonstrate the accuracy gains and discuss the associated increases in costs.