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

arXiv: https://arxiv.org/abs/2409.20156 · HTML: https://mlsanity.com/paper/2409.20156/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 24 of its 43 references that are in the map.
- First posted on arXiv 30 Sep 2024.

## Builds on (most cited first, 24 of 24)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [Sigmoid Loss for Language Image Pre-Training](https://mlsanity.com/paper/2303.15343/) (2023)
- [Supervised Contrastive Learning](https://mlsanity.com/paper/2004.11362/) (2020)
- [Bag of Tricks for Efficient Text Classification](https://mlsanity.com/paper/1607.01759/) (2016)
- [Approximate Nearest Neighbor Negative Contrastive Learning for Dense Text Retrieval](https://mlsanity.com/paper/2007.00808/) (2020)
- [Latent Retrieval for Weakly Supervised Open Domain Question Answering](https://mlsanity.com/paper/1906.00300/) (2019)
- [RocketQA: An Optimized Training Approach to Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2010.08191/) (2020)
- [Efficiently Teaching an Effective Dense Retriever with Balanced Topic Aware Sampling](https://mlsanity.com/paper/2104.06967/) (2021)
- [Smart Mining for Deep Metric Learning](https://mlsanity.com/paper/1704.01285/) (2017)
- [Variance Reduction in SGD by Distributed Importance Sampling](https://mlsanity.com/paper/1511.06481/) (2015)
- [DiSMEC - Distributed Sparse Machines for Extreme Multi-label Classification](https://mlsanity.com/paper/1609.02521/) (2016)
- [Sub-sampled Cubic Regularization for Non-convex Optimization](https://mlsanity.com/paper/1705.05933/) (2017)
- [Bonsai -- Diverse and Shallow Trees for Extreme Multi-label Classification](https://mlsanity.com/paper/1904.08249/) (2019)
- [LightXML: Transformer with Dynamic Negative Sampling for High-Performance Extreme Multi-label Text Classification](https://mlsanity.com/paper/2101.03305/) (2021)
- [ECLARE: Extreme Classification with Label Graph Correlations](https://mlsanity.com/paper/2108.00261/) (2021)
- [DeepXML: A Deep Extreme Multi-Label Learning Framework Applied to Short Text Documents](https://mlsanity.com/paper/2111.06685/) (2021)
- [DECAF: Deep Extreme Classification with Label Features](https://mlsanity.com/paper/2108.00368/) (2021)
- [NGAME: Negative Mining-aware Mini-batching for Extreme Classification](https://mlsanity.com/paper/2207.04452/) (2022)
- [Stochastic Negative Mining for Learning with Large Output Spaces](https://mlsanity.com/paper/1810.07076/) (2018)
- [ELIAS: End-to-End Learning to Index and Search in Large Output Spaces](https://mlsanity.com/paper/2210.08410/) (2022)

## Abstract (arXiv)

`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.