用原型对比学习提升极端多标签分类效率与精度。
Prototypical Extreme Multi-label Classification with a Dynamic Margin Loss
- 将标签聚类为原型,通过浅层Transformer聚合文本和向量信息。
- 在多个数据集上超越现有方法,且推理速度更快。
- 适合需要高效高精度标签预测的工业场景。
极端多标签分类(XMC)需在极大规模标签空间中为查询预测相关标签。现有方法多采用深度编码器将文本映射到嵌入空间以检索最接近标签,但训练成本高,存在性能与效率的权衡。本文提出PRIME方法,通过新型原型对比学习,实现高效与高性能的统一。将XMC建模为从数据到原型的预测任务,利用标签原型聚合相关查询信息。具体地,设计轻量级的标签原型网络(Label Prototype Network),通过浅层Transformer融合文本嵌入、标签中心点及可学习自由向量,丰富标签表示。联合训练深层编码器与原型网络,采用自适应三元组损失,更好适应极端标签空间的细粒度与模糊性。PRIME在多个不同规模与领域的公开数据集上达到最新最优性能,同时保持模型高效。
原文摘要 · Abstract (English)
Extreme Multi-label Classification (XMC) methods predict relevant labels for a given query in an extremely large label space. Recent works in XMC address this problem using deep encoders that project text descriptions to an embedding space suitable for recovering the closest labels. However, learning deep models can be computationally expensive in large output spaces, resulting in a trade-off between high performing brute-force approaches and efficient solutions. In this paper, we propose PRIME, a XMC method that employs a novel prototypical contrastive learning technique to reconcile efficiency and performance surpassing brute-force approaches. We frame XMC as a data-to-prototype prediction task where label prototypes aggregate information from related queries. More precisely, we use a shallow transformer encoder that we coin as Label Prototype Network, which enriches label representations by aggregating text-based embeddings, label centroids and learnable free vectors. We jointly train a deep encoder and the Label Prototype Network using an adaptive triplet loss objective that better adapts to the high granularity and ambiguity of extreme label spaces. PRIME achieves state-of-the-art results in several public benchmarks of different sizes and domains, while keeping the model efficient.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。