Label Disentanglement in Partition-based Extreme Multilabel Classification
Year: 2022

Abstract:
Partition-based methods are increasingly-used in extreme multi-label classification (XMC) problems due to their scalability to large output spaces (e.g., millions or more). However, existing methods partition the large label space into mutually exclusive clusters, which is sub-optimal when labels have multi-modality and rich semantics. For instance, the label “Apple” can be the fruit or the brand name, which leads to the following research question: can we disentangle these multi-modal labels with non-exclusive clustering tailored for downstream XMC tasks? In this paper, we show that the label assignment problem in partition-based XMC can be formulated as an optimization problem, with the objective of maximizing precision rates. This leads to an efficient algorithm to form flexible and overlapped label clusters, and a method that can alternatively optimizes the cluster assignments and the model parameters for partition-based XMC. Experimental results on synthetic and real datasets show that our method can successfully disentangle multi-modal labels, leading to state-of-the-art (SOTA) results on four XMC benchmarks.

heading: 1 Introduction
text: 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.

heading: 2 Related Work

heading: 2.1 XMC literature
text: 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.

heading: 2.2 Overlapped Clustering
text: 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 cannot be optimized with existing partition-based XMC methods.

heading: 3 Background of partition-based XMC
text: 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.

heading: 4 Proposed Method

heading: 4.1 Motivation
text: 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.

