Towards Memory-Efficient Training for Extremely Large Output Spaces -- Learning with 500k Labels on a Single Commodity GPU
This addresses memory efficiency for machine learning practitioners working with large-scale classification tasks, though it is incremental as it builds on sparse connectivity methods.
The paper tackles the memory bottleneck in training classification models with extremely large output spaces by proposing a sparse connectivity approach with a uniform penultimate layer, enabling training on datasets with up to 670,000 labels using only 4GB of GPU memory.
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.