用于极端多标签文本分类的检索增强编码器
计算与语言
2025-02-18 v1
摘要
极端多标签分类 (XMC) 旨在从海量标签集合中为给定文本输入找到相关标签。为解决这一庞大的标签空间，当前的先进方法可分为两类。一种是一对所有 (OVA) 方法，使用可学习的标签嵌入表示每个标签，擅长记忆（即捕获详细的训练信号以实现准确的头部标签预测）。相比之下，双编码器 (DE) 模型将输入和标签文本映射到共享嵌入空间，以实现更好的泛化（即使用有限的训练数据预测尾部标签的能力），但可能在记忆方面不足。为实现泛化与记忆，现有的 XMC 方法常将 DE 和 OVA 模型组合，这涉及复杂的训练管道。灵感来自检索增强语言模型的成功，本文提出检索增强编码器用于 XMC (RAEXMC)，一种新型框架，为 DE 模型提供检索增强功能，以无需额外可训练参数实现高效记忆。在训练期间，RAEXMC 基于由输入实例和标签组成的知识记忆进行对比损失优化。在推断期间，给定测试输入后，RAEXMC 从知识记忆中检索 top-K 个 key，并将相应的 value 聚合作为预测得分。我们在四个公开的 LF-XMC 基准测试中展示了 RAEXMC 的有效性与效率。RAEXMC 不仅推动了先进的 DE 方法 DEXML，也在相同的 8 块 A100 GPU 训练环境下实现了在最大的 LF-AmazonTitles-1.3M 数据集上超过 10 倍的加速。
引用
@article{arxiv.2502.10615,
title = {Retrieval-augmented Encoders for Extreme Multi-label Text Classification},
author = {Yau-Shian Wang and Wei-Cheng Chang and Jyun-Yu Jiang and Jiong Zhang and Hsiang-Fu Yu and S. V. N. Vishwanathan},
journal= {arXiv preprint arXiv:2502.10615},
year = {2025}
}