Don’t Pay Attention, PLANT It: Pretraining Attention via Learning-to-Rank¶
会议: NeurIPS 2026
arXiv: 2410.23066
领域: NLP 理解 / NLP Understanding
关键词: 极端多标签分类, 注意力预训练, 排序学习, 互信息增益, 长尾标签
一句话总结¶
针对极端多标签文本分类中随机初始化标签注意力导致长尾稀有标签难以对齐的瓶颈，本文提出基于互信息增益与排序学习的两阶段预训练注意力方法 PLANT，使小体量大模型在稀有标签分类上大幅超越更大规模基线与现有 SOTA。
研究背景与动机¶
极端多标签文本分类（Extreme Multi-Label Text Classification, XMTC）要求模型从数万甚至数十万个高度倾斜的候选标签集合中为长文本匹配最相关子集，广泛存在于临床医学诊断（ICD 编码）、法律判例归类以及电商内容推荐等领域。现有绝大多数高性能 XMTC 架构均依赖多标签注意力机制（Multi-Label Attention），即为每个标签分配专属的查询向量，在输入文本的长序列中充当“定向探照灯”，抽取该标签最敏感的局部特征。
然而，现有方法在初始化多标签注意力模块时普遍采取随机权重初始化，迫使模型在端到端训练阶段从零开始探索各标签在海量词元中的注意力分布。这种依赖大量训练样本反向传播更新的机制在头部高频标签上尚能收敛，但面对占据标签空间极大比例的尾部长尾稀有标签时，极少甚至不足 5-10 次的样本出现频次根本无法驱动注意力向量探索出合理的语义结构。实证诊断显示，随机初始化的模型对语义高度相近的稀有标签所产生的注意力模式极为发散（类内余弦相似度均值仅约 0.75），甚至在增加轮次后迅速陷入过拟合，直接导致稀有标签预测指标全面崩溃。
本文的核心洞察在于：无需在端到端任务微调时强求稀有标签凭空学出文本注意力，因为标签与词元之间的信息论相关性完全可以先验计算。核心 idea：将标签注意力模块重构为一个可通过排序学习预训练的独立组件，利用互信息增益构建词元级弱监督先验并以 LambdaLoss 预训练注意力层，使模型在正式下游训练前即已“播种”精准的标签聚焦结构。
方法详解¶
整体框架¶
PLANT 包含两个解耦的训练阶段：阶段 1 为“播种注意力”（PLANTing Attention），冻结或热身文本编码器，利用训练语料中标签与词元共现计算得到的互信息增益（Mutual Information Gain, MIG）作为监督信号，将多标签注意力模块作为排序学习（Learning-to-Rank, L2R）组件进行预训练；阶段 2 为“杠杆化复用”（Leveraging Attention），加载预训练好的标签查询与注意力权重，采用分阶段渐进式解冻策略和难负样本挖掘的焦点损失进行端到端微调。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["长文本输入 x 与候选标签全集 L"] --> B["文本编码器与标签查询嵌入 E"]
B --> C["阶段 1：互信息引导排序预训练<br/>MIG 计算 top-k 词元相关度，LambdaLoss 对齐注意力"]
C --> D["阶段 2：自适应特征增强与渐进解冻<br/>可微特征缩放 + 难负样本挖掘焦点损失微调"]
D --> E["极端多标签概率输出"]关键设计¶
1. 互信息引导排序预训练：跳过随机探索直接固化注意力先验 针对稀有标签因样本稀缺无法通过梯度更新学习词元权重的痛点，PLANT 首先通过离线统计全量语料的列联表与条件独立性，计算每个标签 \(l\) 与每个词元 \(t_j\) 之间的互信息增益（MIG）\(r_{l,j}\)： $\(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)\)$ 其中 \(x\) 与 \(y\) 分别代表标签与词元在文档中是否出现。相比朴素词频或共现计数，互信息能有效滤除全语料高频停用词的虚假相关性。对每个标签选取 MIG 最高的 top-\(k\) 个词元构建候选集合 \(\mathcal{T}_l\)。随后，多头注意力模块接收标签嵌入 \(\boldsymbol{Q}=\boldsymbol{E}\) 以及文本隐状态 \(\boldsymbol{K}=\boldsymbol{V}=\boldsymbol{H}_i\)，输出注意力分布矩阵 \(\boldsymbol{S}\in\mathbb{R}^{|\mathcal{L}|\times n}\)。阶段 1 使用基于 nDCG@k 差值的 LambdaLoss 优化注意力投影参数 \(\boldsymbol{W}_{\mathsf{attn}}\)： $\(\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)\)$ 该排序损失显式迫使注意力打分在重要词元对之间遵循信息论相关性，将稀有标签组的类内注意力一致性直接从 0.75 拉升至 0.985。
2. 自适应特征增强与自适应池化：差异化放大微弱尾部表征 在注意力层输出经过关注的文本表征 \(\boldsymbol{A}\in\mathbb{R}^{|\mathcal{L}|\times h}\) 后，为了避免极高维标签空间中稀有标签特征被占据主导地位的高频特征淹没，模型引入了两个可学习的标签特定缩放与偏置矩阵 \(\boldsymbol{B}_a, \boldsymbol{B}_m\in\mathbb{R}^{|\mathcal{L}|\times h}\)，对注意力表征进行逐元素可微调节： $\(\boldsymbol{A}_{\text{boost}}=\boldsymbol{B}_a+\boldsymbol{A}\odot\boldsymbol{B}_m\)$ 随后通过自适应平均池化层（Adaptive Average Pooling）将高维隐空间压缩至 \(p=128\) 维向量 \(\boldsymbol{P}_i\in\mathbb{R}^{|\mathcal{L}|\times p}\)，最终经由共享分类头 \(\boldsymbol{W}_c\in\mathbb{R}^{p\times 1}\) 映射为标签 logits \(\hat{\boldsymbol{y}}_i=\boldsymbol{P}_i\boldsymbol{W}_c\)。该结构赋予网络显式强调或抑制特定标签通道的能力，极大提升了模型区分混淆近义标签的分辨率。
3. 渐进式解冻与难负样本挖掘微调：保护预训练结构并对抗极端不平衡 进入阶段 2 之后，若直接对全部参数施加大规模梯度反向传播，极易冲刷并破坏阶段 1 播种的注意力结构。PLANT 实施了多层级的渐进解冻机制：在前 5 个 epoch 保持注意力模块和标签嵌入完全冻结，仅微调下游投影与分类头；其后按注意力层（lr = \(1\times 10^{-5}\)）、中间层（lr = \(5\times 10^{-6}\)）至全局参数（lr = \(2\times 10^{-6}\)）分阶段逐渐解冻。同时，为了对抗极端的正负样本数量失衡，损失函数在阶段 2 采用带标签平滑（\(\epsilon=0.1\)）与难负样本挖掘（Hard Negative Mining, HNM）的焦点损失（Focal Loss）：对每条样本保留全部正标签，并仅从预测概率最高的负标签中筛选前 \(m=1000\) 个难负例参与损失计算。
损失函数 / 训练策略¶
阶段 1 优化目标为式 (2) 的排序损失，采用 AdamW 优化器，峰值学习率 \(\eta=5\times 10^{-4}\)，配合余弦退火和 10% 热身步数，训练 10 轮。阶段 2 采用带难负样本采样的焦点损失： $\(\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]\)$ 其中聚焦系数 \(\gamma=2\)，平滑标签 \(\tilde{y}_{il}=(1-\epsilon)y_{il}+\epsilon(1-y_{il})\)，候选集合 \(\mathcal{S}_i\) 包含该样本所有正标签与 top-\(1000\) 预测分最高的负标签。整个流程在 8×A100-80GB 环境下结合 QLoRA（4-bit 量化，LoRA rank=16）与 ZeRO-3 开展训练。
实验关键数据¶
主实验¶
在 MIMIC-IV-full（真实临床全集，数千长尾 ICD 编码）以及极度稀有的子集 MIMIC-III-few（出现频次严格少于 5 次的 685 个标签）上评测，PLANT 展现出巨大优势：
| 模型 | 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) |
此外，在 4 种主流开源 LLM 基座（Mistral-7B, LLaMA3-8B, DeepSeek-V3 336B, Phi-3 3.8B）上，PLANT 带来全指标一致提升，平均增益为 Macro-AUC +7.2、Macro-F1 +3.0、P@15 +3.1。极小规模的 LLaMA3-8B + PLANT 即在 Macro-F1（23.5 vs 19.8）和 P@15（57.0 vs 53.5）上反超原版 DeepSeek-V3（336B）。
消融实验¶
在 MIMIC-IV-full 上使用 LLaMA3-8B 评估各个关键组件的影响：
| 消融配置 | Macro-AUC | Macro-F1 | P@15 | 相对下降说明 |
|---|---|---|---|---|
| PLANT (Full setup) | 97.4 | 23.0 | 56.9 | 完整两阶段配置基线 |
| 单阶段 BCE (无阶段 1，仅 BCE) | 91.0 | 18.0 | 52.0 | Macro-F1 暴跌 5.0，验证阶段 1 注意力初始化的核心价值 |
| 单阶段 Focal+HNM (无阶段 1) | 92.5 | 19.5 | 53.5 | 即使有难负样本挖掘，缺少阶段 1 仍导致 Macro-F1 下降 3.5 |
| PLANT 搭配 Vanilla BCE | 95.0 | 21.0 | 54.8 | 保留阶段 1 但阶段 2 仅用简单 BCE，仍大幅超越单阶段 Focal 模型 |
| PLANT 去除 Label Smoothing | 97.0 | 21.8 | 55.8 | 模型轻度过拟合过自信，Macro-F1 下降 1.2 |
| PLANT 去除 Hard Negative Mining | 96.8 | 21.0 | 54.5 | 缺乏难负例监督，P@15 下降 2.4，Macro-F1 下降 2.0 |
| PLANT 阶段 1 使用词频替代 MIG | 96.5 | 20.5 | 54.0 | 粗糙词频受常见停用词干扰，Macro-F1 下降 2.5 |
| PLANT 阶段 1 使用 MSE 替代排序损失 | 96.0 | 20.0 | 53.8 | 丢失相对排位信息，Macro-F1 下降 3.0 |
关键发现¶
- 阶段 1 注意力预训练是性能基石：完全剥离阶段 1 直接从零端到端训练会导致 Macro-F1 发生 3.5 至 5.0 点的断崖式下跌；即使阶段 2 仅采用最简单的 BCE 损失，拥有阶段 1 权重的模型依然超越拥有精细 Focal Loss 但随机初始化的模型。
- 阶段 1 排序质量与下游稀有标签性能强相关：阶段 1 验证集 nDCG@k 与阶段 2 最终 Macro-F1 呈高度正相关（Pearson \(r=0.80, p<0.001\)）。当仅预训练 1 个 epoch 时，nDCG@k 下降 0.24，导致下游 Macro-F1 缩水 2.4 点。
- 跨领域通用性与少样本/零样本泛化突出：在法律长文本 EURLEX-4K 和词条推荐 WIKI10-31K 上，搭载 DistilBERT 骨干的 PLANT 在 P@3 指标上分别实现 +2.47 和 +2.23 的显著飞跃；仅在常见标签（频率 >1%）上训练时，PLANT 对未见过的稀有标签依然能保持 7.3%–8.1% 的零样本 Rare-F1，而随机初始化基线几乎完全崩溃（0.5%–1.1%）。
亮点与洞察¶
- 范式级转变：从“被动学习注意力”到“主动播种注意力”：首次指出极端多标签分类中稀有类别崩溃的病灶在于注意力初始化而非分类器容量，通过信息论（MIG）与搜索排序（Learning-to-Rank）为注意力注入冷启动先验，从根源上化解了长尾稀疏性。
- 跨骨干适配性与极高样本效率：作为即插即用型模块，PLANT 无缝兼容 QLoRA 适配的各种通用大模型乃至小体量编码器；在 MIMIC-IV-full 上仅需约 1,000–2,700 标注样本即可达到基线需超万条样本方能达到的性能，节约了一个数量级的标注成本。
局限与展望¶
- 互信息统计依赖语料覆盖：MIG 依赖离线统计训练集的词元-标签列联表，若出现完全不在训练语料中出现的新型专业词汇或全新跨域文本，其静态先验无法动态扩展。
- 预训练计算额外开销：两阶段方案引入了额外的阶段 1 预训练轮次，且每条样本需要按标签维护 top-\(k\) 词元对进行排位损失计算，在大规模标签全集下的批处理显存占用较高。
- 未来方向：探索将注意力播种思想扩展到多模态图文跨模态检索或视频长序列定位任务中，利用外部知识库或跨模态对比表征直接合成注意力先验。
相关工作与启发¶
- vs XRR / CoRelation: 现有 SOTA 模型试图通过构建标签共现图、设计层次化树状搜索或双编码器结构来捕捉标签间关联，但其内部多头标签注意力仍然是随机初始化、依赖下游监督强行学习的；PLANT 证明仅通过优化注意力初始分布，就能以极简架构压倒复杂的图推理模型。
- vs GPT-4 / 零样本 LLM: 纯提示词零样本方案在超大标签空间下面临严重的“注意力稀释”（Attention Dilution）与长上下文丢失（Lost-in-the-Middle）；PLANT 证明轻量微调结合先验注意力能够将知识沉淀进参数，使 7B–8B 级别开源模型在宏平均 F1 上以 4 倍以上的优势击败闭源大模型。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次提出将注意力层作为 Learning-to-Rank 模块进行基于互信息的独立预训练，视角极为独到。
- 实验充分度: ⭐⭐⭐⭐⭐ 覆盖 4 个经典大模型、3 个跨领域数据集、完整的因果消融与零样本长尾泛化实验，论证缜密。
- 写作质量: ⭐⭐⭐⭐⭐ 逻辑线清晰，诊断性实验设计严谨，图表表达直观且信息密度高。
- 价值: ⭐⭐⭐⭐⭐ 即插即用且显著改善长尾标签性能，对极端多标签分类与长文本理解具备广泛的工业落地价值。