Skip to content

Don’t Pay Attention, PLANT It: Pretraining Attention via Learning-to-Rank

Conference: NeurIPS 2026
arXiv: 2410.23066
Area: Medical LLM
Keywords: extreme multi-label classification, attention pretraining, learning-to-rank, mutual information gain, rare labels

TL;DR

Addressing the failure of random label attention initialization on skewed long-tail classes in Extreme Multi-Label Text Classification, this paper introduces PLANT, a plug-and-play framework that pretrains label attention via Mutual Information Gain and Learning-to-Rank, substantially outperforming larger LLM baselines and prior SOTA models on rare labels.

Background & Motivation

Extreme Multi-Label Text Classification (XMTC) requires assigning the most relevant subset of labels from a massive label taxonomy (often tens or hundreds of thousands of classes) to long text documents. This setting is ubiquitous in clinical diagnosis coding (ICD coding), legal statute categorization, and e-commerce tag recommendation. To navigate such vast and skewed label spaces, modern high-performing architectures almost universally incorporate multi-label attention mechanisms: assigning dedicated query embeddings to each label to act as directional spotlights across lengthy input token sequences, dynamically extracting the salient token cues for each class.

However, standard practice initializes these label attention weights randomly, forcing the model to learn the token-level attention distribution for every label from scratch during end-to-end task supervision. While high-frequency head labels possess sufficient training instances for gradient-based convergence, rare tail labels (which frequently appear fewer than 5–10 times in datasets like MIMIC-III and MIMIC-IV) fail to learn meaningful semantic structures. Diagnostic experiments demonstrate that under random initialization, semantically related rare codes exhibit widely dispersed attention patterns (mean pairwise cosine similarity of only ~0.75), and prolonged training leads to overfitting rather than generalization, causing near-zero Macro-F1 on tail labels.

The core tension is that learning label-specific attention from scratch is fundamentally data-intensive, yet token-label informativeness can be estimated prior to downstream task tuning. Core idea: pretrain the label attention layer as an independent Learning-to-Rank module guided by Mutual Information Gain over training co-occurrences, planting well-structured attention priors before end-to-end task training.

Method

Overall Architecture

PLANT decomposes model training into a two-stage paradigm: Stage 1 ("PLANTing Attention") freezes or warms up the text encoder and optimizes the multi-head attention module using a Learning-to-Rank (L2R) objective with Mutual Information Gain (MIG) ground truth; Stage 2 ("Leveraging Attention") loads the pretrained label attention weights and embedding queries into the full model, optimizing the entire architecture end-to-end using focal loss with hard negative mining and a progressive unfreezing schedule.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Input text x and label taxonomy L"] --> B["Text encoder and label query embeddings E"]
    B --> C["Stage 1: Relevance-guided ranking pretraining<br/>MIG scores over top-k tokens, LambdaLoss on attention scores"]
    C --> D["Stage 2: Feature boosting and progressive unfreezing<br/>Differentiable feature scaling + Hard Negative Mining focal loss"]
    D --> E["Extreme multi-label classification logits"]

Key Designs

1. Relevance-guided ranking pretraining: bypassing random exploration with information-theoretic priors To resolve the failure of tail labels under random initialization, PLANT computes the offline Mutual Information Gain (MIG) \(r_{l,j}\) between each candidate label \(l\) and token \(t_j\) across the training corpus: $\(r_{l,j}=\sum_{(x,y)\in\{0,1\}^{2}}P(x,y)\log_2\left(\frac{P(x,y)}{P(x)\,P(y)}\right)\)$ where \(x\) and \(y\) denote binary occurrences of label \(l\) and token \(t_j\) across documents. Unlike raw frequency or co-occurrence, MIG downweights frequent corpus stopwords that appear spuriously across many labels. For each label, the top-\(k\) tokens by MIG form a candidate set \(\mathcal{T}_l\). The multi-head attention layer maps label queries \(\boldsymbol{Q}=\boldsymbol{E}\) and text hidden states \(\boldsymbol{K}=\boldsymbol{V}=\boldsymbol{H}_i\) to token-level attention scores \(\boldsymbol{S}\in\mathbb{R}^{|\mathcal{L}|\times n}\). In Stage 1, attention parameters \(\boldsymbol{W}_{\mathsf{attn}}\) are trained with a LambdaLoss ranking objective: $\(\mathcal{L}_{\textsf{rank}}^{(i)}(\boldsymbol{S})=-\sum_{l=1}^{|\mathcal{L}|}\sum_{j, h \in \mathcal{T}_l, r_{lj} > r_{lh}} |\Delta\mathsf{nDCG}@k|_{jh} \log_2\left(\frac{1}{1+e^{-\sigma(s_{lj}-s_{lh})}}\right)\)$ This objective explicitly aligns predicted attention weights with information-theoretic token importance, boosting intra-category cosine similarity among rare label groups from 0.75 to 0.985.