heading: 4.2 Optimization-based Label Assignment Approaches
text: How to assign a label to multiple clusters? Previous methods consider clustering as a preprocessing step and apply unsupervised methods, such as K-means, to partition the label space. In contrast, we formulate label assignments as an optimization problem to maximize the precision rate for XMC, which allows learning a matcher-aware clustering assignment as well as alternative updates between cluster assignments and the matcher.
As mentioned before, our goal is to learn S1, . . . ,SK to cover all the labels, and the sets can be overlapping. To find the best label assignments, it is natural to maximize the precision rate. Many previous methods in information retrieval obtain ranking functions by optimizing (surrogates) Precision/Recall [25, 26]. Here our goal is to find a good combination of matcher and cluster assignments in the context of partition-based XMC, so the objective is totally different.
To formally define precision, we first define the output of matcher represented by matrix M ∈ Rn×K :
Mij = { 1, if xi is matched to leaf cluster Sj , 0, otherwise.
(5)
For beam size b, every row of M has exactly b nonzero entries. At the same time, the cluster assignments S1, . . . ,SK can be parameterized by the clustering matrix C ∈ {0, 1}L×K so that
Cij = 1 if and only if label-i is seen in cluster Sj , otherwise Cij = 0. With this notation, the candidate labels generated by matcher is given by Ŷ := Binary(MC>), Ŷ ∈ Rn×L, where Binary(A) = IA is the element-wise indicator function.
Now consider the number of true positives in top-k predictions (TP@k), it is upper bounded by the intersection between matcher predictions Ŷ and the ground truth Y , i.e.
TP@k = |Top-k(Ŷ ) Y | (i) ≤ |Ŷ Y | (ii)= Tr(Y >Ŷ ), (6)
where | · | denotes the #nnz elements in a sparse matrix; the inequality (i)
≤ is due to the error of ranker hidden in Top-k(·); and (ii)= is from the fact that both Ŷ and Y are binary. In this paper, we consider following scenarios that simplifies our analysis
• Perfect ranker: the ranker does not make any mistake, so the gap in Eq. (6) vanishes.
• Probably correct ranker: for any true positive label `+ ∈ Ŷ , it is ranked to top-k with a constant probability p, i.e. P(`+ ∈ Top-k(Ŷ )|`+ ∈ Ŷ ) = p. Then we have TP@k = p · Tr(Y >Ŷ ).
In both cases, the precision is proportional to Tr(Y >Ŷ ). We can then formulate the problem of learning the best cluster assignment as
maximize S1,S2,...,SK⋃K
i=1 Si={1,2,...,L}
Tr ( Y >Binary(MC>) ) , s.t. K∑ i=1 I`∈Si ≤ λ, ∀` ∈ {1, 2, . . . , L}. (7)
Note that {Si, . . . ,SK} are hidden in C. The first constraint ensures we cover the whole label set and the second constraint mitigates the case of degenerated clusters where some Si’s have too many labels, resulting in significantly increased time complexity. The parameter λ is the only hyperparameter to be tuned. In practice, we find λ is stable across datasets and for simplicity we select λ = 2 for all the experiments. We also show the sensitivity of our algorithm with respect to λ in Section 5.2.
The optimization problem (7) is combinatorial and hard to solve. In fact, we show this is NP-complete: Theorem 1. Problem (7) is NP-complete.
This can be proved by a polynomial time reduction from the set cover problem to Problem (7), and the proof is deferred to the Appendix A.1. To develop an efficient algorithm, we approximate the objective function of (7) with a continuous, ReLU-like function
Binary(MC>) ≈ max(MC>,0) = MC>, (8) where the first approximation comes from replacing binary function to ReLU; the second equality is because both M and C are positive in all entries. A special case is when beam size b = 1. In that case MC> is a binary matrix, so Binary(MC>) = MC>. We then consider the following simplified problem:
maximize S1,S2,...,SK⋃K
i=1 Si={1,2,...,L}
Tr ( Y >MC> ) , s.t. K∑ i=1 I`∈Si ≤ λ,∀` ∈ {1, 2, . . . , L}. (9)
This problem can be solved efficiently with a closed form solution: Theorem 2. Problem (9) has a closed form solution
C∗ = Proj(Y >M), (10) where the Proj(·) operator selects the top-λ elements for each row of the matrix.
The above result can be easily derived since the objective function is linear and the constraint is a row-wise `0 norm constraint. The detailed proof can be found in the Appendix A.2.
But is (10) also a good solution to the original problem (7)? In fact, we show that this solution, despite not being optimal for (7), is provably better than any non-overlapping cluster partitioning, which is used in almost all the existing partition-based XMC methods. Theorem 3. For any clustering matrix C corresponding to a non-overlapping partition of the label set, we have Tr(Y >Binary(M(C∗)>)) > Tr(Y >Binary(MC>)).
The proof can be found in the Appendix A.3. This theorem implies our partitioning can achieve higher precision over any existing non-overlapping clustering for any ranker.
Practical Implementation After C (or equivalently {Si}Ki=1) is determined, we finetune the matcherM to accommodate the new deployment of label clusters. Note that most of the existing partition-based XMC starts from unsupervised label clustering and then alternate between getting a new (overlapping) clustering and finetune the matcher. Our algorithm can plug-in into most of these methods by adding one or more loops of alternative cluster assignments and matcher updates. Given the current matcher, the cluster assignments are updated based on the proposed formulation (10), and then the matcher will be retrained following the same procedure used in the original partition-based XMC solver. The whole routine is exhibited in Algorithm 1. Since we use a balanced K-means label clustering as initialization, we found that after one step update in Line 6 of Algorithm 1, {Si}Ki=1 is still not too imbalanced in our experiments. It’s possible to add cluster size constraints in Problem (9) to enforce the label partition to be balanced, as discussed in Appendix A.4, but we do not find a practical need for such constraints.
Algorithm 1 Our proposed framework. 1: Input: training data 〈X,Y 〉, any partition-based XMC algorithm XMC-part. 2: Output: the trained model (i.e., matcherM, rankerR, label clusters {Si}Ki=1). 3: Initialize {Si}Ki=1 with balanced K-means using label features. 4: InitializeM← XMC-part({Si}Ki=1) 5: Compute matcher prediction matrix M =M(X) by Eq. (5). 6: Update label clustering: {Si}Ki=1 by Eq. (10). 7: %% Following two lines are called “alternative update” hereafter. 8: Finetune the matcher given new clusters:M← XMC-part({Si}Ki=1). 9: Train the ranker with updated clusters and matcher: R ← XMC-part({Si}Ki=1,M).
Deduplication at inference time. To perform inference with the new model, we need to deduplicate the scores from the same labels but in different clusters. This happens because we use beam search to efficiently search b paths from root node to leaves. Refer to Figure 1 where b = 2 is shown (S1 and S3 selected), if both S1 and S3 contain a same label `duplicate but different scores s1 and s3, then we average them together with the final score s(`duplicate) = 12 (s1 + s3). In this sense, our algorithm can be interpreted as a finer-grained ensemble which ensembles the scores inside a tree, rather than ensembling multiple, independently trained trees.

