PaperlayerView exact revision on arXiv

Source-linked AI summary

Fast Multi-Resolution Transformer Fine-tuning for Extreme Multi-label Text Classification

Jiong Zhang, Wei-cheng Chang, Hsiang-fu Yu, Inderjit S. Dhillon

arXiv:2110.00685v2cs.LGcs.AIcs.IRstat.ML

TL;DR

XMC asks models to identify relevant labels from enormous collections, and transformer-based methods improve text representation but remain costly to fine-tune at large label scales. XR-Transformer recursively fine-tunes across multi-resolution objectives. It reports state-of-the-art results with less training time, including improved Amazon-3M precision and a 20x speedup over X-Transformer.

  • Problem

    XMC requires selecting relevant labels from enormous collections, while transformer-based XMC methods still take lengthy fine-tuning despite their improved text representations.

    Source: §1 · pp. 1–2

  • Method

    XR-Transformer recursively fine-tunes pre-trained transformers on coarse-to-fine, multi-resolution objectives, using bootstrapping between layers.

    Source: §1 · p. 2, §4 · p. 4, else · p. 6

  • Results

    XR-Transformer establishes state-of-the-art results on public XMC datasets with less training time; on Amazon-3M, Prec@1 improves from 51.20% to 54.04%.

    Source: §1 · p. 2, §6 · p. 10

  • Takeaways & Limitations

    The experiments support multi-resolution objectives and cost-sensitive learning as an approach that accelerates transformer fine-tuning while improving XMC results.

    Source: §6 · p. 10

Abstract

from arXiv · show

Extreme multi-label text classification (XMC) seeks to find relevant labels from an extreme large label collection for a given text input. Many real-world applications can be formulated as XMC problems, such as recommendation systems, document tagging and semantic search. Recently, transformer based XMC methods, such as X-Transformer and LightXML, have shown significant improvement over other XMC methods. Despite leveraging pre-trained transformer models for text representation, the fine-tuning procedure of transformer models on large label space still has lengthy computational time even with powerful GPUs. In this paper, we propose a novel recursive approach, XR-Transformer to accelerate the procedure through recursively fine-tuning transformer models on a series of multi-resolution objectives related to the original XMC objective function. Empirical results show that XR-Transformer takes significantly less training time compared to other transformer-based XMC models while yielding better state-of-the-art results. In particular, on the public Amazon-3M dataset with 3 million labels, XR-Transformer is not only 20x faster than X-Transformer but also improves the Precision@1 from 51% to 54%.

1 Introduction

XMC has enormous label spaces that challenge accurate, efficient training, while transformer-based methods improve text representation but remain costly to fine-tune. XR-Transformer recursively fine-tunes across coarse-to-fine objectives, with reported gains in both training time and Amazon-3M precision.

1 Introduction

  • XMC applications require predicting relevant labels from collections ranging from tens of thousands to millions, making accurate and efficient training difficult.Prior scalable methods shortlist labels using label correlations, but often rely on bag-of-words or pooled pre-generated embeddings.Source: §1 · p. 1
  • Transformer XMC methods improve semantic text representation, but truncation can be needed to reduce GPU memory use and maintain efficiency.X-Transformer, for example, truncates inputs to the first 128 tokens.Source: §1 · p. 1
  • X-Transformer and LightXML cluster labels to reduce fine-tuning cost, but information loss from aggregation leads them to use small clusters and still train for a long time.X-Transformer took 23 and 25 days on Amazon-3M and Wiki-500K, respectively, using eight Nvidia V100 GPUs.Source: §1 · p. 2
  • XR-Transformer recursively fine-tunes a pre-trained transformer on coarse-to-fine, multi-resolution objectives related to the original XMC task.The approach is motivated by multi-resolution learning and curriculum learning.Source: §1 · p. 2
  • Prec@1 rises from 51.20% to 54.04% on Amazon-3M, while training time falls from 23 days to 29 hours on the same hardware.The paper also reports significantly lower training time and better state-of-the-art results across six public XMC datasets.Source: §1 · p. 2

2 Related Works

Prior XMC approaches address large output spaces through label partitioning, shallow or recurrent text encoders, and transformer fine-tuning. Their trade-offs include scalability limits, short-text specialization, and transformer training cost.

Sparse Linear Models with Partitioning Techniques

  • Sparse linear XMC methods use fixed input representations and partitioning or surrogate losses; one-versus-all approaches also use parallelism and weight truncation to reduce complexity.Inference cost linear in output size can be improved by label partitioning or approximate nearest-neighbor indexing.Source: §2 · p. 2
  • Tree-based sparse linear models partition labels into hierarchical trees, reducing inference complexity to logarithmic in the output space.Embedding-based methods instead use approximate nearest-neighbor techniques, including graph-based methods and product quantization variants.Source: §2 · p. 2

Shallow Embedding-based Methods

  • Neural XMC methods include CNNs, BiLSTMs with label-aware attention, and shallow embedding-based MLPs trained on label-cluster subproblems.Shallow embedding methods show competitive performance mainly on short-text XMC problems.Source: §2 · pp. 2–3

