ASTRA:基于ANNS的精确且可扩展的极端分类器训练方法
ASTRA: Accurate and Scalable ANNS-based Training of Extreme Classifiers
AI总结:
ASTRA通过基于分类器向量的ANNS索引检索困难负样本,并混合重要性采样与均匀采样解决索引更新代价问题,在1.2亿标签规模下实现SOTA精度并将训练时间减少4至15倍。
摘要:
“极端分类”(Extreme Classification,XC)是指从极其庞大的L个可能标签集合中,为数据点(查询)标注相关标签(文档)的任务,该任务出现在搜索和推荐场景中。过去十年左右涌现出的最成功的深度学习范式是使用深度编码器(如DistilBERT)对查询(和标签)进行嵌入,并在查询嵌入之上使用线性分类器。这种架构具有吸引力,因为它能够利用近似最近邻搜索(ANNS)实现毫秒级推理。关键问题在于:如何设计出既准确又能扩展到O(100M)标签规模、且在有限数量GPU上运行的训练算法。在标准数据集上表现出高准确率的最先进XC技术(如DEXML、Renée、DEXA),其每轮训练时间随O(L)规模增长,或采用昂贵的负采样策略,这在XC场景中是不可行的。在本工作中,我们开发了一种精确且可扩展的XC算法ASTRA,基于两个关键观察:(a) 在分类器向量上构建ANNS索引,并利用分类器检索困难负样本,使负采样策略与被优化的损失函数保持一致;(b) 随着分类器在各轮训练中不断变化,保持ANNS索引实时更新成本极高,而使用过期的负样本(周期性刷新)会导致准确率下降;为解决这一问题,我们提出了一种混合使用重要性采样和均匀采样的负采样策略。通过在标准XC数据集以及具有1.2亿标签的专有数据集上进行广泛评估,我们证明ASTRA达到了SOTA精度,同时相对于次优方法将训练时间减少了4倍至15倍。
英文摘要:
`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.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。