heading: 5 Experimental Results
text: Our proposed framework serves as a generic plugin for any partition-based XMC methods and we show its efficacy on multiple experiment setups. First, we verify that the proposed method improves over the baselines on synthetic datasets where we simulate labels with mixed semantics. Next, we test the sensitivity of hyperparameter λ, followed by experiments on real-world XMC benchmark datasets. We end this section with an ablation study.
Datasets. We consider four publicly available XMC benchmark datasets [2, 10] for our experiments. See Table 1 for data statistics. To obtain state-of-the-art (SOTA) results in Table 3, we concatenate the dense neural embeddings (from fine-tuned X-Transformer [12, 11]) and sparse TF-IDF features (from AttentionXML [10]) as the input features to train the models.
Evaluation Metric. We measure the performance with precision metrics (P@k) as well as Propensitybased scores (PSP@k) [5], which are widely-used in the XMC literature [27, 8, 28, 18, 12, 11]. Specifically, for a predicted score vector ŷ ∈ RL and a ground truth label vector y ∈ {0, 1}L, P@k = 1k ∑ l∈rankk(ŷ) yl; PSP@k = 1 k ∑k l=1 yrank(l) prank(l) , the latter focuses more on the tail labels.
Models. As we introduce a new technique that is generally applicable, our method must be combined with existing partition-based XMC algorithms. We take XR-Linear [11] as the backbone for all experiments except Section 5.3. To get SOTA results in large-scale real datasets, we change the backbone to X-Transformer [12] in Section 5.3. The implementation details for combining our method with XR-Linear and X-Transformer are discussed in the Appendix A.5.
Hyper-parameters. Our technique depends on hyperparameter λ, which is tested in a standalone section. For hyperparameters in the XMC model, we mostly follow the default settings in the corresponding software. The details about hyperparameters are listed in the Appendix A.6.

heading: 5.1 Synthetic datasets by grouping true labels
text: In this experiment, we create a synthetic data derived from public datasets by deliberately entangling some labels. The motivation is to see the capability of disentangling different semantics from fused labels. Specifically, we designed several methods to mix up labels in the order of easy, medium and hard. The grouping is done without replacement: suppose a dataset has L labels, and we want to entangle every k labels into a fake one, then after this process we will get an artificial dataset with dLk e labels. Below are three implementations we are going to evaluate:
• Easy mode: we run a balanced dLk e-means clustering on the original label space of size L, so that each cluster contains about k labels. After that, we represent each cluster with a separate label, and the new label embedding is the cluster centroid.
• Medium mode: we run a balanced d L32k e-means clustering on the label space. Different from the setting above, now each cluster contains 32k labels. Next, we randomly group the 32k labels in to 32 subsets, with k labels each.
• Hard mode: the most difficult mode is randomly shuffle labels into dLk e subsets, with k labels in each subset.
We label the difficulty of three modes based on the idea that when similar labels are entangled together, it is less a issue for the baseline model, so the accuracy drop is negligible; Whereas if unrelated labels are bundled together, our method starts to shine. See Figure 2 for the relative performance improvement over XR-Linear. Several conclusions can be made from this figure: first, the improvements are all positive, meaning our method consistently beats the baseline XR-Linear model in all scenarios; secondly, by comparing Wiki10-31K with Amazon-670K, we see Amazon670K shows an upwards in performance gain. This is mainly because the larger label space of Amazon-670K renders this problem even harder after combining labels, so our method exhibits a big edge over baselines. How are the labels disentangled to different clusters? We attach a random example from our dataset in Table 2. This example shows that our method can indeed separate labels with fused semantics, improving the quality of both matcher and ranker.

heading: 5.2 Hyperparameter sensitivity (λ)
text: In real applications, we never know the number of semantics a label has - either because the label space is huge, or it is uneconomical. In this section, we check the sensitivity of model performance on hyperparameter λ. To this end, we consider four different datasets and choose λ = {1, 2, . . . , 6}. The results are exhibited in Figure 3. From the figure above, we observe that the performance generally improves as we increase λ – our prior on the number of different meanings of each labels. However, since the benefits diminishes quickly, we choose λ = 2 for all the following experiments in this paper.

