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

2025-10-13

AI Summary

To address the excessive GPU memory consumption and computational overhead caused by linear classification heads in eXtreme Multi-Label (XMC) classification, this paper proposes the first fully low-precision training framework for XMC: end-to-end training entirely in BFloat16/Float8, eliminating FP32 master weights and tensor scaling. We innovatively integrate Kahan summation with stochastic rounding to ensure numerical stability, and introduce gradient fusion and tiling strategies to optimize memory access patterns. On a 3-million-label benchmark, our method reduces GPU memory usage from 39.7 GiB to 6.6 GiB (an 83% reduction), while maintaining training stability and preserving model accuracy—outperforming existing mixed-precision approaches. This work marks the first successful end-to-end Float8 training for XMC, establishing an efficient, scalable low-precision paradigm for models with ultra-large output spaces.

Technology Category

• Machine Learning: Multi-class/Multi-label Learning & Extreme Classification
• Computer Vision: Large Vision Models
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization

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.

Problem

• Optimizing extreme multilabel classification with low-precision training
• Reducing GPU memory usage via gradient fusion and chunking techniques
• Training million-label models without accuracy loss using Float8 precision

Innovation

• Pure low-precision training with BFloat16 and Float8
• Kahan summation and stochastic rounding techniques
• Gradient fusion and chunking reduce memory usage