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: This paper studies the extreme multi-label classification problem. Specifically, on top of the existing partition-based approaches that partition the label space into mutually exclusive clusters, the paper proposes a solution to relax this assumption by disentangle the multi-modal labels with non-exclusive clustering as an optimization approach. Evaluations are reported on four datasets against the existing baselines.