Research map / Label-efficient and robust learning / Recognition and long-tail
Don't Pay Attention, PLANT It: Pretraining Attention via Learning-to-Rank
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes PLANT, a two-stage method that pretrains label-specific attention as a learning-to-rank model guided by mutual information gain, then fine-tunes the LLM-based XMTC model end-to-end, improving rare-label performance on ICD coding, legal, and tag prediction tasks.
Problem
Extreme multi-label text classification relies on label-specific attention, but randomly initialized attention learns poorly for rare labels in skewed, high-dimensional label spaces.
Approach
Initialize the label attention layer by pretraining it as a pairwise learning-to-rank model whose relevance targets are mutual information gain between tokens and labels (top-k tokens, nDCG-weighted pairwise sigmoid loss), then fine-tune the full model end-to-end with focal loss and hard negative mining.
Key results
- On MIMIC-IV, Mistral+PLANT reaches macro F1 23.0 vs 20.0 without PLANT and P@8 72.1 vs 70.7 for best prior baseline GKI-ICD.
- On MIMIC-III few (labels with <5 training samples), PLANT achieves macro F1 66.3 vs 30.2 for the best baseline GP (+36.1).
- On EURLEX, PLANT improves P@3 to 81.35 over the best listed baseline XRR (78.88), and on WikiTen P@3 gains +2.23.
Surprising findings
- A Llama-3.1-8B+PLANT model outperforms the 336B DeepSeek model alone on several MIMIC-III metrics.
- Training with random attention initialization on common labels only collapses rare-label F1 to 0.5%, while PLANT retains 8.0%.
Limitations
- Reported results contain internal inconsistencies (e.g., hardware stated as both 8xA100 and a single RTX 8000 GPU; DeepSeek-R-336B and a mistral model URL that appear nonstandard), raising reproducibility concerns.
- Gains on WikiTen P@5 are negligible or slightly negative.
- Evaluated mainly on ICD coding; generalization beyond text classification is only argued as future work.
Experimental setup
Datasets: MIMIC-III full MIMIC-IV MIMIC-III rare50 MIMIC-III few EURLEX-57K Wiki10-31K
Benchmarks: MIMIC-III full MIMIC-IV MIMIC-III few MIMIC-III rare50 EURLEX WikiTen
Metrics: macro AUC micro AUC macro F1 micro F1 precision@k recall@k nDCG@k Wilcoxon signed-rank test bootstrap 95% CI
Baselines: CoRelation PLM-CA GKI-ICD GPT-4 Zero-Shot GP XRR MatchXML InceptionXML CG DE X-Transformer w/ RDE MSMN + Contrastive
Builds on: LoRA focal loss LAAT CoRelation
Closest papers by meaning
- SAS: Simple Attention Sparsification via End-to-End Optimization of Context Ranking (NeurIPS 2026 · similarity 0.85)
- On Understanding Attention-Based In-Context Learning for Categorical Data (ICML 2025 · similarity 0.84)
- Punctuation-aware Hybrid Trainable Sparse Attention for Large Language Models (NeurIPS 2026 · similarity 0.84)
- Gen-Z: Generative Zero-Shot Text Classification with Contextualized Label Descriptions (ICLR 2024 · similarity 0.83)
- Attention with Trained Embeddings Provably Selects Important Tokens (NeurIPS 2025 · similarity 0.83)
- Scalable In-context Ranking with Generative Models (NeurIPS 2025 · similarity 0.83)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.