heading: 5.3 Real datasets
text: The above experiments showed that our method is indeed better at disentangling complex meanings from labels. In real scenarios, it is certainly not the case that every label will have tens of underlying semantics. On the contrary, they behave more like a mix of both unimodal and multimodal labels. However, in this experiment, we will show that our method is still favorable. We conduct experiments to show that our method can successfully boost the performance of existing partition-based XMC models, including XR-Linear [11] and X-Transformer [11].
In Table 4, we combine the proposed method with competitive linear models XR-Linear on four XMC benchmark datasets. We assume only the TF-IDF features are given and assume the single model setting without any ensemble. Our method consistently improves the performance over the baseline XR-Linear. Moreover, we tested how finetuning increases the accuracies (by comparing No/Yes in “alternative update?” column). The results prove that it is indeed necessary to train both the model parameters and cluster map in an alternative fashion (i.e., see Algorithm 1).
Next, we repeat the same process but based on the current SOTA X-Transformer model. The results are shown in Table 3. Notice that our method reuses the same dense+sparse features as in the XTransformer model, and the results are the ensemble of 9 models [12]. From Table 3 we can observe that our method achieves new SOTA precision results with a significant margin while still maintaining good PSP scores, comparing with other strong algorithms such as Parabel, AttentionXML, XR-Linear and X-Transformer. Notably, the numbers show that our method is more suitable for large label space (such as Amazon-3M) compared with smaller label space (e.g., AmazonCat-13K).
To benchmark the training and inference overhead due to repeated clustering and training in our method, we report the relative training and inference time under different λ’s shown in Table 5. We remark that our method is more economical when the underlying model is inherently slow to train. On Amazon-3M for example, our framework only introduces 1.1% training overhead and 21.1% inference overhead over the X-Transformer models.
In Appendix A.8, we provided one more experiment on the real-world datasets that compares our method with random baseline. The random baseline is constructed by duplicating the same amount of labels to random clusters, as opposed to our precision maximization algorithm.
Method Alternative update? Prec@1 Prec@3 Prec@5
Wiki10-31K XR-Linear N/A 84.14 72.85 64.09
+Ours No 84.52 74.23 65.41 +Ours Yes 84.57 74.55 65.59
AmazonCat-13K XR-Linear N/A 92.53 78.45 63.85
+Ours No 90.63 78.12 63.86 +Ours Yes 92.90 79.04 64.35
Amazon-670K XR-Linear N/A 44.17 39.14 35.38
+Ours No 43.22 38.19 34.10 +Ours Yes 44.82 39.93 36.34
Amazon-3M XR-Linear N/A 46.74 43.88 41.72
+Ours No 46.95 44.24 42.18 +Ours Yes 47.51 44.76 42.67
Table 4: Experiments with XR-Linear. Our method is called “XR-Linear+Ours”, we also tested the finetune step detailed in Algorithm 1(L7). When joinly train option is disabled, the matcherM won’t be updated despite the creation of new cluster map.
Method Extra training time Extra inference time
Wiki10-31K (baseline: Ttrn = 0m30s, Ttst = 0.7ms/pts) +Ours (λ = 1) 2.32× 2.40× +Ours (λ = 2) 5.58× 4.80× AmazonCat-13K (baseline: Ttrn = 1m37s, Ttst = 0.2ms/pts) +Ours (λ = 1) 3.05× 0.71× +Ours (λ = 2) 4.54× 1.59× Amazon-670K (baseline: Ttrn = 1m30s, Ttst = 0.4ms/pts) Ours (λ = 1) 0.47× 0.07× Ours (λ = 2) 0.96× 0.60× Amazon-3M (baseline: Ttrn = 13m49s, Ttst = 0.5ms/pts)
Ours (λ = 1) 0.33× 1.12× Ours (λ = 2) 1.63× 2.51×
Table 5: Extra training time and inference time (in ms/pts, which is millisecond per data point) of our method compared to the baseline XR-Linear. Our method posts negligible overhead for big model. For X-Transformer on Amazon-3M, we only have 1.1% training overhead and 21.1% inference overhead.

heading: 6 Conclusion
text: In this paper, we have proposed a simple way to disentangle the semantics reside in the labels of XMC problems. We have shown how to do this in an indirect way that builds overlapping label clusters. We proposed an optimization algorithm to solve both parts: the matcher model, the ranker and cluster assignments. In experiments, we tested under various cases and the results indicate that our method is exceptionally good when the label space is huge (such as Amazon-3M) or when the label contains many different meanings (such as the artificial data we created).

heading: Acknowledgement
text: This work is supported in part by NSF under IIS-1901527, IIS-2008173, IIS-2048280 and by Army Research Laboratory under agreement number W911NF-20-2-0158.

Review summaries:
The authors propose a new method for partitioning labels space that can be used with a family of currently popular partition-based/label tree-based extreme multi-label classifiers, which organize labels in the form of the tree. The proposed method introduces redundancy of labels in the different label clusters, which are mutually exclusive to the clusters in many methods. The problem of label assignment is formulated as maximization of matcher's (part of classifier predicting clusters for labels) precision, which is optimized by adding new labels to the cluster after initial training with clusters with mutually exclusive labels. The new partitioning method is applied to XL-Linear and X-TRANSFORMERS in the experimental study, resulting in a steady improvement in predictive performance on four popular benchmark datasets.

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. 

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.

In extreme Multi-label classification (XML) a large range of methods based their prediction on partitioning whether the input space or label space.  The paper addresses this partitioning phase, proposing a "disentangled partitioning", allowing overlapping clustering of labels. The proposed model consists to use two matrices, the first one M corresponding to the classification of the examples to each cluster (the matcher) and the second one C indicating the assignment of the labels to each cluster.  The goal is to optimize those matrices ensuring that  MC is close to the label matrix under constraints to ensure that each label is present up to lambda time in different clusters (the disentanglement factor). After a formal analysis of the problem, the authors propose a simple algorithm to optimize the matrix C, using existing approaches to learn the matcher. To validate the approach authors first propose an experiment merging artificially labels for clustering. The second series of experiments evaluate the performances wrt SOTA algorithms and varying the disentanglement parameter.