[Paper Review] X-BERT: eXtreme Multi-label Text Classification with BERT

Wei-Cheng Chang, Hsiang‐Fu Yu+3 more|arXiv (Cornell University)|May 7, 2019

Text and Document Classification Technologies7 citations

TL;DR

X-BERT proposes a fine-tuned BERT-based model for extreme multi-label text classification (XMC), leveraging joint document and label text to learn semantic label clusters and model label dependencies. It achieves state-of-the-art performance on a 0.5M-label Wiki dataset, reaching 67.80% precision@1, a 11.31% relative improvement over Parabel.

ABSTRACT

Extreme multi-label text classification (XMC) aims to tag each input text with the most relevant labels from an extremely large label set, such as those that arise in product categorization and e-commerce recommendation. Recently, pretrained language representation models such as BERT achieve remarkable state-of-the-art performance across a wide range of NLP tasks including sentence classification among small label sets (typically fewer than thousands). Indeed, there are several challenges in applying BERT to the XMC problem. The main challenges are: (i) the difficulty of capturing dependencies and correlations among labels, whose features may come from heterogeneous sources, and (ii) the tractability to scale to the extreme label setting as the model size can be very large and scale linearly with the size of the output space. To overcome these challenges, we propose X-BERT, the first feasible attempt to finetune BERT models for a scalable solution to the XMC problem. Specifically, X-BERT leverages both the label and document text to build label representations, which induces semantic label clusters in order to better model label dependencies. At the heart of X-BERT is finetuning BERT models to capture the contextual relations between input text and the induced label clusters. Finally, an ensemble of the different BERT models trained on heterogeneous label clusters leads to our best final model. Empirically, on a Wiki dataset with around 0.5 million labels, X-BERT achieves new state-of-the-art results where the precision@1 reaches 67:80%, a substantial improvement over 32.58%/60.91% of deep learning baseline fastText and competing XMC approach Parabel, respectively. This amounts to a 11.31% relative improvement over Parabel, which is indeed significant since the recent approach SLICE only has 5.53% relative improvement.

Motivation & Objective

• To address the challenge of modeling complex label dependencies in extreme multi-label text classification (XMC) with large label sets.
• To scale BERT-based models efficiently to extreme label settings where model size grows linearly with output space.
• To improve performance on XMC tasks by jointly modeling document and label text to induce semantic label clusters.
• To develop a scalable, fine-tuned BERT solution that outperforms existing deep learning and XMC-specific baselines.
• To demonstrate significant performance gains on large-scale XMC benchmarks using ensemble learning over heterogeneous label clusters.

Proposed method

• X-BERT constructs label representations by jointly encoding both document text and label text to capture semantic relationships.
• It fine-tunes BERT to model contextual interactions between input text and induced label clusters, enhancing label dependency learning.
• Label clusters are formed based on semantic similarity derived from joint document-label representations, enabling structured modeling of label correlations.
• The model employs an ensemble of BERT variants trained on different heterogeneous label clusters to improve generalization and robustness.
• Fine-tuning is performed end-to-end on the joint representation space to optimize for XMC metrics such as precision@1.
• The approach scales efficiently by reducing the effective label space through clustering while preserving semantic coherence.

Experimental results

Research questions

• RQ1 Can BERT be effectively fine-tuned for extreme multi-label text classification with label sets exceeding 500,000 labels?
• RQ2 How can label dependencies and correlations be modeled effectively in extreme multi-label settings?
• RQ3 Can joint encoding of document and label text improve semantic clustering of labels and downstream classification performance?
• RQ4 What is the performance gain of using ensemble BERT models over heterogeneous label clusters in XMC?
• RQ5 How does X-BERT compare to state-of-the-art XMC methods like Parabel and fastText in terms of precision@1 on large-scale datasets?

Key findings

• X-BERT achieves a precision@1 of 67.80% on a Wiki dataset with approximately 0.5 million labels, setting a new state-of-the-art.
• The model improves by 11.31% in relative performance over Parabel, a strong competing XMC method, demonstrating substantial gains.
• The improvement over Parabel exceeds that of SLICE by more than double (5.53% relative improvement), highlighting X-BERT's effectiveness.
• The use of joint document-label encoding enables better semantic clustering of labels, which enhances label dependency modeling.
• Ensembling multiple BERT models trained on heterogeneous label clusters significantly boosts performance compared to single-model baselines.
• X-BERT successfully scales BERT to extreme label settings by leveraging label clustering and fine-tuning, overcoming the linear model size issue.