用检索增强提升双编码器模型的标签记忆能力,实现高效精准的极端多标签分类。
Retrieval-augmented Encoders for Extreme Multi-label Text Classification
- 引入检索增强机制,让双编码器在不增加参数的前提下获得记忆能力。
- 在四个公开数据集上超越现有最优双编码器方法,最大数据集提速超10倍。
- 适合需要快速部署且兼顾头部与尾部标签预测精度的应用场景。
极端多标签分类(XMC)旨在从极大规模标签集中为给定文本找出相关标签。当前主流方法分为两类:一元对所有(OVA)通过可学习标签嵌入实现强记忆能力,擅长头部标签预测;双编码器(DE)将输入与标签映射到共享空间,具备更好泛化性,但记忆能力弱。现有方法常混合两者,训练复杂。受检索增强语言模型启发,本文提出检索增强编码器框架RAEXMC,赋予双编码器检索增强能力,实现高效记忆而无需额外参数。训练时,基于包含输入实例和标签的知识记忆库,使用对比损失优化;推理时,对测试输入检索知识记忆库中前K个最相似项,并聚合其对应值作为预测得分。在四个公开的LF-XMC基准上验证了该方法的有效性与高效性。RAEXMC不仅超越现有最优双编码器方法DEXML,且在最大规模的LF-AmazonTitles-1.3M数据集上,在相同8张A100 GPU环境下实现超过10倍的加速。
原文摘要 · Abstract (English)
Extreme multi-label classification (XMC) seeks to find relevant labels from an extremely large label collection for a given text input. To tackle such a vast label space, current state-of-the-art methods fall into two categories. The one-versus-all (OVA) method uses learnable label embeddings for each label, excelling at memorization (i.e., capturing detailed training signals for accurate head label prediction). In contrast, the dual-encoder (DE) model maps input and label text into a shared embedding space for better generalization (i.e., the capability of predicting tail labels with limited training data), but may fall short at memorization. To achieve generalization and memorization, existing XMC methods often combine DE and OVA models, which involves complex training pipelines. Inspired by the success of retrieval-augmented language models, we propose the Retrieval-augmented Encoders for XMC (RAEXMC), a novel framework that equips a DE model with retrieval-augmented capability for efficient memorization without additional trainable parameter. During training, RAEXMC is optimized by the contrastive loss over a knowledge memory that consists of both input instances and labels. During inference, given a test input, RAEXMC retrieves the top-$K$ keys from the knowledge memory, and aggregates the corresponding values as the prediction scores. We showcase the effectiveness and efficiency of RAEXMC on four public LF-XMC benchmarks. RAEXMC not only advances the state-of-the-art (SOTA) DE method DEXML, but also achieves more than 10x speedup on the largest LF-AmazonTitles-1.3M dataset under the same 8 A100 GPUs training environments.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。