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

AI summary

Overview

Research area: Machine learning, specifically extreme multilabel classification (XMC) and low-precision (BF16/FP8) neural network training.

Technical level: Intermediate — the paper assumes familiarity with floating-point formats, mixed-precision training, and multilabel classification, but its core ideas are explained from first principles.

Scope: The paper presents ELMO, a training framework that runs both the encoder and the enormous classification layer of XMC models entirely in low precision (BF16 or FP8), combined with memory reorganizations that sharply reduce peak GPU memory at multi-million-label scale.

What This Paper Is About

In extreme multilabel classification, the output space can contain hundreds of thousands to millions of labels, which turns the normally tiny linear classification head into the dominant consumer of memory and compute. Prior state-of-the-art end-to-end XMC training (Renee) relies on FP16–FP32 mixed precision, which the authors show is unstable, and still requires roughly 39.7 GiB of GPU memory for a 3-million-label model. ELMO's goal is to train these models in pure low precision — BFloat16 and Float8 — while cutting peak memory several-fold and preserving classification accuracy.

Key Contributions

1. Purely low-precision training for large classification layers. The authors train XMC models entirely in BF16 and, further, in FP8 (E4M3 classifier weights with BF16 gradients) without single-precision master weights or tensor scaling, using stochastic rounding and Kahan summation to compensate for reduced precision.

2. Peak-memory optimizations. Gradient fusion (a custom Triton kernel that computes classifier gradients and applies the SGD update in SRAM, so classifier gradients are never stored in GPU memory) and chunking of the classifier's forward/backward/update steps (3 to 8 chunks) reduce peak memory by 4x–6x for a 3-million-label dataset.

3. Competitive accuracy against established baselines. ELMO is compared with LightXML, CascadeXML, NGAME, DEXML and Renee across Wiki-500K, AmazonTitles-670K, Amazon-670K, Amazon-3M, LF-WikiSeeAlso-320K, LF-AmazonTitles-1.3M and the newly introduced dataset, maintaining comparable or better Precision@k.

4. A new, larger public benchmark. LF-Paper2Keywords-8.6M, derived from the DBLP-Citation-network V14 dataset, has 8,623,847 labels, 2,020,621 training instances and 2,020,621 test instances — described by the authors as the largest publicly available XMC benchmark.

Main Findings

• Memory reduction at 3 million labels: ELMO requires 6.6 GiB (FP8) or 10.3 GiB (BF16) of GPU memory for a 3-million-label model with batch size 128 and a BERT-base encoder with 768 embedding dimensions, versus 39.7 GiB for the optimized state-of-the-art method Renee. The authors report an overall 4x–6x reduction for that label count.

• Memory scaling with label count: Peak GPU memory is plotted from 131K labels (LF-AmazonTitles-131K) up to 18 million labels (random labels appended beyond 8.6 million). At 3 million labels, ELMO (FP8) lowers memory by 6 times relative to Renee, rising to 11 times at 8.6 million labels and 13 times at 18 million labels. The paper also states the FP8 variant achieves a 13 times reduction compared to sampling-based methods.

• Initialization memory: Renee allocates about 17.9 GiB at initialization (8 GiB each for classifier parameters and momentum, plus a 687 MiB logit buffer and roughly 1.2 GiB for BERT and its optimizer states). ELMO allocates 3.2 GiB in FP8 (5.2 GiB in BF16) at initialization, a 70–80% reduction, because it removes the momentum buffer, stores classifier weights in 16-bit (4 GiB) or 8-bit (2 GiB), and reduces the logit gradient buffer to 86 MiB via 8-way chunking.

• Where the peak goes: Renee's peak comes from a persistent FP16 copy of classifier weights (an additional 4 GiB), classifier gradients allocated at 4 GiB and then upcast to FP32 at 8 GiB. ELMO avoids both, but its FP8 encoder adds 0.5 GiB of buffers, yielding the 6.6 GiB peak. Activation memory is 4.6 GiB in BF16 and 3 GiB in FP8 mixed precision.

• Precision requirements are asymmetric: Simulation on LF-AmazonTitles-131K shows 3 exponent bits suffice for classifier weights while 2 do not, and degradation sets in below 6 mantissa bits but is recovered by stochastic rounding — motivating the choice of E4M3 for weights without tensor scaling. For gradients, the histogram shows about 20% of gradients fall outside the Float8 E5M2 range ([-16, 15]) and nearly 90% drop to zero in Float8 E4M3 ([-9, 8]), so gradients must remain in BF16.

• Benchmark results (Table 2, selected): On Amazon-3M, ELMO (BF16) reaches P@1 53.4, P@3 50.9, P@5 48.8 with 10.39 GiB and 25:15 per epoch, versus Renee at 52.6 / 49.7 / 47.43 with 39.7 GiB and 29:58. On Wiki-500K, ELMO (FP8) records P@1 78.39 with 5.01 GiB and 11:28 per epoch, versus Renee at 78.69 with 12.69 GiB and 18:37. On LF-AmazonTitles-1.3M, Renee reaches P@1 56.04 with 19.9 GiB, while ELMO (BF16) reaches 56.14 with 6.61 GiB and ELMO (FP8) reaches 54.97 with 4.31 GiB.

