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

TL;DR

PLANT is a novel attention initialization method for extreme multi-label text classification that leverages a pretrained Learning-to-Rank model, significantly improving performance especially in few-shot and rare label scenarios.

Contribution

Introduces PLANT, a plug-and-play attention initialization technique using Learning-to-Rank guidance, compatible with various large language models, enhancing classification accuracy.

Findings

Outperforms state-of-the-art methods across multiple tasks.

Significant improvements in few-shot and rare label settings.

Attention initialization is a key factor in performance gains.

Abstract

State-of-the-art Extreme Multi-Label Text Classification models rely on multi-label attention to focus on key tokens in input text, but learning good attention weights is challenging. We introduce PLANT - Pretrained and Leveraged Attention - a plug-and-play strategy for initializing attention. PLANT works by planting label-specific attention using a pretrained Learning-to-Rank model guided by mutual information gain. This architecture-agnostic approach integrates seamlessly with large language model backbones such as Mistral-7B, LLaMA3-8B, DeepSeek-V3, and Phi-3. PLANT outperforms state-of-the-art methods across tasks including ICD coding, legal topic classification, and content recommendation. Gains are especially pronounced in few-shot settings, with substantial improvements on rare labels. Ablation studies confirm that attention initialization is a key driver of these gains. For code and trained models, see https://github.com/debjyotiSRoy/xcube/tree/plant