2. Differentiable feature boosting and adaptive pooling: amplifying delicate tail representations After the multi-head attention module produces attended text representations \(\boldsymbol{A}\in\mathbb{R}^{|\mathcal{L}|\times h}\), subtle signals from sparse tail labels risk being overshadowed by dominant head patterns in the high-dimensional projection. PLANT introduces learnable label-specific bias and scaling matrices \(\boldsymbol{B}_a, \boldsymbol{B}_m\in\mathbb{R}^{|\mathcal{L}|\times h}\), applying element-wise modulation: $\(\boldsymbol{A}_{\text{boost}}=\boldsymbol{B}_a+\boldsymbol{A}\odot\boldsymbol{B}_m\)$ An adaptive average pooling layer reduces the hidden dimensionality to \(p=128\), yielding \(\boldsymbol{P}_i\in\mathbb{R}^{|\mathcal{L}|\times p}\). Finally, a shared linear projection \(\boldsymbol{W}_c\in\mathbb{R}^{p\times 1}\) computes label logits \(\hat{\boldsymbol{y}}_i=\boldsymbol{P}_i\boldsymbol{W}_c\). This architecture allows the model to learn fine-grained emphasis or suppression per label channel, sharpening discrimination between confusing neighboring codes.

3. Progressive unfreezing with hard negative mining: preserving planted attention against extreme skew Transitioning directly into full end-to-end optimization can wash out the delicate attention structures seeded in Stage 1. PLANT adopts a phased, discriminative unfreezing strategy: the attention module and label embeddings remain strictly frozen for the first 5 epochs, allowing downstream classification heads to adapt; subsequent phases gradually unfreeze attention projection layers (lr = \(1\times 10^{-5}\)), intermediate layers (lr = \(5\times 10^{-6}\)), and all model parameters (lr = \(2\times 10^{-6}\)). Furthermore, to mitigate extreme class imbalance where thousands of negative classes dominate each sample, Stage 2 optimizes an HNM-augmented focal loss: keeping all positive labels while mining only the top-\(m\) (\(m=1000\)) hardest negative labels by predicted probability.

Loss & Training

Stage 1 minimizes the LambdaLoss ranking objective (Equation 2) using AdamW with a cosine learning-rate schedule (peak \(\eta=5\times 10^{-4}\), 10% warmup) for 10 epochs. Stage 2 optimizes the subset focal loss with label smoothing (\(\epsilon=0.1\)): $\(\mathcal{L}_{\text{focal}}^{(i)}=-\frac{1}{|\mathcal{S}_i|}\sum_{l\in\mathcal{S}_i}\left[\tilde{y}_{il}(1-\sigma(\hat{y}_{il}))^\gamma\log\sigma(\hat{y}_{il})+(1-\tilde{y}_{il})(\sigma(\hat{y}_{il}))^\gamma\log(1-\sigma(\hat{y}_{il}))\right]\)$ where focusing parameter \(\gamma=2\), smoothed target \(\tilde{y}_{il}=(1-\epsilon)y_{il}+\epsilon(1-y_{il})\), and \(\mathcal{S}_i\) denotes the set of all positive labels alongside the top-\(1000\) hard negative classes. All experiments run on 8×A100-80GB GPUs with QLoRA (4-bit NF4, LoRA rank 16) and ZeRO-3 offloading.

Key Experimental Results

Main Results

On the full MIMIC-IV benchmark (clinical discharge summaries with ICD-10 codes) and the extremely sparse MIMIC-III-few subset (685 codes appearing fewer than 5 times), PLANT achieves massive improvements over prior SOTA baselines:

Model MIMIC-IV Macro-F1 MIMIC-IV P@8 MIMIC-IV P@15 MIMIC-III-few Macro-F1 MIMIC-III-few Macro-Recall
CoRelation (Luo et al., 2024) 6.3 70.0 55.3 25.0 26.5
PLM-CA (Edin et al., 2024) 22.3 70.5 55.8 26.5 28.0
GKI-ICD (Zhang et al., 2025) 20.6 70.7 55.8 24.0 25.5
GPT-4 Zero-Shot (Yuan et al., 2025) 5.0 68.0 53.5 – –
PLANT (Ours, Mistral-7B) 23.0 (+0.7) 72.1 (+1.4) 56.9 (+1.1) 66.3 (+36.1) 81.0 (+48.1)

Across four foundation LLMs (Mistral-7B, LLaMA3-8B, DeepSeek-V3 336B, Phi-3 3.8B), PLANT yields consistent gains across all metrics (average gains: Macro-AUC +7.2, Macro-F1 +3.0, P@15 +3.1). Notably, compact models with PLANT outperform substantially larger unenhanced LLMs: LLaMA3-8B + PLANT (8B) exceeds DeepSeek-V3 (336B) by +3.7 in Macro-F1 and +3.5 in P@15 on MIMIC-IV-full.

Ablation Studies

Ablation configurations on MIMIC-IV-full using LLaMA3-8B isolate each component's contribution:

