# Towards Memory-Efficient Training for Extremely Large Output Spaces -- Learning with 500k Labels on a Single Commodity GPU

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

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 34 references that are in the map.
- First posted on arXiv 6 Jun 2023.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Efficient Estimation of Word Representations in Vector Space](https://mlsanity.com/paper/1301.3781/) (2013)
- [Rigging the Lottery: Making All Tickets Winners](https://mlsanity.com/paper/1911.11134/) (2019)
- [Scalable Training of Artificial Neural Networks with Adaptive Sparse Connectivity inspired by Network Science](https://mlsanity.com/paper/1707.04780/) (2017)
- [Sparse Networks from Scratch: Faster Training without Losing Performance](https://mlsanity.com/paper/1907.04840/) (2019)
- [Deep Rewiring: Training very sparse deep networks](https://mlsanity.com/paper/1711.05136/) (2017)
- [AttentionXML: Label Tree-based Attention-Aware Deep Model for High-Performance Extreme Multi-Label Text Classification](https://mlsanity.com/paper/1811.01727/) (2018)
- [Inferring Networks of Substitutable and Complementary Products](https://mlsanity.com/paper/1506.08839/) (2015)
- [DiSMEC - Distributed Sparse Machines for Extreme Multi-label Classification](https://mlsanity.com/paper/1609.02521/) (2016)
- [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)
- [A no-regret generalization of hierarchical softmax to extreme multi-label classification](https://mlsanity.com/paper/1810.11671/) (2018)
- [Extreme Classification in Log Memory using Count-Min Sketch: A Case Study of Amazon Search with 50M Products](https://mlsanity.com/paper/1910.13830/) (2019)
- [Stochastic Negative Mining for Learning with Large Output Spaces](https://mlsanity.com/paper/1810.07076/) (2018)
- [On Missing Labels, Long-tails and Propensities in Extreme Multi-label Classification](https://mlsanity.com/paper/2207.13186/) (2022)
- [CascadeXML: Rethinking Transformers for End-to-end Multi-resolution Training in Extreme Multi-label Classification](https://mlsanity.com/paper/2211.00640/) (2022)
- [Conditional Probability Tree Estimation Analysis and Algorithms](https://mlsanity.com/paper/1408.2031/) (2014)
- [Disentangling Sampling and Labeling Bias for Learning in Large-Output Spaces](https://mlsanity.com/paper/2105.05736/) (2021)

## Built on by (most cited first, 6 of 6)

- [Sparse maximal update parameterization: A holistic approach to sparse training dynamics](https://mlsanity.com/paper/2405.15743/) (2024)
- [Navigating Extremes: Dynamic Sparsity in Large Output Spaces](https://mlsanity.com/paper/2411.03171/) (2024)
- [SD2: Self-Distilled Sparse Drafters](https://mlsanity.com/paper/2504.08838/) (2025)
- [HASTE: Hardware-Aware Dynamic Sparse Training for Large Output Spaces](https://mlsanity.com/paper/2606.01117/) (2026)
- [NeuroTrails: Training with Dynamic Sparse Heads as the Key to Effective Ensembling](https://mlsanity.com/paper/2505.17909/) (2025)
- [Beyond Masked Sparsity: SNACK Enables Truly Sparse Neural Networks on GPU](https://mlsanity.com/paper/2610.04093/) (2026)

## Abstract (arXiv)

In classification problems with large output spaces (up to millions of labels), the last layer can require an enormous amount of memory. Using sparse connectivity would drastically reduce the memory requirements, but as we show below, it can result in much diminished predictive performance of the model. Fortunately, we found that this can be mitigated by introducing a penultimate layer of intermediate size. We further demonstrate that one can constrain the connectivity of the sparse layer to be uniform, in the sense that each output neuron will have the exact same number of incoming connections. This allows for efficient implementations of sparse matrix multiplication and connection redistribution on GPU hardware. Via a custom CUDA implementation, we show that the proposed approach can scale to datasets with 670,000 labels on a single commodity GPU with only 4GB memory.