ELMO: Efficiency via Low-precision and Peak Memory Optimization in Large Output Spaces

TL;DR

Extreme multilabel classification (XMC) with label spaces of up to millions places the classifier head as the primary memory and compute bottleneck. ELMO proposes a purely low-precision training pipeline using BF16 gradients and FP8 weights, augmented with gradient fusion, chunking, and a decoupled encoder–classifier update flow to achieve end-to-end training with dramatically reduced peak memory while preserving accuracy. The approach yields 4–6× memory reductions on a 3M-label model (down to $6.6\ \mathrm{GiB}$ with FP8) and scales to much larger label sets, with a new 8.6M-label LF-Paper2Keywords dataset demonstrating feasibility at extreme scales. Across public datasets, ELMO maintains competitive precision against state-of-the-art baselines, and its FP8 variant enables substantial efficiency gains, enabling practical end-to-end XMC training on commodity hardware in some cases. The work introduces the LF-Paper2Keywords-8.6M dataset to spur future research in ultra-large XMC benchmarks and highlights the potential for further bit-width reductions with carefully designed numeric techniques.

Abstract

Large output spaces, also referred to as Extreme multilabel classification (XMC), is a setting that arises, e.g., in large-scale tagging and product-to-product recommendation, and is characterized by the number of labels ranging from hundreds of thousands to millions. This means that the linear classification head, usually only a tiny fraction of the overall model, turns into the main driver for compute and memory demand. Current state-of-the-art XMC methods predominantly rely on FP16-FP32 mixed-precision training, which we show can be unstable, and inefficient in terms of memory usage and computational overhead. Meanwhile, existing low-precision methods typically retain higher precision for the classification layer. In this work, we propose ELMO, a pure low-precision training framework for XMC models using BFloat16 and Float8 data types. By leveraging Kahan summation and stochastic rounding, we demonstrate that XMC models can be effectively trained entirely in Float8, without relying on single-precision master weights or tensor scaling. Low-precision training, combined with our proposed memory optimizations -- gradient fusion and chunking -- enables significant reductions in GPU memory usage. For example, we train a 3-million-label XMC model with only 6.6 GiB of GPU memory, compared to the 39.7 GiB required by the optimized SOTA method, Renee without compromising accuracy.

Figures (5)

• Figure 1: Memory trace of Reneejain2023renee at 3 million labels & batch size 128, recorded with Pytorch profiler
• Figure 2: Figure \ref{['fig:sr_performance']}: Precision@1 performance at different exponent and mantissa bit patterns for the classifier weights. The numbers above diagonal is the performance when stochastic rounding is applied. Figure \ref{['fig:gradient']}: histogram of classifier gradients. Around 20% of gradients become zero in Float8 E5M2([-16, 15]), while nearly 90% of gradients drop to zero in Float8 E4M3([-9, 8]). The gradients are sampled from the training of LF-AmazonTitles-131K.
• Figure 3: GPU memory comparison between Renee and the proposed approach (ELMO) at various instances during one round of forward and backward pass. Note the difference in the scale of Y-axis in the two cases. The graphic was created using Pytorch memory viz utility.
• Figure 4: Comparison of peak GPU memory usage across varying label sizes for ELMO and Renee jain2023renee.
• Figure 5: Figures \ref{['fig:weights']} and \ref{['fig:clf_input']} show that most weights and classifier inputs fall within the exponent range of FP8 E4M3 ([-9, 8]), even without quantization. The weights and inputs are sampled from the training of LF-AmazonTitles-131K.