TOWARDS MEMORY-EFFICIENT TRAINING FOR EXTREMELY LARGE OUTPUT SPACES - LEARNING WITH 500 K LABELS ON A SINGLE COMMODITY GPU

Aug 3, 2026

Upload by Himbo

Classification with extremely large output spaces (up to millions of labels) makes the final layer prohibitively memory hungry. The document analyzes sparse connectivity as a way to reduce memory, noting that naive sparse training can substantially harm predictive performance. It proposes mitigating this by adding a penultimate intermediate layer, enabling dense-like results with sparsity. It further enforces uniform sparsity so each output neuron receives the same number of connections, allowing efficient GPU sparse operations. A custom CUDA implementation demonstrates scalability to 670,000 labels within 4GB memory.

LanguageEnglish

CategoryResearch & Report

FormatPDF

Pages13

Reading Time33 min

AccessFree to download

Abstract

Introduction

Setup and Background

arXiv :2306 .03725v 1 [ cs .LG] 6 Jun 2023 TOWARDS MEMORY-EFFICIENT TRAINING FOR EXTREMELY LARGE OUTPUT SPACES – LEARNING WITH 500 K LABELS ON A SINGLE COMMODITY GPU A PREPRINT Erik Schultheis Aalto University Helsinki, Finland [erik.schultheis@aalto.fi](erik.schultheis@aalto.fi) Rohit Babbar University of Bath / Aalto University Bath, UK / Helsinki, Finland [rb2608@bath.ac.uk](rb2608@bath.ac.uk) ABSTRACT In classi􀀂cation 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 ef􀀂cient 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. 1 Introduction In this paper, we present 􀀂ndings towards employing sparse connectivity in order to reduce the memory consumption of the classi􀀂cation layer for problems with extremely large output spaces (XMC). Such problems arise in, e.g., tagging of text documents [8], next-word predictions [22], and different kinds of recommendation tasks [5, 30, 1, 25, 20] . In order to ensure computational tractability of these tasks, which can have up to several millions of labels, one typically builds a hierarchical label tree [24, 33, 13, 31, 14], only exploring branches that are likely to contain relevant labels for the current instance. Even though this is very effective at reducing the computation (from linear in the number of labels to logarithmic), it does not help in addressing the memory consumption, which is still linear in the number of labels times the number of hidden units. As an illustration consider the AMAZON-3M [19] dataset. If we were to map the inputs to a hidden representation of 1024 units, the fully connected last layer for this dataset would need about 2.9 billion parameters, corresponding to 10.7 GiB 1 . Given that modern deep learning optimizers such as ADAM [16] need to keep track of the value, gradient, and 􀀂rst and second moment, this leads to an overall peak memory consumption of over 40 GiB, making it nigh impossible to train such models on commodity hardware. Therefore, we want to investigate possibilities for memory ef􀀂cient sparse training of this huge last layer. There are two pre-existing approaches that serve as an indication that this is an idea that could be successful: First, for DISMEC, a linear model applied to tf-idf representations of input text, it is known that the resulting layer can be sparsi􀀂ed after training to contain less than 1% non-zeros [2] . In a linear model, the different classi􀀂ers for each label can be trained independently. As a result, only the full weights of the label that is currently trained needs to be kept in memory, and can be pruned as soon as the training for that label has 􀀂nished. For non-linear models, the MACH [20] algorithm can be interpreted as a special case of training with static, random sparsity. It works by hashing the labels into different buckets, and performing training and predictions only on the level of buckets. If enough independent hashes are used, 1Assuming 32-bit 􀀃oating point numbers Training for Deep Extreme Classi􀀂ers on Single Commodity GPU A PREPRINT this method allows to solve the original problem in the large output space. However, in practice, the results presented for MACH are not as good as for competing methods. The contributions of this p

FAQs

Why is training difficult when the output space contains millions of labels?

The final classification layer requires an enormous amount of memory because its parameters scale with the number of labels and hidden units.

Why can naive sparse training reduce predictive performance?

Although sparse connectivity lowers memory usage, the resulting models often perform much worse than dense counterparts.

How does the proposed method improve both memory usage and accuracy?

It inserts a larger penultimate intermediate layer and constrains the sparse layer to uniform connectivity, enabling efficient GPU computation while maintaining predictive quality.

Aug 3, 2026

Upload by Himbo