ICML 2025

ELMO : Efficiency via Low-precision and Peak Memory Optimization
in Large Output Spaces
Jinbin Zhang * 1 Nasib Ullah * 1 Erik Schultheis 1 2 Rohit Babbar 3
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 Ka-
han 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 (Jain et al., 2023) with-
out compromising accuracy. Code available at
https://github.com/xmc-aalto/elmo.
*Equal contribution 1Department of Computer Science, Aalto
University, Espoo, Finland 2IST Austria 3Department of Computer
Science, University of Bath, Bath, UK. Correspondence to: Jinbin
Zhang <jinbin.zhang@aalto.fi>.
Proceedings of the 42 nd International Conference on Machine
Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025
by the author(s).
1. Intro...