Source-linked AI summary
AttentionXML: Label Tree-based Attention-Aware Deep Model for High-Performance Extreme Multi-Label Text Classification
Ronghui You, Zihan Zhang, Ziye Wang, Suyang Dai, Hiroshi Mamitsuka, Shanfeng Zhu
TL;DR
XMTC must assign multiple relevant labels from extremely large label sets, while prior approaches struggle with label-specific context and scalability. AttentionXML combines raw-text BiLSTM multi-label attention with a shallow, wide probabilistic label tree. Across six benchmarks, it outperformed competing methods and showed particular strength on long-tailed labels.
Problem
XMTC involves huge label sets and tail labels with few positive examples, while existing methods can miss label-specific text context or scale poorly.
Method
AttentionXML combines raw-text BiLSTM multi-label attention with a shallow, wide probabilistic label tree and level-wise training.
Results
AttentionXML outperformed all competing methods across six benchmark datasets, including extreme-scale datasets, and achieved the best tail-label performance among label tree-based methods.
Takeaways & Limitations
Label-specific attention and a shallow, wide PLT provide a supported approach for accurate and efficient XMTC at extreme label scales, including tail-label prediction.
Abstract
from arXiv · showhide
Extreme multi-label text classification (XMTC) is an important problem in the era of big data, for tagging a given text with the most relevant multiple labels from an extremely large-scale label set. XMTC can be found in many applications, such as item categorization, web page tagging, and news annotation. Traditionally most methods used bag-of-words (BOW) as inputs, ignoring word context as well as deep semantic information. Recent attempts to overcome the problems of BOW by deep learning still suffer from 1) failing to capture the important subtext for each label and 2) lack of scalability against the huge number of labels. We propose a new label tree-based deep learning model for XMTC, called AttentionXML, with two unique features: 1) a multi-label attention mechanism with raw text as input, which allows to capture the most relevant part of text to each label; and 2) a shallow and wide probabilistic label tree (PLT), which allows to handle millions of labels, especially for "tail labels". We empirically compared the performance of AttentionXML with those of eight state-of-the-art methods over six benchmark datasets, including Amazon-3M with around 3 million labels. AttentionXML outperformed all competing methods under all experimental settings. Experimental results also show that AttentionXML achieved the best performance against tail labels among label tree-based methods. The code and datasets are available at http://github.com/yourh/AttentionXML .
1 Introduction
XMTC assigns multiple relevant labels to each text from label sets reaching millions, creating major computational challenges. AttentionXML addresses limitations in existing deep and label tree-based methods through label-specific attention and a shallow, wide PLT.
- Task and challenges: XMTC tags each text with multiple relevant labels from extremely large label sets, unlike multi-class classification.Applications include e-commerce item categorization, web page tagging, and news annotation.
- Task and challenges: Extreme-scale XMTC requires efficient learning with huge numbers of samples and labels, including tail labels with very few positive examples.
- Existing limitations: Label tree methods partition labels with a PLT, but Parabel’s deep balanced tree can accumulate errors and group dissimilar tail labels together.These design choices deteriorate overall and tail-label performance according to the paper’s motivation.
- Proposed approach: AttentionXML uses raw text, BiLSTM representations, and multi-label attention to capture label-relevant text while representing each label differently.The paper motivates this design as especially helpful for tail labels.
- Proposed approach: AttentionXML also uses a shallow and wide PLT to handle extreme-scale datasets, addressing scalability and tail-label challenges.
2 AttentionXML
AttentionXML combines a shallow, wide probabilistic label tree with level-wise attention-aware deep models. Its BiLSTM and label-specific attention process raw text, while the PLT reduces the candidate-label burden for scalable training and prediction.
- PLT construction and training: AttentionXML first builds a shallow and wide PLT, then trains an attention-aware deep model for each tree level in a top-down manner.Each level remains a multi-label classification problem, and deeper models are trained on candidate nodes selected from the preceding level.
- Probabilistic label tree: The PLT assigns original labels to leaves and pseudo-labels to internal nodes, estimating conditional node probabilities whose products yield marginal label probabilities.The marginal probability follows the chain rule along the path from a node to the root.
- PLT construction and training: AttentionXML compresses an initially clustered binary tree into a K-way tree with height H by repeatedly removing intermediate nodes.The initial tree uses balanced k-means over normalized label representations and a cluster-size threshold M.
- Attention-aware deep model: The attention-aware model uses word representations, a BiLSTM, multi-label attention, fully connected layers, and an output layer.The BiLSTM captures left- and right-side context, while attention forms label-specific combinations of contextual vectors.
- Attention-aware deep model: Each label has distinct attention parameters, allowing the model to focus on different relevant text contexts for different labels.Shared fully connected and output-layer parameters reduce model size and help avoid overfitting.
- Complexity: With a PLT, AttentionXML reduces the first model’s label size to L/K^H and deeper models’ candidate labels to C × K.This lowers time and space costs relative to applying multi-label attention across all L labels, enabling operation with limited GPU memory.
3 Experimental Results
Across six XMTC benchmarks, AttentionXML achieved the strongest reported performance while combining label-tree scalability with attention-based text representations. Its ensemble improved accuracy further, whereas the single-tree model offered a more efficient accuracy–resource trade-off.
- Datasets: Six benchmarks span large-scale datasets with 4K–30K labels and extreme-scale datasets with 500K–3M labels.The extreme-scale set includes Amazon-670K, Wiki-500K, and Amazon-3M.
- Overall comparison: AttentionXML with a three-PLT ensemble outperformed all eight competing methods by P@k.At P@5, it was at least 4% above the second-best method across datasets and over 17% higher on Wiki-500K.
- Overall comparison: AttentionXML-1 outperformed all eight competing methods except in one case, with especially strong gains on longer-text datasets.At P@5, it reached 44.41 on Wiki-500K, 68.78 on Wiki10-31K, and 61.10 on EUR-Lex.
- Extreme-scale performance: On the three extreme-scale datasets, AttentionXML was the best competing method, while AttentionXML-1 was faster than DiSMEC and used less memory than Bonsai.The ensemble produced more significant improvements than AttentionXML-1 while retaining those efficiency advantages.
- Tail-label performance: AttentionXML outperformed Parabel and Bonsai in PSP@k on all datasets, indicating strong performance on tail labels.The paper attributes this result to its shallow and wide PLT and multi-label attention.
- Ensemble analysis: More PLTs improved prediction accuracy but increased training and prediction time, creating an accuracy–time trade-off.AttentionXML can therefore trade ensemble size for computational cost.
4 Conclusion
AttentionXML combines multi-label attention with a shallow and wide PLT for scalable XMTC, including datasets with millions of labels. It outperformed competing methods across six benchmarks and showed an advantage on long-tailed labels.
- AttentionXML combines multi-label attention, which captures text parts relevant to each label, with a shallow and wide PLT for handling millions of labels.
- AttentionXML was evaluated against eight state-of-the-art methods across six benchmark datasets, including three extreme-scale datasets.
- AttentionXML outperformed all competing methods across all six datasets, particularly on datasets with long texts.
- AttentionXML showed an advantage in predicting long-tailed labels.
C.1 1-vs-All Methods
1-vs-All methods train an independent classifier for every label, creating substantial computational and model-size challenges at extreme scale. Sparse learning and distributed parallelization reduce these costs in later approaches.
- 1-vs-All methods train one classifier independently for each label.
- Their computational complexity and model size can become very high because XMTC contains extremely large numbers of labels and instances.
- PD-Sparse uses L1-penalized sparse learning to obtain sparse primal and dual solutions without sacrificing predictor expressive power.
- DiSMEC uses distributed computing and double-layer parallelization, then prunes near-zero coefficients to make the model thousands of times smaller.
C.2 Embedding-based Methods
Embedding-based methods address huge label spaces by compressing labels into lower-dimensional representations and decompressing them for prediction. SLEEC and AnnexML differ mainly in how they preserve label relationships and perform prediction.
- Embedding-based methods compress high-dimensional label vectors into lower-dimensional embeddings for training and decompress them during prediction.
- Their main differences concern the compression function f_C and decompression function f_D.
- SLEEC learns embeddings that capture nonlinear label correlations while approximately preserving pairwise distances among nearby label vectors.
- AnnexML extends SLEEC by addressing unlabeled clustering, distance-free KNN prediction, and slow prediction through a KNN graph in embedding space.
C.3 Tree-based Methods
Tree-based methods recursively partition instances using features and typically use ensembles of trees, but their selection process can reduce performance. FastXML and PfastreXML use ranking-based objectives at tree nodes.
- Tree-based methods recursively partition instances by features, producing classifiers at leaves with only a few active labels.
- Many tree-based methods form ensembles by randomly sampling feature subsets at each node.
- A stated disadvantage is low performance caused by the selection process in tree construction.
- FastXML learns hyperplanes to split instances and optimizes an nDCG-based ranking loss at each node.
- PfastreXML retains FastXML’s architecture but replaces its objective with a propensity-scored objective function.
C.4 Deep learning-based Methods
Sequence-to-sequence methods use attention-based architectures with an RNN encoder to represent source text and predict labels.
- MLC2Seq, SGM, and SU4MLC are pioneering sequence-to-sequence approaches using attention-based architectures.
- These architectures encode source-text representations with an RNN encoder before predicting labels.
- Figure 3 visualizes AttentionXML’s attention from a Wiki10-31K GMail Driver entry to the label “gmail”.
D.1 Evaluation Metrics
The evaluation uses Precision at k and normalized Discounted Cumulative Gain at k, with the latter accounting for the ranking of top predictions.
- P@k and N@k are used because both are widely used for evaluating multi-label classification methods.
- P@k measures precision among the top k predicted labels using the true binary label vector.
- N@k evaluates ranking quality by considering the order of the top k predictions, unlike P@k.
- P@1 and N@1 are identical.
D.2 Performance Results
AttentionXML is compared with competing methods across six benchmark datasets, while additional experiments examine tree hyperparameters and attention behavior. Smaller cluster sizes and tree heights improve performance but increase time cost, and attention helps rank a relevant label highly in a representative case.
- AttentionXML and seven competing methods are compared across six benchmark datasets.
- Impact of height and maximum cluster size: Smaller maximum cluster sizes improve AttentionXML’s performance across datasets, especially for tail labels.
- Impact of height and maximum cluster size: Smaller tree heights improve performance, but smaller heights and cluster sizes require more time.
- Effectiveness of Attention: In a Wiki10-31K example, attention ranks the true label “gmail” at the top, whereas the model without multi-label attention ranks it below 100th.
- Effectiveness of Attention: The attention mechanism assigns high scores to “Gmail”, “e-mail”, and “attachments”, which are relevant to the label “gmail”.