Deep Transformer Models

  • Transformer XMC methods include clustered-label training, adaptive imbalanced clusters, and end-to-end training with dynamic negative sampling, but require larger models and longer training.X-Transformer uses transformer encoders followed by sparse TF-IDF features in linear one-versus-all models.Source: §2 · p. 3

3 Background Material

XMC learns document–label relevance scores and retrieves the highest-scoring labels. Partitioning and clustering make large output spaces tractable, but cluster size creates a trade-off between fine-tuning efficiency and label resolution.

3 Background Material

  • XMC learns a relevance score for each document–label pair, then predicts the k labels with the largest scores.Training examples pair documents with multi-hot label vectors indicating relevant labels.Source: §3 · p. 3
  • An OVA model scores each label using its weight vector and a text vectorizer, but training and inference without sampling become prohibitive as L grows.The vectorizer may be deterministic, such as BOW or TF-IDF, or have learnable parameters.Source: §3 · p. 3
  • Partitioning groups labels into clusters, trains a shortlisting model to match documents to clusters, then trains a classifier on shortlisted labels.This three-stage framework is used to reduce the output space considered during training and inference.Source: §3 · p. 3
  • Large clusters can blur the signals of their constituent labels, while small fixed clusters leave roughly L/B clusters for transformer shortlisting and prolong training.The figure contrasts faster but potentially lower-quality low-resolution tasks with higher-resolution tasks that take longer.Source: §3 · pp. 3–4, Figure 1 · p. 4

4 Proposed Method: XR-Transformer

XR-Transformer recursively fine-tunes a transformer across coarse-to-fine label objectives, using a hierarchical label tree to make progressively larger output spaces tractable. It combines statistical text features with transformer embeddings and reuses representations across training stages.

4 Proposed Method: XR-Transformer

  • Recursive fine-tuning applies the same framework to progressively larger output spaces, until the shortlisting problem reaches a reasonably small size.The number of considered labels is O(B log_B(L)) under recursive shortlisting.Source: §4 · p. 4
  • The hierarchical label tree supplies multi-resolution objectives, with coarser label vectors formed by max-pooling labels in the hierarchy.The paper frames XMC labels as a binary vector and propagates true labels from child clusters to coarser layers.Source: §4 · pp. 4–5
  • Relevance matrices retain information lost when positive child labels are merged, assigning greater relevance to clusters with more positive children.The matrices are recursively induced by the hierarchical label tree for cost-sensitive learning.Source: §4 · p. 5
  • Training shortlists each layer using both model-selected top-k parent clusters and positive parent clusters, approximating beam search while balancing exploration and exploitation.Only shortlisted labels are included in each instance's training objective.Source: §4 · p. 5
  • Combining statistical text features with transformer embeddings complements information lost through input truncation while retaining semantic representations.Statistical features can use the whole input, whereas transformer inputs are commonly truncated for efficiency.Source: §4 · p. 5

else

  • Training begins with a preliminary tree from raw statistical features, then fine-tunes recursively from low to high resolution, initializing each layer's transformer from the previous layer.The initial transformer weights at the root are the pretrained weights.Source: else · p. 6
  • After fine-tuning, XR-Transformer refines the tree using combined text features and predicts through beam search, computing the transformer embedding once per instance.Its inference complexity is O(T_dnn + kd log(L)).Source: else · p. 7
  • Unlike AttentionXML, XR-Transformer progressively trains one transformer across tree layers rather than using a separate encoder per layer.The paper reports that this difference leads to longer inference time for AttentionXML because multiple encoders must be queried.Source: else · p. 7

5 Experimental Results

Across six public XMC benchmarks, XR-Transformer is compared with recent methods for Precision@k and training time. The reported results combine broad performance gains with substantially shorter training, alongside analyses of cost-sensitive learning, label resolution, and text representations.

5 Experimental Results

  • 14 out of 18 evaluation columns achieve new state-of-the-art results, with XR-Transformer outperforming competitive methods on large datasets.The comparison covers six public datasets.Source: Table 2 · p. 8, §5 · p. 9
  • On Amazon-3M, XR-Transformer is 20x faster than X-Transformer while achieving better P@k.In a single-model comparison, XR-Transformer also consistently outperforms LightXML in P@k and training time.Source: Table 3 · p. 8, §5 · p. 9, Table 4 · p. 9
  • On large datasets, Pre-Trained+TF-IDF yields only marginal improvement over TF-IDF, while XR-Transformer gains significantly; the authors attribute the main improvement to stronger neural semantic embeddings.Table 5 compares XR-Transformer with Pre-Trained and word2vec embeddings concatenated with TF-IDF features.Source: Table 5 · p. 9, §5 · p. 9
  • On most of the four datasets with the largest output spaces, cost-sensitive learning via aggregated labels improves performance and also benefits XR-Linear.Table 6 ablates cost-sensitive learning on the single XR-Transformer model.Source: Table 6 · p. 10, §5 · p. 10
  • Larger clusters, corresponding to lower label resolution, yield worse semantic features, while combining neural and TF-IDF features generally improves performance.Figure 2 compares BERT models fine-tuned at different label resolutions; Φcat is the normalized concatenation of Φtfidf and Φdnn.Source: Figure 2 · p. 10, §5 · p. 10

