Abstract
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.
Abstract (translated)
极端多标签分类（XMC）的目标是从一个极其庞大的标签集合中找出与给定文本输入相关的标签。为了应对这种大规模的标签空间，当前最先进的方法主要分为两大类：一种是使用可学习的标签嵌入的“一对所有”（OVA）方法，这种方法在记忆化方面表现出色，即能够捕捉到详细的训练信号以准确预测头部标签；另一种是双编码器（DE）模型，该模型将输入文本和标签映射到一个共享的嵌入空间中，从而更好地泛化（即，在有限的数据下也能有效预测尾部标签），但可能在记忆化方面略有不足。现有的XMC方法通常通过结合使用DE和OVA模型来实现泛化和记忆化的效果，这涉及复杂的训练流程。 受检索增强语言模型成功的启发，我们提出了用于XMC的检索增强编码器（RAEXMC）框架。该框架为双编码器模型配备了检索增强能力，以高效地进行记忆化，而不增加额外可学习参数的数量。在训练阶段，通过对比损失优化RAEXMC，并利用一个包含输入实例和标签的知识内存来实现这一目的。在推理阶段，给定一个测试输入后，RAEXMC从知识内存中检索出前K个键，然后将对应的值聚合为预测得分。 我们在四个公开的LF-XMC基准数据集上展示了RAEXMC的有效性和效率。不仅比现有的最先进的双编码器方法DEXML有所改进，而且在最大的LF-AmazonTitles-1.3M数据集中，在相同条件下使用8块A100 GPU进行训练时，速度提高了超过10倍。
URL
https://arxiv.org/abs/2502.10615
