Source-linked AI summary

LightXML: Transformer with Dynamic Negative Sampling for High-Performance Extreme Multi-label Text Classification

Ting Jiang, Deqing Wang, Leilei Sun, Huayi Yang, Zhengyang Zhao, Fuzhen Zhuang

arXiv:2101.03305v1cs.CLcs.LG

TL;DR

Extreme multi-label classification must handle huge label sets while existing methods incur high computational costs and rely on static negative sampling. LightXML uses a single end-to-end transformer with generative cooperative networks and dynamic negative sampling. It reports best results across five benchmark datasets with lower complexity and smaller models, including a 72% model-size reduction on Amazon-670K versus AttentionXML.

  • Problem

    XMC methods face very large label sets, multiple models or training stages, and static negative sampling in label ranking.

  • Method

    LightXML fine-tunes a single transformer end to end and uses generative cooperative networks to recall and rank labels with dynamic negative sampling.

  • Results

    LightXML achieves the best results on five benchmark datasets with smaller model size and lower computational complexity than current state-of-the-art methods.

  • Takeaways & Limitations

    LightXML supports efficient large-scale XMC, finding relevant labels from more than 0.5 million labels in 5 milliseconds with raw text as input.

  • Takeaways & Limitations

    The authors do not compare LightXML with X-Transformer because reproducing X-Transformer is difficult and inefficient.

Abstract

from arXiv · show

Extreme Multi-label text Classification (XMC) is a task of finding the most relevant labels from a large label set. Nowadays deep learning-based methods have shown significant success in XMC. However, the existing methods (e.g., AttentionXML and X-Transformer etc) still suffer from 1) combining several models to train and predict for one dataset, and 2) sampling negative labels statically during the process of training label ranking model, which reduces both the efficiency and accuracy of the model. To address the above problems, we proposed LightXML, which adopts end-to-end training and dynamic negative labels sampling. In LightXML, we use generative cooperative networks to recall and rank labels, in which label recalling part generates negative and positive labels, and label ranking part distinguishes positive labels from these labels. Through these networks, negative labels are sampled dynamically during label ranking part training by feeding with the same text representation. Extensive experiments show that LightXML outperforms state-of-the-art methods in five extreme multi-label datasets with much smaller model size and lower computational complexity. In particular, on the Amazon dataset with 670K labels, LightXML can reduce the model size up to 72% compared to AttentionXML.

Introduction

XMC must select relevant labels from extremely large label sets, creating computational challenges for existing methods. LightXML addresses model complexity and static negative sampling through end-to-end training and generative cooperative networks with dynamic sampling.

  • XMC identifies the most relevant labels for each text from an extremely large label set, with applications including Wikipedia tagging and ecommerce advertising.
  • Existing semantic-feature methods improve accuracy but require substantial computation, including multiple models or training stages for large datasets.AttentionXML needs four separate models for datasets such as Amazon-670K, while X-Transformer uses a large transformer for label recall.
  • Static negative sampling restricts label ranking to selected negatives and can make training difficult to converge.
  • LightXML fine-tunes a single transformer end to end and uses generative cooperative networks for label recalling and ranking.The framework uses transformer representations, a generator for recalling labels, and a discriminator for distinguishing positive from recalled labels.
  • Dynamic negative sampling lets label ranking learn from easy to hard negatives and avoid overfitting.
  • LightXML achieves the best results on five benchmark datasets while using smaller models and lower computational complexity than current state-of-the-art methods.

Related work

Related XMC methods trade off accuracy, scalability, model size, and computational cost across traditional and deep-learning approaches. Deep models improve text representations but remain challenged by very large label spaces and resource constraints.

  • XMC methods broadly use sparse traditional features or raw text processed by deep learning models.
  • One-vs-all methods: One-vs-all methods treat each label as an independent binary classification problem but retain expensive computation and large model sizes.
  • Tree-based methods: Tree-based methods reduce one-vs-all complexity by organizing labels or sparse features into hierarchical trees.
  • Embedding-based methods: Embedding-based methods compress the label space, but compression inevitably loses information and can reduce accuracy.
  • Deep learning methods: XML-CNN uses end-to-end CNN text representations and a bottleneck layer, but scoring all labels with a fully connected layer is difficult at large scale.
  • Deep learning methods: AttentionXML combines raw-text modeling with a probabilistic label tree but requires several models and has slower prediction and larger overall model size.
  • Deep learning methods: X-Transformer uses deep transformers for label-cluster matching and linear classifiers for ranking, but its computational complexity and model size limit feasibility.

Methodology

LightXML combines transformer text representations with generative cooperative networks that recall and rank labels end to end. Its generator dynamically samples label candidates, while the discriminator learns to distinguish positive from negative labels.

  • Problem formulation: LightXML formulates XMC as learning f(x_i) scores for all labels, then reduces computation by ranking only recalled labels.The target is to assign high scores to labels with y_il = 1 and obtain top-K predictions.
  • Framework: The framework clusters labels, embeds raw text with a transformer, and feeds the shared representation to label recalling and ranking.Label clustering uses a two-layer probabilistic label tree with balanced recursive partitioning.
  • Generative cooperative networks: Generative cooperative networks use a generator to sample label clusters and a discriminator to distinguish positive from negative labels.The generator supports label recalling, while the discriminator scores individual candidate labels.
  • Label recalling: During training, positive labels are added to generator-produced candidates so the ranking model learns to separate them from sampled negatives.At prediction time, the recalled subset is not modified and may omit some positive labels.
  • Dynamic negative sampling: Dynamic negative sampling resamples negatives for each training instance and progresses from easy to difficult distinctions during generator fitting.The paper states that this helps the discriminator converge and avoid overfitting to fixed negatives.
  • Label ranking: The discriminator embeds candidate labels and uses a hidden bottleneck to reduce label-related model size from O(L × 5k) to O((L + 5k) × b).The bottleneck also lets the model fit within limited GPU memory.
  • Training: LightXML adds generator and discriminator losses directly, allowing the transformer to receive gradients from both recalling and ranking.This supports end-to-end training instead of separately handling the two stages.

