ASTRA：基于ANNS的极端分类器精准可扩展训练
机器学习
2024-10-01 v1 信息检索
摘要
“极端分类”（或XC）是将数据点（查询）与来自极大标签集中相关标签（文档）进行注释的任务，这些标签在搜索和推荐中常见。过去十年来最成功的深度学习方法是使用深度编码器（例如DistilBERT）对查询（和标签）进行嵌入，并在查询嵌入上使用线性分类器。这种架构具有吸引力，因为它可以使用近似最近邻搜索（ANNS）实现毫秒级推断。关键问题是如何设计既准确又能扩展到O(100M)标签的训练算法。展示出高准确性（例如DEXML、Renée、DEXA）的当前最先进XC技术在标准数据集上具有O()或采用昂贵负采样策略的每epoch训练时间，这在XC场景中是不可接受的。在本工作中，我们开发了一个准确且可扩展的XC算法ASTRA，基于两个关键观察：（a）在分类器向量上构建ANNS索引并使用分类器检索硬负样本将负采样策略与优化的损失函数对齐；（b）随着分类器通过各epoch变化而保持ANNS索引是代价高昂的，而使用陈旧负样本（定期刷新）会导致准确性差；为缓解这一问题，我们提出一种结合重要性采样和均匀采样的负采样策略。通过在标准XC以及具有120M标签的专有数据集上进行广泛评估，我们展示ASTRA在精度上实现了SOTA，同时将相对于第二名的训练时间缩短了4倍至15倍。
引用
@article{arxiv.2409.20156,
title = {ASTRA: Accurate and Scalable ANNS-based Training of Extreme Classifiers},
author = {Sonu Mehta and Jayashree Mohan and Nagarajan Natarajan and Ramachandran Ramjee and Manik Varma},
journal= {arXiv preprint arXiv:2409.20156},
year = {2024}
}