Ablation Config Macro-AUC Macro-F1 P@15 Performance Drop & Insight
PLANT (Full setup) 97.4 23.0 56.9 Full two-stage pipeline baseline
Single-Stage BCE (No Stage 1) 91.0 18.0 52.0 Macro-F1 drops by 5.0; confirms critical need for attention pretraining
Single-Stage Focal+HNM (No Stage 1) 92.5 19.5 53.5 Macro-F1 drops by 3.5; loss engineering cannot fix unseeded attention
PLANT w/ Vanilla BCE (Stage 2) 95.0 21.0 54.8 Retaining Stage 1 still beats single-stage Focal models despite simple BCE
PLANT w/o Label Smoothing 97.0 21.8 55.8 Overconfidence causes a mild 1.2 drop in Macro-F1
PLANT w/o Hard Negative Mining 96.8 21.0 54.5 Lack of hard negative focus drops P@15 by 2.4 and Macro-F1 by 2.0
PLANT w/ Term Frequency in Stage 1 96.5 20.5 54.0 Coarse frequency signals degrade tail quality; Macro-F1 drops by 2.5
PLANT w/ MSE Loss in Stage 1 96.0 20.0 53.8 Removing pairwise ranking signals leads to a 3.0 drop in Macro-F1

Key Findings

  • Stage 1 attention pretraining is the primary driver of performance: Completely omitting Stage 1 and training from scratch causes severe drops of 3.5 to 5.0 in Macro-F1; even when Stage 2 uses basic BCE, the pretrained attention maintains a +1.5 Macro-F1 advantage over single-stage Focal training.
  • Attention initialization quality directly predicts downstream tail accuracy: Stage 1 validation nDCG@k strongly correlates with final downstream Macro-F1 (Pearson \(r=0.80, p<0.001\)). Training Stage 1 for only 1 epoch reduces nDCG@k by 0.24 and costs 2.4 points in downstream Macro-F1.
  • Superior data efficiency and zero-shot transfer: PLANT matches the performance of baseline models with an order of magnitude less annotated data (e.g., reaching P@5=0.50 with 1,090 samples vs. 10,337 samples for the baseline). When trained exclusively on common labels (>1% frequency), PLANT retains 7.3%–8.1% Rare-F1 on unseen tail classes, whereas random attention initialization collapses entirely to 0.5%–1.1%.

Highlights & Insights

  • Paradigm shift from passive learning to active planting: Demonstrates that rare-label failure in extreme multi-label classification is fundamentally an attention initialization bottleneck rather than an encoder capacity limit; converting attention into a pretrainable component resolves the long-tail discovery barrier.
  • Model-agnostic applicability with extreme sample efficiency: Integrates seamlessly with diverse quantized LLM backbones as well as compact encoders like DistilBERT, cutting annotation requirements by roughly 10× while outperforming massive closed-source or specialized few-shot baselines.

Limitations & Future Work

  • Dependence on corpus token coverage: MIG scores require offline contingency statistics from the training set, meaning purely out-of-vocabulary technical terminology or novel cross-domain distributions cannot be dynamically updated without recomputing statistics.
  • Stage 1 pretraining overhead: The two-stage protocol introduces an auxiliary pretraining phase and pairwise ranking computations across top-\(k\) tokens, which increases training wall-clock time and memory usage during Stage 1.
  • Future directions: Exploring how attention planting can be generalized to cross-modal tasks (such as vision-language retrieval or video grounding), where cross-attention maps could be seeded with external multimodal alignment priors.
  • vs XRR / CoRelation: Prior state-of-the-art methods construct graph neural networks, co-occurrence matrices, or hierarchical search trees to model code relationships, yet their label attention layers remain randomly initialized and struggle to discover tail structures; PLANT shows that proper attention initialization with a standard architecture outperforms intricate graph reasoning pipelines.
  • vs GPT-4 / Zero-Shot Prompting: Prompting LLMs with long flat label lists suffers from severe "attention dilution" and lost-in-the-middle issues on tail classes; PLANT demonstrates that parameter-efficient fine-tuning with pretrained attention enables smaller 7B–8B models to surpass GPT-4 by over 4× on Macro-F1.

Evaluation

  • Novelty: ⭐⭐⭐⭐⭐ Pioneering perspective treating multi-label attention as an independently pretrainable Learning-to-Rank module driven by information theory.
  • Experimental Rigor: ⭐⭐⭐⭐⭐ Comprehensive coverage across 4 foundation LLMs, 3 diverse cross-domain datasets, exhaustive causal ablations, and zero-shot rare label evaluations.
  • Writing Quality: ⭐⭐⭐⭐⭐ Highly coherent narrative, elegant diagnostic experiments, and transparent presentation of results.
  • Practical Value: ⭐⭐⭐⭐⭐ Plug-and-play compatibility delivering massive improvements on rare tail classes with high utility for clinical and legal AI pipelines.