Experiments

Experiments use five widely used XMC benchmark datasets and evaluate predictions with precision at k. P@k measures the percentage of accurate labels among the top-k scored labels.

  • Datasets: Five XMC benchmarks are evaluated: Eurlex-4K, Wiki10-31K, AmazonCat-13K, Wiki-500K, and Amazon-670K.Sparse text features used for clustering are included in the datasets.
  • Evaluation measures: P@k measures the percentage of accurate labels among the top k scored labels.The metric is defined using the prediction vector and the indices of its k highest-scoring elements.

Baseline

The baseline set spans one-vs-all, label-tree, extreme-text, and deep-learning approaches. Reported baseline results come from prior studies.

  • Baselines: Baselines include DiSMEC, Parabel, ExtremeText, Bonsai, XML-CNN, AttentionXML, and X-Transformers.These methods cover one-vs-all, label-tree, and deep-learning approaches.

Experiment Settings

The experiments use five benchmark datasets and compare LightXML with existing XMC methods under accuracy, scalability, and computational-resource considerations.

  • LightXML uses end-to-end training, while the implementation applies dropout, stochastic weight averaging, AdamW, and automatic mixed precision.
  • LightXML outperforms all methods on four datasets and achieves better accuracy than X-Transformer on all datasets.
  • Table 3 compares LightXML with state-of-the-art XMC methods across the five benchmark datasets.

Performance on single model

The single-model evaluation compares LightXML-1 with AttentionXML-1 on Amazon-670K and Wiki-500K. LightXML-1 achieves better accuracy than AttentionXML-1.

  • LightXML-1 shows better accuracy than AttentionXML-1 in the single-model comparison.The comparison uses results on Amazon-670K and Wiki-500K.

Effect of the dynamic negative sampling

The ablation compares dynamic negative sampling with two static alternatives. Dynamic negative sampling outperforms both static methods while requiring less training time.

  • Two static baselines are evaluated: S uses one text representation, while BS adds a separately fine-tuned text representation.
  • Dynamic negative sampling outperforms both static negative sampling methods in the comparison.
  • Both static negative sampling methods take longer to train than dynamic negative sampling.
  • BS performs better than S when using an additional text representation.

Effect of the multi layers text representation

The experiment compares multi-layer and single-layer text representations on Wiki-500K and Amazon-670K. Multi-layer representations accelerate convergence and improve final P@5 accuracy.

  • Multi-layer and single-layer text representations are compared using training loss on Wiki-500K and Amazon-670K.
  • Multi-layer text representation reaches the same training loss as the single-layer representation using half the total epochs.The reported comparison concerns convergence during training.
  • Multi-layer text representation improves final P@5 accuracy by more than 1%.

Computation Time and Model Size

LightXML is evaluated against AttentionXML on training time, prediction speed, and model size, while X-Transformer is excluded because of its costly and difficult reproduction. The reported results show substantially faster prediction and smaller models for LightXML.

  • Comparison scope: X-Transformer takes more than 35 hours to train on Wiki-500K using eight Tesla V100 GPUs, and reproducing it would require more than 100 hours.The authors therefore do not compare LightXML with X-Transformer.
  • Evaluation measures: Table 6 reports overall training hours, per-sample prediction time in milliseconds, and model size in gigabytes.Ttrain denotes training hours, Stest denotes milliseconds per sample, and M denotes model size in GB.
  • Prediction speed: LightXML finds relevant labels from more than 0.5 million labels in 5 milliseconds using raw text as input.This prediction-speed comparison uses the same hardware for AttentionXML and LightXML.
  • Model size: LightXML reduces model size by 72% on Amazon-670K and 52% on Wiki-500K compared with AttentionXML.Both comparisons use one Tesla V100 GPU and the same hardware for the two methods.
  • Training time: LightXML and AttentionXML train quickly relative to X-Transformer, saving more than three times the training time.The passage reports this comparison after noting X-Transformer's substantially higher training cost.

Conclusion

LightXML combines a transformer with generative cooperative networks for end-to-end XMC training and dynamically samples negative labels. Experiments report high efficiency and best accuracy on large-scale datasets while reducing model size by more than 50%.

  • Conclusion: LightXML combines a transformer with generative cooperative networks for end-to-end training in extreme multi-label classification.The model uses these networks to support label recalling and ranking.
  • Conclusion: Dynamic negative sampling makes LightXML more robust in prediction by resampling negative labels through generative cooperative networks.The conclusion presents dynamic negative sampling as a component of the model's prediction robustness.
  • Conclusion: LightXML achieves high efficiency and best accuracy on large-scale datasets, with experiments executable on a single GPU card in reasonable time.The conclusion attributes this result to the reported model design and experiments.
  • Conclusion: LightXML reduces model size by more than 50% while retaining accuracy compared with current state-of-the-art deep learning methods.The conclusion links the reduction to redundant parameters in existing deep learning methods.
Loading 2101.03305v1…