Source-linked AI summary
Bonsai -- Diverse and Shallow Trees for Extreme Multi-label Classification
Sujay Khandagale, Han Xiao, Rohit Babbar
TL;DR
Extreme multi-label classification must scale to very large label spaces while handling many sparsely observed tail labels and preserving prediction quality. Bonsai combines generalized input, output, and joint label representations with shallow trees learned through flexible multi-way partitioning. It achieves prediction diversity and tail-label coverage better than state-of-the-art tree methods, with comparable training speed and prediction accuracy at par with one-vs-rest methods.
Problem
Extreme multi-label classification must handle hundreds of thousands or more labels, including many tail labels with few training examples, while existing approaches trade accuracy, label diversity, or computational efficiency.
Method
Bonsai partitions labels in input, output-co-occurrence, or joint representation spaces using flexible multi-way clustering and relaxed balancedness constraints to learn shallow trees.
Results
Bonsai achieves prediction diversity and tail-label coverage better than state-of-the-art tree-based methods, with comparable training speed and prediction accuracy at par with one-vs-rest methods.
Takeaways & Limitations
Generalized label representations and shallow trees together let Bonsai combine the training efficiency of tree-based methods with the prediction quality and diversity associated with stronger alternatives.
Takeaways & Limitations
Unlike Parabel, Bonsai does not have logarithmic dependence on the number of labels for prediction complexity, although its prediction time is typically measured in milliseconds.
Abstract
from arXiv · showhide
Extreme multi-label classification (XMC) refers to supervised multi-label learning involving hundreds of thousand or even millions of labels. In this paper, we develop a suite of algorithms, called Bonsai, which generalizes the notion of label representation in XMC, and partitions the labels in the representation space to learn shallow trees. We show three concrete realizations of this label representation space including : (i) the input space which is spanned by the input features, (ii) the output space spanned by label vectors based on their co-occurrence with other labels, and (iii) the joint space by combining the input and output representations. Furthermore, the constraint-free multi-way partitions learnt iteratively in these spaces lead to shallow trees. By combining the effect of shallow trees and generalized label representation, Bonsai achieves the best of both worlds - fast training which is comparable to state-of-the-art tree-based methods in XMC, and much better prediction accuracy, particularly on tail-labels. On a benchmark Amazon-3M dataset with 3 million labels, \bonsai outperforms a state-of-the-art one-vs-rest method in terms of prediction accuracy, while being approximately 200 times faster to train. The code for Bonsai is available at \url{https://github.com/xmc-aalto/bonsai}
1 Introduction
Extreme multi-label classification must handle enormous label spaces and highly imbalanced, tail-heavy distributions while balancing accuracy against computational efficiency. Bonsai addresses this challenge with generalized label representations and shallow, diverse trees, combining fast training with stronger prediction and tail-label coverage.
- XMC challenge: Extreme multi-label classification assigns a small subset of relevant labels from hundreds of thousands or more possible labels.Such problems arise in product categorization, hashtag suggestion, encyclopedia annotation, image classification, advertising, and recommendation.
- XMC challenge: XMC datasets contain many tail labels with very few training instances, following power-law or Zipf-like distributions.In WikiLSHTC-325K, only approximately 150,000 of 325,000 labels have more than five training instances.
- Existing approaches: One-vs-rest methods offer strong accuracy and label diversity but are difficult to use in resource-constrained environments because they require distributed training.Tree-based methods train faster, but their cascades can propagate errors; embedding methods rely on low-rank label spaces that break down for power-law labels, while deep models face tail-label data scarcity.
- Existing approaches: Parabel improves computational scalability through recursive binary partitioning but remains sub-optimal in prediction performance, especially for tail labels.Its binary and equal-sized partition constraints can cause tail labels to be subsumed by head labels.
- Bonsai: Bonsai generalizes label representation using input, output-co-occurrence, or joint spaces, then learns shallow trees with flexible multi-way, unconstrained partitions.These design choices use K > 2 clustering and relax balancedness constraints to preserve diverse label subgroups.
- Bonsai: Bonsai combines prediction diversity and improved tail-label coverage with training speed comparable to tree-based methods and prediction accuracy at par with one-vs-rest methods.The paper attributes this outcome to the synergistic effects of richer label representations and shallow trees.
2 Formal description of Bonsai
Bonsai combines generalized label representations with unconstrained multi-way partitioning to build shallow tree cascades. Its input, output, and joint representation spaces support diverse partitions intended to preserve tail-label information and reduce propagation error.
- 2 Formal description of Bonsai: Bonsai follows a tree-structured label partitioning approach for scalable extreme classification.The model learns a vector-valued multi-label classifier from feature and label vectors while using tree structure for scalability.
- 2.1 Label representation: Labels can be represented in input, output, or joint spaces based on features, label co-occurrence, or both.The joint representation concatenates input and output representations and has dimensionality D+L.
- 2.1 Label representation: Input-space representations aggregate training-instance feature vectors associated with each active label.Each label representation remains D-dimensional and can be normalized to unit Euclidean length.
- 2.1 Label representation: Output-space representations encode each label through its co-occurrence counts with other labels.The resulting representation has the same dimensionality as the output space, η = L.
- 2.1 Label representation: The output representation matches Parabel competitively, while the joint representation surpasses state-of-the-art methods in prediction performance and label diversity.These results are reported for the generalized representations combined with shallow tree cascades.
- 2.2 Label partitioning via K-means clustering: Bonsai repeatedly applies K-means to partition labels into disjoint subsets, using relatively large K values such as K ≥ 100 to produce shallow trees.The partitioning is performed in an input, output, or joint representation space.
- 2.2 Label partitioning via K-means clustering: Unlike Parabel’s binary balanced partitions, Bonsai permits varied cluster sizes and lets labels, including diverse tail labels, occupy separate data-dependent clusters.The method removes the balancedness constraint and supports K-way partitioning.
- 2.4 Prediction error propagation in shallow versus deep trees: With large branching factors, Bonsai makes trees shallower and reduces cascade error that compounds along long root-to-leaf paths.For example, repeated conditional probabilities of 0.95 over depth 16 yield approximately 0.46 overall probability.
3 Experimental Evaluation
The evaluation uses public XMC datasets spanning different scales and compares Bonsai variants with representative label-embedding, tree-based, and one-vs-all methods. Performance is measured with top-k ranking metrics for k ∈ {1, 3, 5}.
- 3 Experimental Evaluation: The experiments use publicly available XMC datasets from Amazon and Wikipedia, ranging from EURLex-4K to Amazon-3M.The datasets vary in numbers of labels, instances, and features.
- 3 Experimental Evaluation: Evaluation reports Precision@k and normalized Discounted Cumulative Gain, using percentage scores for k ∈ {1, 3, 5}.These metrics target the top-k recommendation, ranking, and advertising setting.
- 3 Experimental Evaluation: Bonsai-i, Bonsai-o, and Bonsai-io partition the input, output, and joint representation spaces, respectively.The variants are compared with methods from label-embedding, tree-based, and one-vs-all XMC families.
- 3 Experimental Evaluation: The comparison includes SLEEC, RobustXML, and LEML as label-embedding baselines.LEML is not compared explicitly because it gives much worse results, according to the passage.
- 3 Experimental Evaluation: FastXML, PFastXML, and Parabel represent the tree-based baselines used in the evaluation.PFastXML uses a propensity-scored variant intended to reward accurate tail-label predictions.
- 3 Experimental Evaluation: PD-Sparse and DiSMEC represent one-vs-all baselines using sparse or pruned per-label classifiers.The evaluation also notes that Bonsai is implemented in C++ on a 64-bit Linux system and learns three trees.
4 Experimental results
Across the experiments, Bonsai variants outperform Parabel, improve tail-label performance and coverage, and show that shallow trees are central to prediction quality. These gains come with higher training time, while prediction remains practical.
- Overall results: Bonsai variants show competitive performance, supporting generalized label representations and their potential for improved label partitioning.The results also suggest that these representations can be enriched further.
- Overall results: Bonsai consistently improves over Parabel across datasets, validating higher fanout and shallow-tree construction.The comparison averages results over five clustering initializations.
- Overall results: Higher label cardinality favors Bonsai-io’s joint representation, which further improves the strong performance of input-only Bonsai-i.This pattern is reported for Wikipedia-31K, Amazon-670K, and Amazon-3M.
- Tail-label performance: Approximately 6.7% relative improvement over Parabel on WikiLSHTC-325K prec wt@5 demonstrates stronger tail-label performance for Bonsai-i.The propensity-scored metrics prec wt@k and nDCG wt@k emphasize tail labels; Bonsai’s flexible, unbalanced partitions can separate tail labels from head labels.
- Unique label coverage: Bonsai discovers more correct unique labels than Parabel on coverage@k across five datasets.Coverage@k measures the proportion of normalized unique predicted labels relative to propensity-scored ground-truth labels.
- Tree depth: As tree depth increases, prediction accuracy tends to drop, while Bonsai-o substantially outperforms Parabel-o using the same label representation.These findings identify shallow architecture as an integral part of Bonsai’s success and associate it with reduced prediction error.
- Training and prediction time: Bonsai training takes approximately 2–3x longer than Parabel, but prediction typically remains in milliseconds for real-time applications.On WikiLSHTC-325K using three cores, Parabel takes one hour and Bonsai approximately three hours; training can be performed offline.
5 Conclusion
Bonsai learns shallow trees for label partitioning in extreme multi-label classification, using generalized label representations to improve prediction accuracy and tail-label coverage while retaining comparable training speed. Figure 6 compares Bonsai-o and Parabel-o on precision metrics across three benchmark datasets.
- Bonsai is a class of algorithms for learning shallow trees for label partitioning in extreme multi-label classification.
- Figure 6 compares prec@k and prec wt@k scores of Bonsai-o and Parabel-o over three benchmark datasets.
- Generalized label representations and shallow trees improve prediction accuracy and tail-label coverage while retaining training speed comparable to tree-based methods.The representations extend beyond input space, while shallow trees reduce error propagation in the tree cascade.