arXiv:2409.20156cs.LGcs.IR2024-09

提出高效训练超大规模分类模型的新方法,显著提升精度与速度。

ASTRA: Accurate and Scalable ANNS-based Training of Extreme Classifiers

  • 基于近似最近邻索引动态筛选难负样本,对齐损失函数优化目标。
  • 采用混合采样策略,在保持高精度的同时将训练时间减少4至15倍。
  • 适用于百万级标签场景,特别适合搜索与推荐系统中的极端分类任务。

极端分类(Extreme Classification, XC)旨在从海量标签集合($L$可达数千万)中为数据点(查询)标注相关标签,广泛应用于搜索与推荐系统。当前主流方法使用深度编码器(如DistilBERT)嵌入查询与标签,并在查询嵌入上构建线性分类器,结合近似最近邻搜索(ANNS)实现毫秒级推理。然而,现有高精度方法(如DEXML、Renée、DEXA)存在每轮训练时间随标签数$O(L)$增长或依赖昂贵负样本采样策略的问题,难以在有限GPU资源下扩展。本文提出ASTRA算法,基于两个关键观察:(a) 在分类器向量上构建ANNS索引并用其检索难负样本,可使负样本采样策略与优化目标对齐;(b) 持续更新索引成本过高,而仅使用过时负样本会导致精度下降。为此,提出融合重要性采样与均匀采样的混合负样本策略。在包含120M标签的标准及私有数据集上的大量实验表明,ASTRA在达到最先进精度的同时,训练时间相比次优方法降低4倍至15倍。

原文摘要 · Abstract (English)

`Extreme Classification'' (or XC) is the task of annotating data points (queries) with relevant labels (documents), from an extremely large set of $L$ possible labels, arising in search and recommendations. The most successful deep learning paradigm that has emerged over the last decade or so for XC is to embed the queries (and labels) using a deep encoder (e.g. DistilBERT), and use linear classifiers on top of the query embeddings. This architecture is of appeal because it enables millisecond-time inference using approximate nearest neighbor search (ANNS). The key question is how do we design training algorithms that are accurate as well as scale to $O(100M)$ labels on a limited number of GPUs. State-of-the-art XC techniques that demonstrate high accuracies (e.g., DEXML, Renée, DEXA) on standard datasets have per-epoch training time that scales as $O(L)$ or employ expensive negative sampling strategies, which are prohibitive in XC scenarios. In this work, we develop an accurate and scalable XC algorithm ASTRA with two key observations: (a) building ANNS index on the classifier vectors and retrieving hard negatives using the classifiers aligns the negative sampling strategy to the loss function optimized; (b) keeping the ANNS indices current as the classifiers change through the epochs is prohibitively expensive while using stale negatives (refreshed periodically) results in poor accuracy; to remedy this, we propose a negative sampling strategy that uses a mixture of importance sampling and uniform sampling. By extensive evaluation on standard XC as well as proprietary datasets with 120M labels, we demonstrate that ASTRA achieves SOTA precision, while reducing training time by 4x-15x relative to the second best.

极端分类ANNS训练加速推荐系统

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。