• The 8.6M-label dataset (Table 3): ELMO requires 18.8 GiB (BF16) or 9.02 GiB (FP8), compared with 105.64 GiB for Renee and 58.44 GiB for a Float32 baseline; other XMC baselines did not scale to 8.6 million labels. BF16 ELMO (P@1 45.4) even exceeds the Float32 baseline (43.60), which the authors attribute to possible regularization effects of stochastic rounding. Renee scores only P@1 17.65, which the authors suggest is due to gradient overflow in the classifier input from FP16.

• Encoder precision trade-off (Table 4): With the classifier fixed in FP8, a BF16 encoder and an FP8 encoder give similar precision; on LF-AmazonTitles-1.3M the FP8 encoder took slightly longer (17:44 versus 17:26) and used less peak memory (4.63 GB versus 5.50 GB) — the added time attributed to FP8–BF16 mixed-precision recipe overhead.

• Hardware and software: All BF16 experiments were run on an A100 GPU; FP8 experiments on an H100 (Tables 2, 3, 8) and an RTX 4060Ti (Table 5). The encoder uses AdamW with Kahan summation via the optimi library, while the classifier uses in-place SGD with stochastic rounding implemented in custom Triton and CUDA kernels; the FP8 encoder comes from torchao.

Methodology in Plain English

The authors start by auditing why the previous leading method, Renee, uses so much memory. Renee keeps a full-precision master copy of weights for mixed-precision training, keeps classifier gradients in full precision, and lets memory-heavy operations pile up at one point in the training step. The authors then make a series of targeted changes.

First, they abandon mixed precision for pure 16-bit training. They switch from FP16 to BF16 because BF16 keeps the same exponent range as FP32, which avoids the overflow in classifier-input gradients and the underflow in weight gradients that plague FP16. Because BF16 has fewer mantissa bits, small updates can round away to nothing; they counter this with stochastic rounding for the classifier weights (rounding up or down probabilistically so updates are unbiased on average) and Kahan summation for the encoder (keeping a compensation term that tracks and corrects accumulated rounding error).

Second, they simplify the classifier's optimization: they drop the momentum buffer entirely and use plain large-learning-rate SGD. They reorder the computation so the classifier-input gradient is computed only at the end of the backward pass, after encoder activation memory has been freed. They split labels into k equal chunks (using 3 to 8) and run the classifier's forward, backward and update per chunk, cutting transient memory by a factor of k without harming latency.

Third, they push to 8 bits. Simulations with varying exponent and mantissa bits show that classifier weights need 3 exponent bits but can survive with fewer than 6 mantissa bits if stochastic rounding is used, so weights go to FP8 E4M3 with no tensor scaling. Gradients, however, frequently exceed FP8's representable range, so they stay in BF16, and inputs are cast to FP8 only for the matmul that produces logits. A custom Triton kernel handles the mixed FP8-weight/BF16-logit matrix multiplication, and another fuses the classifier gradient computation with the SGD update inside SRAM so gradients are never written to GPU memory. Finally, they adopt the torchao FP8 encoder so the whole pipeline is nearly pure FP8.

Why This Matters

Impact on research: The work shows that the classification layer of an XMC model — the part that everyone previously protected at higher precision — can itself be trained in FP8 without tensor scaling, and that end-to-end XMC training is feasible at label counts (8.6 million, and simulated up to 18 million) far beyond the 3-million-label ceiling of prior public benchmarks. It also releases a large new public benchmark and open code at https://github.com/xmc-aalto/elmo, lowering the compute barrier for reproducing XMC research.

Real-world applications (all listed in the paper as XMC use cases):

• Product-to-product recommendation, such as identifying frequently bought together items.
• Wikipedia article tagging with category labels.
• Matching search queries to advertisements.
• Automated keyword suggestion and paper recommendation for research articles.

Industry relevance: Reducing a 3-million-label training run from 39.7 GiB to 6.6 GiB, and per-epoch time on Wiki-500K from 18:37 to 11:28 (FP8), makes it possible to train extreme classifiers on far more modest hardware and to scale to label spaces that were previously impractical — directly relevant to recommender systems, ad targeting and large-scale tagging pipelines.

Future Directions

• Extending below 8 bits. The authors state that their recipe needs no tensor scaling, but their Figure 2(a) analysis indicates that moving to FP6 or FP4 datatypes will require such scaling strategies to be taken into account.
• Mitigating the FP8–BF16 mixed-precision overhead. FP8 encoders saved memory but were slightly slower per epoch than BF16 encoders, so removing that recipe overhead is an open engineering question.
• Scaling beyond the new 8.6M benchmark. The 18-million-label memory measurements were made by appending random labels rather than training on a real dataset of that size, so validation on genuinely larger label spaces remains to be done.
• Reconciling the accuracy trade-offs. The BF16 variant outperformed the Float32 baseline on LF-Paper2Keywords-8.6M while FP8 was slightly below it, leaving open the question of exactly when stochastic rounding's regularization benefit outweighs FP8's reduced precision.

Target Audience

Researchers and practitioners working on extreme multilabel classification, large-output-space retrieval and recommendation, and efficient large-scale training. It is also useful for engineers interested in FP8 training recipes, stochastic rounding, Kahan summation, fused optimizer kernels in Triton, and peak-memory profiling of very large layers. Readers should have some background in deep learning training loops and floating-point number formats.

Authors’ 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.