6 Conclusion and Future work

XR-Transformer combines multi-resolution objectives and cost-sensitive learning to accelerate fine-tuning of pretrained transformers for XMC. Experiments report state-of-the-art results on public XMC datasets with significantly less training time than earlier transformer-based methods.

6 Conclusion and Future work

  • XR-Transformer achieves new state-of-the-art results on public XMC datasets while taking significantly less training time than earlier transformer-based methods.Source: §6 · p. 10
  • The authors suggest that the approach could also apply to information retrieval and to other deep models such as CNNs and ResNets.Source: §6 · p. 10
  • Further study is needed to understand why coarse-to-fine learning may improve both training speed and quality; the authors hypothesize that learning at multiple scales yields more robust transformer learning.Source: §6 · p. 10

A.1 Method Overview

XR-Transformer combines text vectorizers, a recursive multi-resolution fine-tuning curriculum, and multi-layer ranking models. Its training builds a preliminary hierarchical label tree and recursively fine-tunes the transformer vectorizer on labels at multiple resolutions.

A.1 Method Overview

  • The architecture combines a learnable transformer semantic vectorizer, a deterministic statistical vectorizer, and multi-layer ranking models.Source: §A.1 · p. 15
  • The training pipeline builds a preliminary hierarchical label tree with hierarchical k-means, then recursively fine-tunes the transformer vectorizer on multi-resolution labels.Source: Figure 3 · p. 15
  • The ranking models are learned using refined multi-resolution labels, and the final prediction is produced together with a refined hierarchical label tree.Source: §A.1 · p. 15

A.2 Evaluation Metrics

The evaluation section introduces Precision@k and Recall@k as widely used XMC metrics and also reports propensity-scored precision for tail-label performance. The propensity-score metric uses the same application-specific parameter setting as AttentionXML across datasets.

A.2 Evaluation Metrics

  • Precision@k and Recall@k are the widely used XMC metrics, defined with respect to ground-truth labels and predicted-label rank.Source: §A.2 · p. 15
  • Propensity-scored precision is reported to compare performance on tail labels.Source: §A.2 · p. 15
  • The propensity metric uses application-specific parameters A and B, with the same setting as AttentionXML used across datasets.Source: §A.2 · p. 15

A.3 Model Hyperparameters

For each sub-XMC problem, XR-Transformer fine-tunes the transformer encoder with Adam, a linear learning-rate schedule, and a batch size of 32 per GPU (256 total). Dropout is 0.1, and bootstrapping is used except at the root layer.

A.3 Model Hyperparameters

  • Each sub-XMC problem uses Adam with a linear learning-rate schedule and a batch size of 32 per GPU, or 256 total.Source: §A.3 · p. 16
  • Dropout is set to 0.1, and bootstrapping is used for all fine-tuning problems except the root layer.Source: §A.3 · p. 16

A.3.1 Model Ensemble Setup

The appendix presents ensemble and training-time information alongside analyses of cost-sensitive learning, inference speed, propensity-scored precision, and potential societal impacts.

A.3.1 Model Ensemble Setup

  • The ensemble setup uses three transformer encoders for LightXML and a nine-model ensemble for X-Transformer, with XR-Transformer ensembles varying by dataset.XR-Transformer uses BERT, RoBERTa, and XLNet ensembles on Eurlex-4K, Wiki10-31K, and AmazonCat-13K; three BERT models on Wiki-500K, Amazon-670K, and Amazon-3M.Source: §A.3.1 · p. 16

A.4 More Empirical Results

  • Table 8 compares training time in hours for DNN-based methods producing the state-of-the-art results in Table 2; model labels indicate ensemble size.Source: Table 8 · p. 16

A.4.2 Inference Speed

  • Inference times are reported in milliseconds per sample; AttentionXML uses multiple GPUs on the three largest datasets because it may otherwise run out of memory on one GPU.Table 10 records most methods on a single V100 GPU with batch size 1; starred results use eight GPUs.Source: §A.4.2 · p. 16, Table 10 · p. 17

A.4.1 Cost-sensitive Learning

  • Cost-sensitive learning via recursive label aggregation benefits both XR-Linear and XR-Transformer, with better performance on most datasets.Table 9 reports Precision@k and Recall@k for single models with and without cost-sensitive learning.Source: §A.4.1 · p. 16, Table 9 · p. 17

A.4.3 Propensity-score Precision

  • The appendix reports PSP@1, PSP@3, and PSP@5 for four methods across three large public datasets using AttentionXML's setting.Table 11 includes PfastreXML, Parabel, AttentionXML, and XR-Transformer.Source: Table 11 · p. 17, §A.4.3 · p. 17

A.5 Potential Negative Societal Impacts

  • The authors state that the paper focuses on accelerating XMC training algorithms and that, to their knowledge, it poses no negative societal impacts.Source: §A.5 · p. 17
Loading 2110.00685v2…