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

TL;DR

ASTRA introduces a scalable and accurate training algorithm for extreme classifiers that leverages approximate nearest neighbor search with a novel negative sampling strategy, achieving state-of-the-art precision and significantly faster training on large datasets.

Contribution

The paper proposes ASTRA, a new method that improves training efficiency and accuracy for extreme classifiers by aligning negative sampling with the loss function and using a mixed sampling approach.

Findings

ASTRA achieves state-of-the-art precision on large-scale datasets.

Training time is reduced by 4x-15x compared to previous methods.

Effective negative sampling strategy improves classifier accuracy.

Abstract

`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\'ee, 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.