ELMO : Efficiency via Low-precision and Peak Memory Optimization in Large Output Spaces
Authors: Jinbin Zhang
Nasib Ullah
Erik Schultheis
Rohit Babbar
Venue: ICML 2025 poster (2025)
Keywords: Extreme Multi-Label Classification
Low-Precision Training
Peak Memory Optimization
FLOAT8 training
Primary area: general_machine_learning->scalable_algorithms
TL;DR: We propose a pure low-precision training framework for XMC models using BFLOAT16 and FP8, achieving significant GPU memory savings  while competing baselines on most public datasets under low-bitwidth constraints.

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.7GiB required by the optimized SOTA method, Renee without compromising accuracy.

ELMO: Efficiency via Low-precision and Peak Memory Optimization in Large Output Spaces
Jinbin Zhang"
Nasib Ullah
Erik Schultheis
12
Rohit Babbar
Abstract
1. Introduction
Large output spaces, also referred as Extreme multil- abel classification (XMC) (Bhatia et al., 2016; Babbar & Schölkopf, 2017; Prabhu et al., 2018) refers to the task of predicting a sparse subset of relevant labels from an exceed- ingly large set, often ranging from hundreds of thousands to millions of potential classes. XMC has gained promi- nence due its applicability in real-world scenarios such as product recommendations, Wikipedia tagging, and match- ing search queries to advertisements. From the perspective of standard deep learning literature, this problem appears solvable using encoder such as CNN (Hu et al., 201 4), LSTM (Hochreiter, 1997; Chung et al., 2014), or more com- monly, a Transformer (Vaswani, 2017; Devlin, 2018) fine- tuned with a linear output layer in what is typically referred to as an end-to-end approach. However, contrary other domains, where transformer model such as BERT (De- vlin, 2018) would account for vast majority of all model parameters, in XMC, it is the classifier layer that becomes the bottleneck. For example, with an embedding dimension of 768, and three million labels, classifier weights alone would consume approximately GiB of memory. When accounting for gradients and optimizer states (Kingma Ba, 2014), the memory footprint expands around 32 GiB, Furthermore, loss computation in the last layer of the net- work, involving billions of parameters for larger datasets, also entails an enormous computational challenge.
Large output spaces, also referred to as Extreme multilabel classification (XMC), is 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-ard XMC methods predominantly rely on FP16-FP 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 BFloat1 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 gradienn fusion and enables significant reductions in GPU memory usage. For example, we train a 3-million-label XMC model with GPU memory, compared
Renee (Jain et al., 2023) shows that full end-to-end training can feasible with appropriate memory and computational optimizations in training the model, and results in classifi- cation performance superior to the approaches employing negative sampling (Jiang et al., 2021; Zhang et al., 2021; Kharbanda et al., 2022). It exploits the fact that, particularly for neural networks in large output spaces, the loss and gra- dient computation using automatic differentiation engines involves maintaining large buffers for intermediate states. Therefore, as long as one can directly compute the gradients required for the backward pass, explicitly computing the loss can be forgone, hence avoiding to materialize the mem- ory allocations for the intermediate variables altogether. In this way, Renee achieves a significant reduction in activation memory required for the classification layer.
only 6.6 to the 39.7 GiB required by the optimized SOTA method, Renee (Jain al., 2023) with- out compromising accuracy. Code available https:///ithhi hoo/xmc--ae wwwm///caaaatooaaoccccaaaaaaaaaa ceoc--llto/eeeeoo
'Equal contribution 'Department of Computer Science, Aalto University, Espoo, Finland *IST Austria 'Department of Computer Science, University of Bath, Bath, UK. Correspondence to: Jinbin Zhang <jinbin zhang @ filto.fi>.
Proceedings of the nnnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

ELMO: Efficiency via Low-precision and Peak Memory Optimization in Large Output Spaces
Despite the model optimizations in Renee, several memory bottlenecks remain. Firstly, Renee employs mixed precision (Micikevicius et al., 2018) training for the encoder, coupled with standard gradient scaling, which requires maintaining a full precision copy of the parameters. Furthermore, this approach mandates that input gradients be kept in full pre- cision, leading to the classifier layer's gradients also being cast to full precision, which further inflates memory usage. Secondly, through memory snapshot analysis 3, we observe that in Renee, the order of execution in the computation graph causes memory-intensive operations to accumulate at a single point in time, leading to excessive peak mem- ory demand. Thirdly, neither Renee, nor any of the label shortlisting approaches, result in a reduction in memory requirements for the classification layer weights.
2012; 2007; Sinha et al., 2015). On this larger dataset, our low-precision training method demonstrates significant memory and computational savings. We believe that low- precision training will become a standard in the XMC field, given the growing demand for handling datasets with an immense number of labels.
To summarize, this paper makes the following contributions: (I) We introduce purely low-precision training approach using BF 6 and FP8 for models with large classification layers. (2) Combined with peak memory optimizations,
this reduces memory usage by 4x-6x for a 3-million dataset; (3) Apart from efficiency gains, we compete with existing XMC baselines on most public datasets, illustrat- ing that purely low-precision training can preserve similar performance: (4) We introduce LF-Paper2Keyyordd-88.. 6.,,, 8.6- 2.6-million-labll dataset that, to our knowledge, is now the largest publicly available XMC benchmark.
To address these challenges, we first move from mixed precision to pure 16-bit training, achieving concurrent re- ductions in both memory and computation time. We adopt BF BFloatl 6 (BF1 6) for gradients, leveraging its extended range to mitigate the training instability caused by potential overflows. To compensate for precision loss, we employ Kahan summation (Kahan, 1965) for the encoder optimizer and stochastic rounding (Zhang et al., 2018 8) for the classifier optimizer, compensating for inaccuracies and rounding bias during parameter updates. Going further, we demonstrate that classifiers can be trained in pure Fl Float& (FP8) with- training, as long as gradients
2. Related Work
We presented discussion of Renee, the most relevant XMC method for our work, in the previous section. Next, we provide brief overview of different categories of other XMC methods and low-precision training.
Extreme Classification. Initial approaches in extreme clas- sification utilized linear classifiers with bag-of- and TF-IDF features (Babbar& & Schölkopf. 2017; 2019). As the computational cost of these methods scales linearly with the number of labels, tree-based methods were introduced to reduce computation complexity to logarithmic scale with respect to label count (Prabhu et al., 2018; Khandagale et al., 2020; Wydmuch et al., 2018; Jasinska- Basinska-Kobus et al., 2021). Subsequent advancements incorporated task-specific feature learning using deep encoders atop label trees, under the as- sumption that joint training of both encoder and classifier layers would be computationally expensive without some form of label shortlisting for negative sampling when eval- uating the loss function in the last layer (You et al., 2019; Jiang et al., 2021; Kharbanda et al., 2022; Zhang et al., 2021; Kharbanda et al., 2023). While tree-based methods provided a negative sampling mechanism, alternative ap- proaches employed nearest-neighbor approach for negative sampling and multi-stage training to first train the encoder, followed by the classifier (Dahiya et al., 2021; 2023a; Mittal et al., 2021; Dahiya et al., 2023b; Kharbanda et al., 2025). DEXML. (Gupta et al., 2024) eliminates the classifier using dual-encoder training with all-negative labels, but at the cost of increased compute and memory usage. More recently, Renee (Jain et al., 2023) demonstrated the fcasibility of a fully end-to-end approach while optimizing memory con- sumption, revealing its potential to outperform traditional sampling-based methods. Our work advances this end-to- end approach, focusing enhancing computational and
out scaling or mixed-precisiee remain in BF16. Finally, by integrating a FP8 encoder from torchao, we achieve a nearly pure FP8 (excluding gradients of the transformer backhone) training pipeline for XMC tasks.
Adopting pure BF16 6 and FP8 training reduces parameter memory requirements by 50-75% relative to full precision. In order to address memory accumulation in Renee, we re- organize the computation flow and decouple encoder and classifier updates, leading to more evenly distributed mem- ory allocations throughout each training iteration and re- ducing peak memory usage. Furthermore, our chunking strategy for classifier updates curtails transient memory de- mands, and by fusing classifier gradient computation with the optimizer step in a custom Triton kernel, we effectively eliminate the need to store classifier gradients in memory. With these optimizations, the proposed method, ELMO, requires only 10.3 GiB' (BF19 6) or 6.6 GiB (FP8) of mem- ory for 3-million- labell model, significantly lower than 39.7 GiB necessitated by Renee.
We evaluate our low-precision training method on datasets of varying labels' sizes, demonstrating results comparable to Renee. To further assess the efficiency of our approach, we derived a new dataset with 8.6M labels from the DBLP- Citation-network V dataset (Tang et al., 2008; 2010; 2011;
1 1BBB & 0.93GBB
2

---

ELMO: Efficiency via Low-precision and Peak Memory Optimization in Large Output Spaces
memory efficiency by leveraging recent developments in low-precision training.
representation, defined in IEEE 754 (IEE, 2019), is speci- fied by the number of exponent bits E,, mantissa bits M. the format includes special bit patterns for values like inf iafinitiess NaNs, and subnormals, which repre- sent very small magnitudes near zero. The common 32-bit floating-point (FP format uses 23 mantissa bits and 8 exponent bits. Lower-precision formats, like FP16 (half- precision) and BF16 (BF16), reduce mantissa and/or expo- nent bits. BF 6, for example, retains the 8 exponent bits of FP32 but has fewer mantissa bits, offering similar dy- namic range with reduced precision. Recently, FP8 8 formats have also been proposed, with 4 or 5 exponent bits and 3 or mantissa bits (referred to as E4M3 and E5M2 respec- tively), using the IEEE 754 structure but differ in how they handle special values.
Low Precision Training. With the success of foundation models (Brown et al., 2020: Achiam et al.. 2023; Zhang et al., 2022; Touvron et al., 2023), there has been a strong push toward scaling both model size and training data. Nev- ertheless, this scalability is limited by memory and compu- tational constraints, where low-precision training provides notable benefit. Micikevicius et al. (2018) introduced method for training in mixed precision (FP1 6-FP32) by maintaining a master copy in full precision and using loss scaling to counteract gradient underflow. The restricted range of FP16 6 may cause instability in large models (Rae et al., 2021; Zeng et al., 2022), necessitating a transition to BF 16 6 F2 mixed precision (Rae et al., 2021; Smith et al., 2022) for enhanced stability. An alternative to mixed- precision training, Zamirai et al. (2020) employs pure BF1 6 with methods such as Kahan summation (Kahan, 1965) and stochastic rounding (Forsythe, 1950; Croci et al., 2022) to reduce rounding errors, while COLLAGE (Yu et al., 2024) utilizes pure BF 6 with a colpi-componett floating-pont point representation (Yu et al., 2022) for enhanced precision. The FP8 format is promising next step for further reduction in precision, with early successes in using tensor scaling, as demonstrated by Micikevicius et al. (2022), and hybrid formats (Sun et al., 2019) for weights and gradients with a different exponent bias. Currently, FP8 8 support (torchao maintainers & contributors, 2024) is limited primarily to matmul( () operations. Towards this Peng et al. (2023) introduces a comprehensive FP8 training suite, including an FP8 optimizer, communication, and distributed train- ing capabilities. Another line of work focuses on reducing memory in Adam optimizer states (Dettmers et al., 2022: Zhao et al.) or using stateless optimizer (Lv et al., 2024). In the XMC setup, recent works have successfully achieved compression of the last layer using dynamic sparse training (Schultheis& & Babbar, 2023; Ullah al.). As an orthogonal approach, in this paper, we focus on now-precision, specif- ically FP8 training for XMC models, where the classifier can benefit significantly from the memory and computation efficiency of FP8 representation.
The reduction in exponent and mantissa bits in lower- precision formats introduces quantization errors. Clipping error occurs when the exponent range is insufficient, caus- ing values outside the representable range to be "clipped" to the nearest maximum or minimum representable value. Meanwhile, rounding error arises from having fewer man- tissa bits, which forces values to be rounded the nearest representable value within the available precision. As a result, small weight updates are canceled due to nearest rounding (Zamirai et al., 2020).
Stochastic Rounding and Kahan summation. For any finite subset F of the reals R. the stochastic round- ing (Forsythe, 1950; Croci et al., 2022; Zamirai et al., 2020) of x E R, is defined as follows:
SR(x)= = { [>]
with probability p(x)
(I)
[x]]
with probability 1 - p(())
where [x] = min(= E F = 2 x), [x] = max f= E
F : 2 S x} denote rounding up and down, respectively, The probability is based on the distance p(x) --|-|=] While stochastic rounding does not make any individual rounding operation more accurate, its result is an unbiased estimate of the true number. Thus stochastic rounding can prevent the catastrophic accumulation of rounding errors, e.g., when adding many small numbers sequentially to one large number, as likely to happen when adding small gradient updates to the network's 's weights.
-- |x|
T
3. Problem setup and preliminaries
Problem setup For a multi-label dataset with N training samples, D = {(x,, P.) J: L as the total number of labels, and P: C [z]] denotes subset of relevant labels associated with x; E X' such that |P/ < Vi. Typically, the in- stances are text based, such as the contents of an article on Wikipedia or the description of a product on Amazon with labels being Wikipedia categories and frequently bought together products, respectively (Bhatia et al., 2016).
An alternative is to employ Kahan summation (Kahan, 1965). In this case, an additional buffer keeps track of the rounding error, and is used to correct subsequent additions:
y t- v C c -- ((s y) 5) - + y,
where s is the current value of the sum, cs the compensation term, and v the number to be added.
Floating-poin point formats The standard binary floating-poii
3

---

ELMO: Efficiency via Low-precision and Peak Memory Optimization in Large Output Spaces
At first glance, this might seem counterintuitive, by re- quiring the compensation term e, the memory benefit of the smaller representation for is negated. However, if s were kept in high precision, that would require making an ad- ditional low-precision copy for fast matrix multiplications, which is unnecessary with Kahan summation.
mediately reduce that by at least 8 GiB. In the next section, we will present such change, together with further memory optimizations, before finally reducing parameters to FP8 8 for even more memory savings.
4. ELMO : Low-precision and Peak Memory Optimization in XMC
12.00
Grad FP16
4.1. Pure 16-Bit Training
Grad FP32
The most obvious inefficiencies visible in Figure 1 are the multiple copies of weights and gradients in different preci- sions. Thus, our first step towards memory-efficient training is the implementation of pure 16-bit training.
FP16 weights (e)
FP weights (f-
Bert Activations
1500
Momentum huffer
Classifier weights
This is a bit more involved than just changing the datatype of the weights tensor, as two problems inherent in lower bit- width representations need to be overcome. FP1 6 numbers, as used in Renee's 's mixed precision training, have a much lower range of representable values than the single preci- sion baseline. Typically, this manifests in underflows during the backward pass, leading to gradients being rounded to zero, and is mitigated by loss-scaling techniques that shift the overall range of gradients into the representable inter- val (Micikevicius et al., 2018).
Figure I: Memory trace of Renee (Jain et al., 2023) at 3 million labels & batch size 128, recorded with Pytorch profilerr
Shortcomings of mixed-precision training in Renee Mixed precision training (MPT) aims at improving through- put by ensuring that the most compute-intensive operations, matrix multiplications (matmul), are performed in lower- precision representations that enjoy accelerated hardware implementations. This means that weights are kept at their full precision, and a second, ephemeral low-precision copy of the weights is created for the purpose of matmul. In contrast, activations are consistently kept in lower precision, and the Fl P1cat16 gradients are cast into Float 32 during weights update.
Gradient computation for the classifier input (i.e. classi- fier logit gradient x classifier weights), however, involves matrix multiplications over a large inner dimension, corre- sponding to the number of labels. This large accumulation over labels often results in overflow outside the FP1 range. At the same time, gradients of the classifier weights them- selves do not grow in magnitude as the label space increases. This contrasting nature (overflow & underflow) of gradient behavior makes Renee's 's training unstable and sensitive to both encoder and label space sizes despite using mixed precision with loss scaling (see Section 5 in Renee).
Consequently, mixed precision training greatly benefits mod- els with moderate parameter sizes across multiple layers (e.g.. layers in the BERT-base encoder), and leads to a sig- nificant reduction in memory requirements if they are dom- inated by activation memory. However, its application to single layer with extremely large parameter sizes signifi- cantly increases the peak memory consumption because it creates a second (albeit smaller) copy of the huge classifi- cation weights. Specifically, the memory trace for Renee presented in Figure shows that the low-precision copies of the classifier weights persists for the entire step.? Addition- ally, we can see that the gradient is first calculated in 16-bit precision, but then upcast 32 3
One possible solution is to adopt separate scaling factors but this requires sweeping over scale combinations (Blake et al., 2023). Even simpler, though, overflow and underflow can be easily addressed by switching from FP 6 to BF 6, sac- rificing some precision for an extended range of values that matches FP32. This removes the additional high-precision copy of the classifier's gradients.
The reduced precision of BF1 6 comes price. If one just replaces the 32-bit versions of weights in the optimizer with corresponding BF 1 6 weights, training can no longer progress successfully. This is because round-to- -earesst rounding in the optimizer update can end up canceling the update step, if it is less than half the distance to the next representable number.
As consequence, Renee's 's memory consumption is con- siderably higher than one might naïvely suspect, especially at its peak, with total of about 40 GiB at 3 million labels. Switching to pure BF oat:1 6 training and fusing the SGD update so it does not need an additional buffer would im-
There are two common ways for dealing with this prob- lem, and we employ both. For the classifier weights, where memory efficiency is paramount, we use stochastic round-
Thiis particular prohlem could casily he fixed hy de 'ing them after the calculating of the classifier layer's gradient. Renne official code
4

---

ELMO: Efficiency via Low-precision and Peak Memory Optimization in Large Output Spaces
Classifier gradient
39.40
41.60
44.00
45.22
45.43
60.55
45.00
12.0%
Percentage the classifer gradient
42.85
41.70
43:23
45.23
40.58
10.0%
Exponent Bits
77.00
12
42.80
35
8.0%
37.50
45.43
22
6.0%
41.50
44.10
43.23
45.27
45.43
49.38
43.90
4.0%
3.49
41.29
64.47
4430
40.30
45.50
45.39
Mantissa Bits
23-22-22 277--2120--100
19-18-17-16-
Exponent
(a) Precision@ (Without With SR)
(b) The histogram of the classifier gradient
Figure 2: Figure 2(a): Precision 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 2(b): histogram of classifier gradients. Around 20% of gradients become zero in Float8 #5M2([-16, 15]), while nearly 90% of gradients drop to zero in Float8 E4M3([-9, 8)). The gradients sampled from the training of Ain-emzzoTTills-1133
ing (Forsythe, 1950; Croci et al., 2022; Zamirai al., 2020) as defined in (1). This reduces the amount of memory re- quired for classifier weights to one third of the Renee set- ting and halves the optimizer state. The weights of the Bert encoder, in contrast, only make up a tiny fraction of the total memory consumption. Therefore, we opt use Kahan summation (Kahan, 1965) for the encoder's parameters.
consumption; however, as label size increases, memory allo- cated for classifier logits and gradients becomes the primary bottleneck. We employ chunking (Rabe & Staats, 2021; Hsu et al., 2024) for classifier parameters: We first execute the encoder's forward pass, then divide the labels into equal-sized chunks, processing the classifier's forward pass, backward pass, and optimization step sequentially for each chunk. The encoder's backward pass and update occur after all classifier operations are complete. This reduces transient memory requirements for classifier operations by factor of k. We use between 3 and chunks for XMC datasets, observing no impact on training latency (see Appendix F).
4.2. Architectural improvements
After fixing the most egregious inefficiencies stemming from standard mixed-preciste training. we now fur- ther architectural enhancements to reduce the (peak) mem- ory consumption of the training algorithm.
4.3. 8-Bit Training
noticeable chunk of the remaining memory taken up by the momentum buffer for the classifier weights. Our experi- ments showed that momentum is not required so we remove entirely, switching to pure large-learning-eee SGD for the classifier layer.
For further memory savings, we turn towards reducing the number of bits allocated to each parameter further. To that end, we first investigate how much precision and dynamic range are required for the training process, by simulating floating ppint numbers with specific number of mantissa and exponent bits.
Finally, there remains the spike in memory consumption due to the classification layer's gradient. There are several ways this could be addressed. A first option would be to change the order of operations and calculate this gradient only at the end of the backward pass, when activation memory for the Bert encoder has already been freed. While this does not actually reduce the memory requirements for this operation, it is moved point in time with less memory pressure, thus reducing peak memory consumption.
Figure 2(a), which illustrates the performance for LF- AmazonTittes-1333 dataset in terms of precision@1 k. shows that exponent bits provide sufficient range to represent classifier weights, whereas 2 bits are not enough. In the man- tissa, see significant degradation starting to set in shrinks below 6 bits, but this be counteracted by stochas- tic rounding, which recovers the original performance. This leads us to choose the E4M3, which is directly supported in Hopper, Ada, and Blackwell 11 GPUs, for representing
Reorganizing computation flow helps alleviate peak memory
5

---

ELMO: Efficiency via Low- Low-precision and Peak Memory Optimization in Large Output Spaces
Renee
ELMO
Computation Flow
Encoder Forward pass
Encoder
Inttallzttion
Forward pasa
Bockwarrd and optimizer stapa
Inittaiizatio
dasssifier clturedd upáote (churik sizeni))
beckward para
F1
B1 B2 B3 B4
B5 02
F1
B2- B3
+002
XMC Momentum (8 GB
12
JMCC Weight (208)
Detailed Computation Flow
0f
XMC Weight
12
Weight (a GB)
SGD step
[BGB)
ChunkedXM
Chunked and Fused XMC
Weight GB)
H
Ao((loatt6)
XMC Weight (4GB)
Embed
XMC Wwight Grad (4GB)
B4
B3
Aoffleat32)
BR0 Loss
Encoder
Weight [2IK Gi)
grad (logits) B1
logits
SGD
step
Embed
B3
Chunkoo Skip Love
Encoder
grad (logit))
B1
Input grad
B5
B2
logite
K times classifier chunking update
48.0
Memory Snapshot
B3
3.0 G
4.0
16.0
11:
G.0
5.000
0.00
0.00
Point momory piling
Figure 3: GPU memory comparison between Renee and the proposed approach (ELMO) 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.
classifier weights; notably, we do not introduce any addi- tional tensor scaling (Micikevicius et al., 2022), the native dynamic range of E4M3 sufficient, c.f. also Figures 5(a) and 5(b).
GPU memory. This eliminates the need to store classifier gradients in GPU memory, reducing its memory footprint to nearly zero
Further, we take on the recently developed torchao (tor- chao maintainers & contributors, 2024) framework for the encoder training, which leverages FP training for trans- formers, reducing the encoder's activation memory require- ments compared to BF1 6 training. By combining FP 8 train- ing for both the encoder and classifier, end-to-end training of FP8 XMC models becomes feasible.
In contrast, Figure 2(b) indicates that exponents of around 20% of the gradients exceed even the representable range of FP8 E5M2, necessitating the use of BF16 6. To address the different precision requirements for weights (in FP8) and gradients (BF16), we cast classifier inputs from BF16 to FP8 E4M3 when computing logits, as FP8 sufficiently cov- the inputs range (Figure 5(b)). We then perform matrix multiplication between FP8 E4M3 inputs and FP8 E4M3 weights, but obtain logits in BF 6 for higher-precision gra- dient computation. For input gradient computation, which involves matrix multiplication between FP8 weights and BF16 6 logits, we developed a Triton (Tillet et al., 2019) ker- nel to manage this, avoiding additional HBM memory usage due to data type differences.
4.4. Comparison of ELMO and Renee
Finally. let us consider training a 3 million label model with a batch size of 128, using Bert-base as an encoder with 768 embedding dimensions. How much memory do the different parts of the models take? We use the above Figure 3, which shows a comparison of the memory snap- shot for the two approaches, and corresponding order of operations during initialization, forward and backward pass.
To reduce the memory overhead caused by the gradient, ELMO fuses gradient computation and SGD updates into a single Triton (Tillet et al., 2019) kernel, optimized for execution entirely in SRAM. Within this kernel, classifier weights, logits, and inputs are loaded from GPU memory into SRAM. Classifier weight gradients are computed via matmul operation between the logits and inputs. The classifier weights are then updated directly in SRAM using SGD with stochastic rounding before being written back to
At initialization (denoted by II, 12 etc. in Figure 3), Re- nee allocates memory for encoder parameters, its optimizer states, and a buffer store logit gradients for the labels in last layer. For this example setting, parameters and momen- tum for the classifier layer amount to 8 GiB each, the logit
18 2..122812.8113 T024X X 1024
cuthed Jabels py:
1024X 1024 1024
6

---

ELMO: Efficiency via Low-precision and Peak Memory Optimization in Large Output Spaces
Table 1: Statistics of XMC Datasets with and without Label Features. This table presents a comparison across various datasets, detailing the total number of training instances (N), unique labels (L), number of test instances (N"), average label count per instance (T.), and average data points per label (f).
Dataset
N
L
N
L
i.
Datasets without Label Features
Wiki-500K
1,779,881
501,070 670,091
769,421
4.75
16.86
AmazonTitles-66
185,176
150,875
5.39
5.11
Amazon-670K
490,449
670,091
153,025
5.45
3.99
Amazon-3M
1,717,899
2,812,281
742,507
36.17
31.64
Datasets with Label Features
294,805
131,073
134,835
5.15
2.29
LF- Als-SSee
693,082
312,330
177,515
4.67
2.11
AmazzonTitee-LL.33
2,248,619
1,305,265
970,237
22.2
38.24
2,020,621
8,623,847
2,020,621
9.03
2.12
buffer consumes 687 MiB. ELMO gets rid of the momen- tum buffer altogether, and allocates weights in either 16-bit (4 GiB) or 8-bit (2 GiB). Due to chunking, the size of the logits' gradients gets divided by the number of chunks, in this case 8, but stays in 16-bit representation for both BF16 and FPB training, leading to 86 MiB. Additionally, the Bert model and its optimizer states are allocated. These are the same size for both ELMO and Renee, amounting to Rs 1.2 GiB. In total, ELMO allocates 3.2 GiB (5.2 GiB for BF16) at initialization, 70-80% reduced compared to Renee's s 17.9 GiB.
bottlenecks that become significant at larger label sizes,
We curated LE- Paper2Keywords-88 6M using the DBL.PCitation-networr VI4 dataset (Tang al., 2008; 2010: 2011; 2012: 2007; Sinha al., 2015). This dataset comprises the titles and abstracts of research papers sourced from DBLP, ACM, and MAG (Microsoft Academic Graph), with each paper's keywords serving as labels. The dataset includes 8.6 million labels and can support tasks such as automated keyword suggestion and paper recommendation for research articles. Complete dataset details are provided in Table
JF-Paper2Keyword--S.MM.
During the forward propagation steps (denoted as F1, F2 etc. in Figure 3), activation memory accumulates for the Bert transformer, 4.6 GiB in BF 6 and 3 GiB in FP8 mixed pre- cision. For the classifier layer, Renee also needs to create the FP1 116 6 copy of its weights, an additional 4 GiB. During the backward pass Renee (B1, B2 etc.), gradients are allo- cated (4 GiB) and converted to FP32 GiB, with all these allocations stacking up peak memory consumption of 39.7 GiB. Much more efficiently, ELMO does not need to make copy of the classifier weights, nor does it materi- alize classifier gradients. However, the FPS encoder uses additional buffers of 0.5 GiB, bringing the peak memory consumption to 6.6 GiB.
6. Experiments and Discussion
Datasets. We validated our approach on broad suite of XMC datasets spanning both non-label-feature based (Amazon-670K, Wiki-500K, Amazon-3M, AmazonTitles 670K) and label- (LF-AmazonTTttes-- 131K, LF- WikiiSeee LF-AAmazoTTilles-1...M and our newly curated w.F=-apprr2KywwodssS.kkk All except L. Puper2Keywords-88 66M are publicly accessible through the Extreme Classification Repository (Bhatia et al.. 2016). De- tailed descriptions are provided in the Table 1.
Baselines and Evaluation Metries. We compared our method with two categories of baselines: (i) Sampling- Based XMC, focusing on Transformer-based methods (c.g.. LightXML (Jiang et al., 2021), CascadeXML. (Kharbanda et al., 2022)) for non-label-featur datasets and NGAME (Dahiya et al., 2023a) and DEXML (Gupta et al., 2024) for label-feature datasets; (ii) End-to-End XMC, with Renee (Jain et al., 2023) as the principal baseline. Following stan- dard XMC practices, we evaluate all methods using top-k metrics, specifically Precision@A k and its propensity-scom variant. A detailed overview is given in the Appendix A.
5. Contributed Dataset
Motivation. Performance evaluation for XMC algorithms traditionally relies on public benchmarks (Bhatia et al., 2016), where the largest dataset contains 3 million labels. Consequently, many recent large-scale XMC experiments have employed proprietary datasets (Mehta et al., 2024; Jain et al., 2023; Dahiya et al., 2021). A more expansive public dataset would facilitate the exploration and comparison of modern, efficient training strategies, while also revealing

---

ELMO: Efficiency via Low-precision and Peak Memory Optimization in Large Output Spaces
Table 2: Comparison of the precision performance of our proposed ELMO method with state-of-the-art XMC methods on the Wiki-500K, Ais and FF-AmazonTitless fittle--.......... 1.3M. Bold indicates the best results, and underline indicates the second-hest. Mr denotes peak training memory and METHOD-N denotes method with ensembles.
Amazon-670K, and Amazon-3MV -MM datasets, äs well as the label feature datasets
Mr (GiB)
Epoch Time (mm:ss)
Mr (GiB)
Epoch Time
Method
Encoder
P@1
P@3
P@5
P@
P@3
P@S
(mmsss)
Wiki-500K
Amazon MamazonTitls--oooo Titles-670K
LIGHTXML CASCADEXML
BERT Basc
76.19
57.22
44.12
15.72
147:27
41.7
37.3
34.2
13.99
19:02
BERT-Bas
77.0
58.3
45.1
18.8
50:00
42.1
37.5
34.1
22.3
11:32
LIGHTXML-3
BERT Base
77.78
58.85
45.57
15.72
147:22
43.1
38.7
35.5
13.99
3 19:02
CASCADEXML-3
BERT-Base
78.39
59.86
46.49
18.8
3x50:00
43.5
39.0
35.6
22.3
3x 11:32
RENEE
BERT Base
78.69
60.03
46.46
12.69
18:37
43.78
39.17
35.91
12.46
2:04
ELMO (BF16)
BERT- BERT-Base
78,61
60.04
46.59
7.21
17:01
44.3
39.7
36.4
5.12
1:47
ELMO (FP8)
BERT Base
78.39
59.64
46.12
5.01
11:28
44.39
39.75
36.43
3.37
1:18
Amazon-670s
Amazon-3M
LIGHTXML
BERT BERT-Basc
47.3
42.2
38.5
11.5
53:30
OOM
CASCADEXML
BERT Base
48.5
43.7
40.0
18.3
16:46
51.3
49.0
46.9
87.0
90:00
LIGHTXML-3
BERT-Basc
49.1
44.17
40.25
11.5
3x53:30
OOM
CASCADEXML-3
BERT
50.22
45.20
41.45
18.3
3x16:46
53.10
50.64
48.49
87.0
3x90:00
RENEE
BERT Base
50.6
45.16
41.13
11.91
7:14
52.6
49.7
47.43
39.7
29:58
ELMO (BF16)
BERT Base
50.7
45.27
41.29
5.29
6:00
53.4
50.9
48.8
10.39
25:15
ELMO (FP8)
BERT-B.n Base
50.34
44.91
40.97
3.3
4:05
52.73
50.38
48.27
6.6
18:02
WikiSeeAlso-333
LF-Amazon Titles 1.3M
LIGHTXMI.
Distil-BERT Distil-BRT
34.5
22.31
16.83
24.46
61:09
0OM
NGAME
45.72
29.61
22.06
16.63
57:19
54.99
48.09
43.11
11.03
40:09
DEXML
Distil-BERT
46.78
30.32
22.59
38.6
242:51
58.4
45.46
75.53
1054:21
RENEE ELMO (BF16)
Distil-Bee BERT
47.86
31.91
24.05
13.89
10:50
56.04
49.91
45.32
19.9
9:12
Distil-BERT
47.84
31.99
24.12
6.57
10:08
56.14
49.86
45.25
6.61
9:10
ELMO (FP8)
Distil-BERT
47.88
31.92
24.09
5.2
6:22
54.97
48.41
43.82
4.31
17:44
Precision performance comparison on LF- Paper2Keywords-88 dataset. Other XMC baselines do not scale to 8.6 million label size.
Table 3:
plemented via custom Triton and CUDA kernlss. We also employ Triton kemelss for the FP 8 classifier and gradient fusion. All BF1 6 experiments are conducted on an A 1100 GPU, while FP8 experiments are run on an H100 (Table 2, and 8) and RTX 4060Ti (Table 5) GPU. Further details implementation and hyperparameters provided in the Appendix G.
Method
P@ 1
P@3
P@5
MGiB)
FLOAT32
43.60
32.13
26.02
58.44
RENEE
17.65
11.78
9.23
105.64
Empirical Performance. Table 2 reports Precision @1 re- sults for ELMO on non-label-feature datasets, benchmarked against state-of-the-an XMC baselines. Notably, ELMO significantly reduces peak memory usage-its FP8 variant achieves 6 times reduction relative to Renee (Jain et al., 2023) (the current end-to-end optimized approach) and baaeed based methods. Despite these efficiency gains, ELMO shows competitive or improved Precision @ k performance. For label-fcature datasets, we apply the standard augmentation (Kharbanda et al., 2024; Jain et al., 2023) strategy summarized Ta- ble (LF--Wik WikiSeeAAlso3220 LF-AmazonTitllss 1.3M) and Table L.F-AmazonTiiless 131K), Our results demonstrate
ELMO 6)
45.4
33.58
27.18
18.8
ELMO(FP5)
43.4
31.59
25.38
9.02
Implementation Details. Low-precision training with chunking is implemented using the PyTorch framework (Paszke et al., 2017). For the encoder, we used A (Loshchilov & Hutter, 2017) optimizer with Kahan summa- tion provided by the iptimi imi S The in-place SGD optimizer with stochastic rounding for the classifier is im-
13 times reduction compared sampling
Shttps: thub. com rner-benjamin h/optimii
8

---

ELMO: Efficiency via Low-precision and Peak Memory Optimization in Large Output Spaces
that ELMO maintains robust performance in both FP8 and BF16 modes, underscoring the versatility of our approach.
Table 4: Comparison of precision performance between the BFloat16 and Float 8 encoders with the classifier fixed at Float8.
Finally, on our newly introduced LF-Paper2KKewoord-2 8.6 dataset as shown in Table 3, ELMO provides substan- tial memory savings, requiring only 18.8 GiB (BF16, 6) or 9.02 GiB 8F compared to 105 GiB for Renee. Notably, BF16 ELMO even outperforms the Fl baseline, likely benefiting from the regularization effects of stochastic rounding (Ozkara et al., 2025). FP8 also delivers perfor- mance close to that of Float32. Renee underperforms, likely due to gradient overflow in the classifier input caused by its use of FP16 6 data types.
Mi. (GB)
Epoch Time (mm:ss)
Encoder
P@I 1
P@3
P@5
AFF-Amazon Titles- 1.3M
BF16 FP8
55.08 54.97
48,47
43.87
5,50
17:26
48.41
43,82
4.63
17:44
Amazon-3MM 3M
BF 6 FP8
52.60
50.23
48.11
8.51
15:56
52.73
50.38
48.27
7.16
18:02
EL (BF1 CFhukkhh ELMD (TFB4 Cag.naaeeee....oom
SHA100,
Berre
Peak Memorr
250
yond the conventional practice of relaxing the classifica- tion layer's precision. In contrast to FP1 6-FP 32 MPT which can be memory-inefficient and unstable, our method is robust and, with gradient fusion and chunking, reduces memory by 4x-6x for a 3-million- dataset. Notably, these efficiency gains do not compromise performance; approach competes with existing XMC baselines most public datasets. Furthermore, we introduce new XMC dataset TT-Paper2Keyywords-8 -...doddd-.6M.. with 8.6 million labels, which, upon its release, will be the largest publicly avail- able XMC dataset. We anticipate that this new dataset will spur further innovation in extreme classification. While our proposed training recipe does not require (micro-) scaling, our investigation in 2(a) indicates that future work aiming to further reduce representation bitwidth to FP6 or FP4 datatypes will have to take such strategies into account.
3+A180 M100
150
2xh180,N100
50
A1:004000
0.14
2,5
18.0H
Label Size (No. of Labels)
Figure 4: Comparison of peak GPU memory usage across varying label sizes for ELMO and Renee (Jain et al., 2023).
Label Size vs Peak GPU Memory. Figure plots peak GPU memory usage as label sizes grow from 131K (LF- AmazonTTtles-1111 to 18 million. While the largest public dataset has 3 million labels, we introduce one with 8.6 million. Beyond 8.6 million, random labels are appended to measure peak usage. 3F16 6 and FP8 significantly reduce GPU memory compared to Renee (Jain et al., 2023); for instance, at million labels, ELMO (FP8) lowers memory by 6 times, increasing to 11 times at 8.6 million and times 18 million.
Impact Statement
We do not anticipate any negative societal impact of our work. It is expected that the affordable training method- ologies developed in this work will further enable the ex- ploration of similar methodologies for other deep networks which are more affordable and easily accessible to broader research community.
BFloat16 vs. Float8 Encoders. For our encoder, we can also use BF16 instead of the FP8 encoder from torchao (torchao maintainers & contributors, 2024). while the XMC classifier continues to use FP8. We compare the performance with different precision settings Table which shows similar precision but longer epoch time for FP8 due to the overhead in the FP8-BF16 mixed precision recipe.
Acknowledgements
We acknowledge the support of the Academy of Finland (Re- search Council of Finland) via grants 347707 and 348215 15 and the support of computational resources provided by the Aalto Science-IT project, and CSC IT Center for Science, Finland.
7. Conclusion
References
We present a low-precision training framework with BFloat16 6 and Float8 for XMC models, moving be-
Icee standard for floating point arithmetic. IEEE Std 754- 2019 (Revision of IEEE 754-2008), pp. 1-84, 2019.

---

ELMO: Efficiency via Low-precision and Peak Memory Optimization in Large Output Spaces
Achiam, J., Adler, S., Agarwal, S., Ahmad, L.,, Akkaya, I., Aleman, F. L.,, Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al. Gpt-4 technical report. arXiv preprint 1/Xi::2303.0877:: 2023.
Dettmers, T., Lewis, M., Shleifer, S., and Zettlemoyer, L. 8-bit optimizers via block-wise quantization. In International Conference on Learning Representations, 2022. URL https: http///ppeeeeeeeeeeeeeeeeeeeeeee net. htt.://wwwwee/oounn forum? i id=shpkpVxzo3h.
Babbar, R. and Schölkopf, B. Dismec: Distributed sparse machines for extreme multi-label classification. In Pro- ceedings of the tenth ACM international conference web search and data mining, pp. 721-729, 2017.
Devlin, J. Bert: Pre-training of deep bidirectional trans-
formers for language understanding. arXiv preprint
arXiv:1810.0480 2018.
Babbar, R. and Schölkopf, B. Data scarcity, robustness and extreme multi-label classification. Machine Learning, 108(8):1329-133) 2019.
Forsythe, G. E. Round-off errors in numerical integration on automatic muchinnry-ppeeliiing report. Bull. Amer. Math. Soc, 56:61-666 62, 1950.
Bhatia, K., Dahiya, K., Jain, H., Kar, P., Mittal, A., Prabhu, Y., and Varma, M. The extreme classification repository: Multi-label datasets and code, 2016. URL http:/ /mwnikmannnaaea.. org/ downloads / XC/ XCC/XMLeepositozy html.
Gupta, N., Devvrit, F., Rawat, A. A. Bhojanapalli, S., Jain, P., and Dhillon, S. Dual-encoders for extreme multi- label classification. In The Twelfth International Confer- ence Learning Representations, 2024. URL https: / /ppenrevieww net/ / forum? wir.-iiddeeeeeeeeeee
Blake, C., OIT, D., and Luschi, C. Unit scaling: Out-of-the- box low-pre cow-prccision training. In International Conference on Machine Learning, pp. 2548-2576. PMLR, 2023.
Hochreiter, S. Long short-term memory. Neural Computa- tion MIT-Press, 1997.
Hsu, P.-L.., Dai, Y., Kothapalli, V., Song, Q. Tang, S., Zhu, S., Shimizu, S., Sahni, S., Ning, H., and Chen, Y. Liger kernel: Efficient triton kernels for Ilm training. arXiv preprint arXiv:2410.10989, 2024.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877-1901, 2020.
Hu, B., Lu, Z., Li, H., and Chen, Q. Convolutional neu- ral network architectures for matching natural language sentences. Advances in neural information processing systems, 27, 2014.
Chung, J., Gulcehre, C., Cho, K., and Bengio, Y. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint XXXi::442235555 2014.
Jain, H., Prabhu, Y., and Varma, M. Extreme multi-label loss functions for recommendation, tagging, ranking & other missing label applications. In Proceedings of the 22nd ACM SIGKDD international conference on knowledge discovery and data mining, pp. 935-944, 2016.
Croci, M., Fasi, M., Higham, N. J., Mary, T., and Mikaitis, M. Stochastic rounding: implementation. error analy- sis and applications. Royal Society Open Science, 9()): 211631, 2022.
Jain, V., Prakash, J., Saini, D., Jiao, J., Ramjee, R., and Varma, M. Renee: End-to-end training of extreme classi- fication models. Proceedings of Machine Learning and Systems, 5, 2023.
Dahiya, K., Agarwal, A., Saini, D., Gururaj, K., Jiao, J., Singh, A., Agarwal, S., Kar, P., and Varma, M. Siame- sexml: Siamese networks meet extreme classifiers with 100m labels. In International conference on machine learning, pp. 2330-2340.. PMLR, 2021.
Jasinska- Jasinska-Kobus, K., Wydmuch, M., Thiruvenkatachari, D., and Dembezynski, K. Online probabilistic label trees. In International Conference on Artificial Intelligence and Statisties, pp. 1801-1809. PMLR, 2021.
Dahiya, K., Gupta, N., Saini, D., Soni, A., Wang, Y., Dave, K., Jiao, J., K, G., Dey, P., Singh, A., et al. Ngame: Negative mining-aware mini-batching for extreme clas- sification. In Proceedings of the Sixteenth ACM Interna- tional Conference on Web Search and Data Mining, pp. 258-266, 2023a.
Jiang, T., Wang, D.,, Sun, L., Yang, H., Zhao, Z., and Zhuang, F. Lightxml: Transformer with dynamic negative sam- pling for high-performance extreme multi-label text clas- sification. In Proceedings of the AAAI conference on artificial intelligence, volume 35, pp. 7987-7994, 2021.
Dahiya, K., Yadav, S., Sondhi, S., Saini, D., Mehta, S., Jiao, J., Agarwal, S., Kar, P., and Varma, M. Deep encoders with auxiliary parameters for extreme classification. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 358-367, 2023b.
Kahan, W. Pracniques: further remarks on reducing trun- cation errors. Commun. ACM, 8(1):40, January 1965. ISSN 0001-0782, doi: 10.1145/3637777363323.. URL https://don .rr..org//11 11.00..145/3200 httg//00.114//33777.66666
10

---

ELMO: Efficiency via Low-precision and Peak Memory Optimization in Large Output Spaces
Khandagale, S., Xiao, H., and Babbar, R. Bonsai: diverse and shallow trees for extreme multi-label classification. Machine Learning, 109((1):2099-2111 2020.
Mittal, A., Sachdeva, N., Agrawal, S., Agarwal, S., Kar, P., and Varma, M. Eclare: Extreme classification with label graph correlations. In Proceedings of the Web Conference 2021, pp. 3721-3732, 2021.
Kharbanda, S., Banerjee, A., Schultheis, E., and Babbar, R. Cascadexml: Rethinking transformers for end-to-end multi-resolutin training in extreme multi-label classifica- tion. Advances in neural information processing systems, 35:2074-2087, 2022.
Ozkara, K., Yu, T., and Park, Y. Stochastic rounding for Ilm training: Theory and practice. arXiv preprint arXiv:2502.20560 2025.
Paszke, A., Gross, S., Chintala, S., Chanan, G., Yang, E., DeVito, Z., Lin, Z., Desmaison, A., Antiga, L.,, and Lerer, A. Automatic differentiation in pytorch. 2017.
Kharbanda, S., Banerjee, A., Gupta, D., Palrecha, A., and Babbar, R. Inceptionxml: A lightweight framework with synchromized negative sampling for short text extreme classification. In Proceedings of the 46th International ACM SIGIR Conference on Research and Development in Information Retrieval, pp. 760-769, 2023.
Peng, H.,, Wu, K., Wei, Y., Zhao, G., Yang, Y., Liu, Z., Xiong, Y., Yang, Z., Ni, B., Hu, J., et al. Fp8-Im:
Training fp8 large language models. arXiv preprint arXiv:2310.183. 2023.
Kharbanda, S., Gupta, D., Schultheis, E., Banerjee, A., Hsieh, C.-J., and Babbar, R. Gandalf: Learning label- label correlations in extreme multi-label classification via label features. In Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 13360-137,,22 2024.
Prabhu, Y., Kag, A., Harsola. S., Agrawal, R., and Varma, M. Parabel: Partitioned label trees for extreme classifica- tion with application to dynamic search advertising. In Proceedings of the 2018 World Wide Web Conference, pp. 993-1002, 2018.
Kharbanda, S., Gupta, D., K, G., Malhotra, P., Singh, A., Hsieh, C.-J., and Babbar, R. Unidee: Unified dual encoder and classifier training for extreme multi-label classifica- tion. In Proceedings of the ACM on Web Conference 2025, pp. 4124-4133, 2025.
Rabe, M. N. and Staats, C. Self-attention does not need o(n2) memory. arXiv preprini arXiv:2112.05662, 2021.
Rae, J. W., Borgeaud, S., Cai, T., Millican, K., Hoffmann, J., Song, F., Aslanides, J., Henderson, S., Ring, R., Young, S., et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv: 12. 1446, 2021.
Kingma, D. and Ba, J. Adam: A method for stochastic optimization. arXiv preprint arXiv: 2014.
Loshchilov, and Hutter, F. Decoupled weight decay reg- ularization. In International Conference Learning Representations, 2017.
Schultheis, E. and Babbar, Towards memory-efficient training for extremely large output spaces-leaming with 670k labels a single commodity gpu. In Joint Euro- pean Conference on Machine Learning and Knowledge Discovery in Databases, pp. 689-704. Springer, 2023.
Lv, K., Yang, Y., Liu, T., Guo, Q., and Qiu, X. Full param- eter fine-tuning for large language models with limited resources. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers). pp. 8187-818 8198, 2024.
Sinha, A., Shen, Z., Song, Y.. Ma. H., Eide, D., Hsu, B.- j. P., and Wang, K. An overview of microsoft academic service (mas) and applications. In Proceedings of the 24th international conference on world wide web, pp. 243-246. ACM, 2015.
Mehta, S., Mohan, J., Natarajan, N., Ramjee, R., and Varma, M. Astra: Accurate and scalable anns-based training of extreme classifiers, 2024. URL https: //rxiv. org/ hogg/bb//4400 200bb//200.2066.
Smith, S., Patwary, M., Norick, B., LeGresley, P., Rajbhan- dari, S., Casper, J., Liu, Z., Prabhumoye, S., Zerveas, G., Korthikanti, V., et al. Using deepspeed and megatron to train megatron-turing nlg 530b, large-scale generative language model. arXiv preprint 2rXiv:2201.11999 2022.
Micikevicius. P., Narang, S., Alben, J., Diamos, G., Elsen, E., Garcia, D., Ginsburg, B., Houston, M., Kuchaiev, O., Venkatesh, G., et al. Mixed precision training. In International Conference on Learning Representations, 2018.
Sun, X., Choi, J., Chen, C.-Y., Wang, N., Venkataramani, S., Srinivasan, V. V., Cui, X., Zhang, W., and Gopalakr- ishnan, K. Hybrid 8-bit floating point (hfp8) training and inference for deep neural networks. Advances in neural information processing systems, 32, 2019.
Micikevicius, P.. Stosic, D., Burgess, N., Cornea, M., Dubey, P., Grisenthwaite, R., Ha, S., Heinecke, A., Judd, P., Kamalu, J., et al. Fp8 formats for deep learning. arXiv preprint arXiv:2009.054:2 2022.

---

ELMO: Efficiency via Low-precision and Peak Memory Optimization in Large Output Spaces
Tang, J., Zhang, D., and Yao, L. Social network extraction of academic researchers. In /CD///C/2/ 00. pp. 292-301, 2007.
You, R., Zhang, Z., Wang. Z., Dai, S., Mamitsuka, H., and Zhu, S. Attentionxml: Label tree-based attention- aware deep model for peg-operformnnc extreme multi- label text classification. Advances in neural information processing systems, 32, 2019.
Tang, J., Zhang, J., Yao, L.., Li, J., Zhang, L., and Su, Z.
Ametminer: Extraction and mining of academic social
networks. In KDD 08, pp. 990-998, 2008.
Yu, T., Guo, W., Li, J. C., Yuan, T., and De Sa, C. Mctensor: A high-precision deep learning library with multi-component arXi::2207.0886) 2022.
Tang, J., Yao, L., Zhang, D., and Zhang, J. A combination approach to web user profiling. ACM TKDD, 5(1):1-44, 2010.
arXiv preprint
Tang. J., Zhang, J., Jin, R., Yang, Z., Cai, K., Zhang, L., and Su, Z. Topic level expertise search over heterogeneous networks. Machine Learning Journal, 82(2):211-237, 2011.
Yu, T., Gupta, G., Gopalswamy, K., Mamidala, A. R., Zhou, H., Huynh, J., Park, Y., Diamant, R., Deoras, A., and Huan, L. Collage: Light-weight low-precision strategy for Ilm training. In Forty-firsi International Conference on Machine Learning, 2024.
Tang. J., Fong, A. C., Wang, B., and Zhang. J. A unified probabilistic framework for name disambiguation digi- tal library. IEEE Transactions on Knowledge and Data Engineering, 24(6)-975-987, 2012.
Zamirai, P., Zhang, J., Aberger, C. R., and De Sa, C. Revis- iting bffloat1 training. 2020.
Tillet, P., Kung, H.-T., and Cox, D. Triton: an intermediate language and compiler for tiled neural network computa- tions. In Proceedings of the 3rd ACM SIGPLAN Interna- tional Workshop on Machine Learning and Programming Languages, pp. 2019.
Zeng, A., Liu, X., Du, Z., Wang. Z., Lai, H., Ding, M.,
Yang, Z., Xu, Y., Zheng, W., Xia, X., al. Glm-130b: An open bilingual pre-trained model. arXiv preprint arXiv:2210.02 2022.
Zhang, J., Yang, J., and Yuen, H. Training with low- precision embedding tables. In Systems for Machine Learning Workshop NeurlPS, volume 2018, 2018.
torchao maintainers and contributors. torchao: Pytorch native quantization and sparsity for training and infer- ence, October 2024. URL https/ /github. com/ pytorch/torcha.
Zhang. J., Chang, W.-C., Yu, H.-F., and Dhillon, I. Fast multi-resolution transformer fine-tuning for extreme mubti-label text classification. Advances in Neural In- formation Processing Systems, 34:7267-7280, 2021.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., al. Llama: Open and efficient foundation lan- guage models. arXiv preprint arXiv:2302.13377 2023.
Zhang, S., Roller, S., Goyal, N., Artetxe, M., Chen, M., Chen, S., Dewan, C., Diab, M., Li, X., Lin, X. V., et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.0100 01068, 2022.
Ullah, N., Schultheis, E., Lasby, M., Ioannou, Y., and Bab- bar, R. Navigating extremes: Dynamic sparsity in large output spaces. In The Tiirty-eighth Annual Conference on Neural Information Processing Systems.
Zhao, J., Zhang, Z., Chen, B., Wang. Z., Anandkumar, A., and Tian, Y. Galore: Memory-effiiiens 11m training by gradient low-rank projection. In Forty-first International Conference on Machine Learning.
Vaswani, A. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Wan, L., Zeiler, M., Zhang, S., Le Cun, Y., and Fergus, R. Regularization of neural networks using dropconnect. In Dasgupta, S. and McAllester, D. (eds.), Proceedings of the 30th International Conference on Machine Learn- ing, volume 28 of Proceedings of Machine Learning Research, pp. 1058- 1066, Atlanta, Georgia, USA, 17- 19 Jun 2013. PMLR. URL http::s:///occeenneg/// mlr presss/ v28 w////. html.
Wydmuch, M., Jasinska, K., Kuznetsov, M., Busa-Fekete, R., and Dembezynski, K. A no-regret generalization of hierarchical softmax to extreme multi-label classification. Advances in neural information processing systems, 31, 2018.
12

---

ELMO: Efficiency via Low-precision and Peak Memory Optimization in Large Output Spaces
A. Baselines and Evaluation Metrics
where y is the true label vector, ý is the predicted score vector, and top; (t) identifies the indices with the top-k highest predicted scores.
We compare our method with deep XMC methods with mainly transformer encoder.
Propensity-Scored Precision k (PSP@k): Given the long-tailes tailed label distribution in many XMC datasets, PSP@k incorporates propensity score y to weight the precision contribution of each label, thereby emphasizing the tail labels' performance:
LightXML (Jiang et al., 2021): The method employs transformer encoder to concurrently train both the retriever and ranker, which incorporates dynamic nega- tive sampling to enhance the model's 's efficacy.
CascadeXML (Kharbanda et al., 2022): This method separates the feature learning of distinct tasks across various layers of the Probabilistic Label Tree (PLT) and aligns them with corresponding layers of the trans- former encoder.
PSPakly: 9) = Σ
VE
33)
Pt
where pt corresponds to the propensity score for the label yi (Jain al., 2016).
NGAME (Dahiya et al., 2023a): NGAME enhances transformer-based training for extreme classification by introducing a negative mining- aware mini-batching technique, which supports larger batch sizes and ac- celerates convergence by optimizing the handling of negative samples.
Algorithm 1 Float8 XMC classifier
class rP8CLassiiten Inn, Monule) iMC_update/aelf, # x [b, m] imputs for the clansifier # Rows and cols of the positive labels
labels):
X_gradient 2oxchh.ezos__Jkke1
rows, cols labels[:, labels(:,11 for in range (fell.numm fasmml.fnn__cnnkss)
Renee (Jain et al., 2023): The Renee model employs an integrated end-to-end training approach for extreme classification, using a novel loss shortcut for memory optimization and a hybrid data-model parallel architec- ture to enhance training efficiency and scalability.
self.R[i] [n, m), classitier weights one chunk
lagits = motmul_fptl f__
self.N[i], t. <.10l.tcttarre..laaee_eemmb (l03tt
} # The logits are in RF16
rows_i, cols_i fiiter_ehurk__La_____ whwukhhuuussaaeeeeeeeeeeeeeeeeee [labels/rowe cols)
lagits logits..iggoidd_
logitalcols, rowe Xgradient +- merge_k_matmul (legits, self.w(1])
DEXML (Gupta et al., 2024): The DEXML model aims eliminate the need for an explicit classifier, instead relying solely on dual-encoder-based training with all negative labels within each batch. While the motivation to remove the classifier (often the primary bottleneck) is sound, their approach ultimately incurs higher computational and memory costs compared to methods that retain a classifier.
(sse_uddate self.1s, logits, X, X. shape self.bs, self.bs, self.bs) return X_gracient
def puse_updazedoo (M,
logits, x, bk, bm, bn): # the pseudo code of the trison kernel grad = = tllzerss ( [bm, bn), dtype=FF32)
bk): in range From_HBM too (BF16)
x lead_block
1 = bac__blocc__oooo += block__sstttte (), x)
HBM fooos__froo_ieloliiieao (Logits)))) (BF16)
" foed_block_ftom Het_______0XM__ too (FP32)
To evaluate the performance of our Extreme Multi-labe Classification model, which incorporates low-precision training, we use set of metrics designed to provide com- prehensive analysis of both overall and label-specific model performance. The primary metries we employ is Precision at k (P@A), which assess the accuracy of the top-k predictions. Additionally, we incorporate Propensity- -corre Precision at (PSP@k) k) to gauge the uniformity of the model's effec- tiveness across the diverse range of labels typical XMC problems.
# 1regrad
stochast.ce_ree rounding_to..ooooooo (P8())
write_to__BMiC
Table 5: Epoch level training statistics on the RTX 4060 Ti. My denotes peak training memory.
Dataset
Epoch Time (mmsss)
(GB)
Fi-mmzzonTilles- 1.3M
S7:36
5.45
Amazon-3M
121:17
8.46
Precision at (P@ k): Precision is the fundamental metric for evaluating the top-k predictions in XMC appli- cations such e-commerce product recommendation and document tagging:
LF- Hi-Puper2Keewoods-8.kk
229:24
10.49
B. Background of Skipping Loss Computation
P@k(y, 9)
Σ
Va
2 For the XMC problems in this paper, we use the binary cross-entropy loss. Following Renee (Jain et al., 2023), we
EEtoP (9)
13

---

ELMO: Efficiency via Low-precision and Peak Memory Optimization in Large Output Spaces
Classifier input
Classifier weights
80%
70%
17.5%
Auranraaa of the clacitfer wellgnt
input
60%
50%
12.5%
40%
Percantage
30%
7.5%
20%
2.5%
10%
0.0%
-25
-2.0
-1.5
-1.0
-0.5
0.0
0.5
1.0
3.5
2.0
2.5
-3
-2
Exponent
Exponent
(a) The histogram of the classifier weights
(b) The histogram of the classifier input
Figure 5: Figures 5(a) and 5(b) show that most weights and classifier inputs fall within the exponent range of FP8 B4M3 (I-9, 8)), even without quantization. The weights and inputs are sampled from the training of LF-AmazonTittles 131K.
Table 6: Precision performance with Post-Hoc classifier refinement (on top of Float& checkpoint from ELMO and Kahan summation for head labels (20% head) on the LF- AmazonTitles: 1.3M. My denotes peak training memory.
C. Deployment Commodity Hardware
Although our main Float 8 experiments use H100 GPUs, we also demonstrate that the approach runs efficiently on commodity hardware, such as the GeForce RTX 4060 Ti, As shown in Table 5, the 4060 uses slightly more mem- ory, primarily due to the torch. supporting Float 8 on commodity hardware. As result, we use a BF1 6 encoder while still applying Float 8 to the XMC classifiers.
P@1
p@3
P@5
Mri(i
package not yet fully
Renee
56.04
49.91
45.32
19,9
BF16 (ELMO)
56.14
49.86
45.25
6.61
Float8 (ELMO)
54.97
48.41
43.82
4.31
Post-Hoc
55.4
48.87
44.34
4.31
D. Precision Recovery for Sensitive Applications
Head Kahan
55.6
49.38
44.88
4.65
For applications where recovering the last bit of accuracy is critical, two potential practical mitigation strategies that still operate within similar memory budgets are:
apply skipping loss computation, so the gradient for the classifier input is given by:
-1+e-v ( W
1. Post-hoc Classifier Refinement: A simple approach is fine-tune the classifier in higher precision top of an ELMO-trained (low-precision) model using frozen encoder features. This allows a partial recovery of the lost precision while staying within a constrained memory budget by loading only subsets of labels at time. This strategy introduces an additional training phase and hyperparameters be tuned for the second stage.
InputGrad
where y represents the logits, Y is the ground truth, and W denotes the classifier weights.
Similarly, the gradient for W is:
VW = 1
(x) X
2.
Kahan summation for head labels: To address accu- racy drops without additional training stages, we also outline another approach that leverages label statistics inherent in XMC tasks. By exploiting the long-tailed label distribution. one can apply Kahan summation with BF1 1 compensation only the top-p% most fre-
where X is the input embedding for the classifier.
We refer to ( - Y) as the "classifier logit gradient"
in our paper.
14

---

ELMO: Efficiency via Low-precision and Peak Memory Optimization in Large Output Spaces
Table 7: Comparison of Propensity based Precision @k for different XMC methods. The best results are denoted by bold and second best results are denoted by underline.
Method
PSP@1
PSP@3
PSP@5
PSP@1
PSP@3
PSP@5
PSP@1
PSP@3
PSP@5
Wiki-500K
AmazonTTiles-6706
Amazon-670K
ATTENTIONXML.
30.69
38.92
44
24.24
26.43
28.39
30.29
33.85
37.13
AR-TRANSFORMER
32.1
39.41
43.75
29.21
33.49
37.65
CASCADEXML
31.25
39.35
43.29
30.23
34.93
38.79
RENEE
32.9
42.31
46.78
27
31.1
34.89
31.45
36.16
40.15
ELMO (BF16)
33.32
42.56
47.03
28.62
32.13
35.27
30.84
35.69
40.06
ELMO (FP8)
32.40
41.68
46.17
28.24
31.88
35,26
30.57
35.33
39.67
Amazon-3M
LF- LF-WikiSeeAlso -300K
LF-AmazonTitue Titles-I.a. 1.3M
LIGHTXMI
17.85
21.26
24.16
XR- TRANSFORMER NGAME-2
25.18
30.13
33.79
20.06
24.85
27.79
33.83
37.79
41.03
29.18
33.01
35.36
RENBE ELMO (BF16)
14.39
17.47
19.80
32.02
37.07
40.9
28.54
33.38
36.14
15.65
19.05
21.6
31.65
37.08
41.04
30.38
34.59
37.09
ELMO (FP8)
16.06
19.48
21.98
31.87
36.98
40.90
26.72
31.58
34.46
Table 8: Precision and propensity scored precision comparison of ELMO with state of the art XMC baselines on LF- AmazonTitles- 131K dataset. The best results are denoted by bold and second best are by underline. Mrr denotes peak training memory.
Method
P@1
P@3
P@5
PSP@1
PSP@3
PSP@5
Mir (GiB)
Epoch Time (mm:s))
LIGHTXML
35.6
24.15
17.45
25.67
31.66
36.44
16.79
10:27
NGAME
44.69
29.89
21.21
38.81
44.4
49.43
11.03
5:15
DEXML
42.52
20.64
48.7
29.22
14:08
RENEE
46.05
30.81
22.04
39.08
45.12
50.48
5.53
0:33
ELMO (BF16)
45.6
30.6
21.9
38.84
45.02
50.39
3.41
0:31
ELMO (FP8)
45.45
30.53
21.87
38.75
44.98
50.41
2.75
0:22
quent labels. This approach selectively boosts pre- cision@k, with minimal memory overhead, approxi- mately 2xp% (where p% memory for p% label pa- rameters in Float8 8 more than the Float baseline. Importantly, this strategy preserves end-to-end training and avoids the complexity of multi-stage pipelines. For example, as shown in Table 6, on (mazon'Titles-1 1.3M with top 20% head labels, this method achieves com- petitive performance as Reene with a total classifier memory footprint of just 4.99 GB, still significantly below the baseline (6.61 GB).
when going down in bit-width. Similar to the precision @ performance, our approach shows competitive performance with existing state-of-the-art XMC baselines, showcasing low precision training can be robust tail labels. Perfor- mance of 1---AmazoTTitle--11111 is shown in Table 8.
F. Chunking Classifier Update: Latency vs Peak Memory
Table 10 shows the latency (epoch time) peak GPU mem- ory usage comparison for different chunk sizes. The training run is with ELMO BF16 6 data types for 3 million data size with batch size 128 on H100 GPU. We see chunking doesn't t affect the latency until some point. In fact, we see the la- tency improves as the chunking increases from 8.
E. Tail Label Performance
Table 7 compares the propensity scored precision@1 @k values, which is an indication of tail label performance. It would be interesting explore the effect of tail label performance
15

---

ELMO: Efficiency via Low-precision and Peak Memory Optimization in Large Output Spaces
Table 9: The hyper-paramettrs of the BFloat 6 and F Float8 models. Dropout is the embedding dropout and Encoder LR and XMC LR are the learning rates for the encoder and classifiers. WD denotes weight decay.
Dataset
Encoder
Batch Size
Seq. Length
Dropout
Encuder IR
XMC L.R.
Epochs
Warmup
WD (Encoder, XMC)
Dataset without Label Features
Wiki-500K (BF16)
BERT-Rasc Base
128
128
0.65
0.00002
0.05
35
1000
(0.01, 0.0001)
Wiki-500K (F
DERT-Baas
128
128
0.65
0.00002
0.15
70
1000
(0.01, 0.0001)
AmuzonTidles-6 (a=16)
BERT-Base
0,00005
100
0.0001)
AmazoTTids-6670K (226)
BERT-Base
256
32
0.75
0,00005
0.05
150
1000
(0.01. 0.0001)
Amazon-670K (BF26)
BERT-Basc
64
128
0.75
0.00002
0.06
100
500
00.01, 0.0001)
Amazon- 670K (FP8)
128
(0.01, 0.0001)
Amazon 3M (8F16)
BERT Base
128
128
0.5
0.00005
0.05
90
10000
(0.001. 0.001)
Amazon-3M (226)
BERT-Bas Base
128
128
0.65
0,00002
0.05
150
10000
(0.001, 0.001)
Dataset with Label Features
1K-Ammaze//itccc-131K L. 131 (7P8)
DiST
512
32
0.84
0.00001
100
15000
(0.
Dishi-BERY
512
32
0.85
0.00001
0.05
100
5000
(0.01, 0.0001)
L. WiliSeeeAloo (SP1 5)
Disil RoBERTa
128
256
0.75
0.00002
0.08
100
5000
(0.01, 0.0001)
LE WiddSecAboo3300K(2e55
RoBERTa
256
0.0001)
LF- AmnzzeeThiees-..3M
Disti- BERT
512
32
0.65
0,000005
0.05
100
5000
(0.0001, 0.0001)
1-0 (TP8)
Distil-BER
512
32
0.6
0.000005
0.2
150
15000
(0.01. 0.0001)
AeeeeeeerKKyy (BF1
BERT
0.0001)
LF Ke.PPer2KKyywuss 8.6M (FF8)
Distil BERT
128
128
0.70
0.00002
0.05
12
5000
(0.01, 0.0001)
Table 10: Peak memory vs Latency (in terms of epoch time) for different chunk size when chunking classifier update is used with BF on M dataset.
weights before performing matrix multiplication. This op- eration is handled within the Triton kernel, ensuring there is no additional memory overhead. We apply this for LE- Amazon Titles-1.33.
Chunk Size
Epoch Time (mm:ss)
Peak Mem (GiB)
13:22
14.74
2
12:20
14.40
4
12:12
12.22
11:09
11.13
16
11:23
10.59
32
12:39
10.32
64
14:19
10.20
128
19:44
10.20
G. Hyper-parameters and Implementation Details
We detail the els in the table 9.
of the BF16 6 and FP 8 mod-
H. Low-Memory Dropout for Classfier
To improve the classifier's robustness, we apply dropout (Wan et al., 2013) it. However, limitation of tradi- tional dropout is its requirement for copy of the classi- fier's weights, doubling memory consumption. To address this, we implement dropout directly within the matrix mul- tiplication process. After loading the classifier's weights from HBM into SRAM, we apply dropout to the loaded
16


Reviews:
[{"review_id": "jFAA4usj22", "reviewer": "Reviewer_QZwr", "summary": "This paper presents ELMO, an efficient training method designed for solving extreme multilabel classification (XMC) problems by leveraging low-precision computation. The key techniques that ELMO leverages are pure 16-bit training, classifier parameter chunking, and 8-bit training to reduce memory usage and improve training speed. Comprehensive experimental results have been demonstrated to justify the effectiveness of the proposed ELMO method, as well as showing that it outperforms the prior state-of-the-art Renee framework by a significant margin.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> It seems that using the ELMO method inevitably causes a precision drop compared to FP32 methods, which can potentially hurt the method's practicality for precision-sensitive applications.\n\nOn 5 out of 7 datasets in the paper, pure FP8/FP16 training with ELMO already achieves similar prediction performance as full/mixed precision methods such as Renee. The precision drop observed in the LF-Paper2Keywords-8.6M dataset was due to a preprocessing error on our part. Specifically, as a result of this oversight, we omitted the paper titles from the input to the ELMO model, whereas the baseline Float32 Renee utilized this information. After correcting this by concatenating the title and abstract as input to ELMO, the updated results for FP8 ELMO (in table below) are now very close to the Float32 Renee model. \n\n\n| LF-Paper2Keywords-8.6M | P@1 | P@3 | P@5 |\n| --- | --- | --- | --- |\n| Float32 | 43.6 | 32.13  | 26.02 |\n|  Previous FP8|  39.93| 29.88 |24.33  |\n| Updated FP8 | 43.4 |  31.59| 25.38 |\n\nOn the LF-AmazonTitles-1.3M dataset, the difference is within ~1.5% point in P@k, which can further be reduced to be within ~0.5% as mentioned in the answer to the next question.\n\n> For precision-sensitive applications, it is also not clear how to mitigate the precision/accuracy drop caused by the ELMO method.\n\nFor such applications where recovering the last bit of accuracy is critical, two potential practical mitigation strategies that still operate within similar memory budgets are:\n\n(i) Post-hoc Classifier Refinement: A simple approach is to fine-tune the classifier in higher precision on top of an ELMO-trained (low-precision) model using frozen encoder features. This allows a partial recovery of the lost precision while staying within a constrained memory budget by loading only subsets of labels at a time. This strategy introduces an additional training phase and hyperparameters to be tuned for the second stage.\n\n(ii) Kahan summation for head labels: To address accuracy drops without additional training stages, we also outline another approach that leverages label statistics inherent in XMC tasks. By exploiting the long-tailed label distribution, one can apply Kahan summation with BF16 compensation only to the top-P% most frequent labels. This approach selectively boosts precision@k, with minimal memory overhead—approximately 2×P% (where P% is memory for P%  label parameters in FP8) more than the FP8 baseline. Importantly, this strategy preserves end-to-end training and avoids the complexity of multi-stage pipelines. For example, on AmazonTitles-1.3M with top 20% head labels, this method achieves a competitive performance as Reene with a total classifier memory footprint of just 4.99 GB, still significantly below the BF16 baseline (6.61 GB).\n\nThe results of the above mentioned mitigation strategies are shown in the table below.\n\n\n| LF-AmazonTitles-1.3M | P@1 | P@3 | P@5 | Memory(GB) |\n| --- | --- | --- | --- | --- |\n| Renee | 56.04 | 49.91 |45.32  | 19.9 |\n| Original FP8 | 54.97 | 48.41 | 43.82 | 4.63 |\n|  Post Hoc| 55.4 | 48.87 | 44.34 | 4.63 |\n| Head Kahan |55.6  | 49.38 | 44.88 | 4.99 |\n\n> FP8 computation is only supported for GPUs using the Hopper architecture and newer, which limits the usability of the method for users with older generations of GPUs.\n\nFP8 computation is supported on Ada, Hopper, and Blackwell tensor cores. Notably, our model enables Float8 XMC classifiers to run on commodity GPUs, including the RTX 4000 series. For example, we successfully ran ELMO on an RTX 4060 with a memory footprint of just 10.49GB (on LF-Paper2Keywords-8.6M), achieving a training time of 3.5 hours per epoch. In contrast, Renee (58.4GB) and other state-of-the-art models cannot run on these affordable consumer GPUs.\n\n\nFor users with GPUs that do not natively support FP8 computation, we can store classifier weights in FP8 and upcasts them chunkwise to BF16 during forward and backward passes. This reduces memory usage while maintaining compatibility with older hardware. For example, on the LF-Paper2Keywords-8.6M dataset, ELMO with this strategy runs successfully on an A100 GPU with a memory footprint of 13.5 GB, compared to 20 GB for a pure BF16 and 9.6 GB for the FP8 version.\n\n> Is there a way to write a fused CUDA kernel for the proposed ELMO method to further improve computation and memory efficiency?\n\nA fused CUDA kernel would maintain similar memory efficiency to our fused Triton kernel. However, additional low-level warp- and thread-based optimizations might further improve computational performance.", "claims_and_evidence": "The major claims of speed-up and memory savings made by ELMO are well supported by the experimental evaluations.", "methods_and_evaluation_criteria": "Both the proposed methods and the evaluation criteria make sense for solving XMC with compute and memory efficiency.", "theoretical_claims": "There are no theoretical claims in this paper.", "experimental_designs_or_analyses": "The experimental designs and analyses make sense.", "relation_to_broader_scientific_literature": "The key contributions of this paper are related to low-precision efficient training more broadly.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths: \n- This paper is well-written and well-motivated in general.  \n- Exploring computation and memory-efficient methods for solving the XMC problem is a promising research direction.  \n- The experimental results of this paper seem to be convincing.  \n\nWeaknesses: \n- It seems that using the ELMO method inevitably causes a precision drop compared to FP32 methods, which can potentially hurt the method's practicality for precision-sensitive applications.  \n- It is also not clear how to mitigate the precision/accuracy drop caused by the ELMO method.  \n- FP8 computation is only supported for GPUs using the Hopper architecture and newer, which limits the usability of the method for users with older generations of GPUs.", "other_comments_or_suggestions": "Please see details provided in \"Other Strengths And Weaknesses\".", "questions_for_authors": "- Is there a way to write a fused CUDA kernel for the proposed ELMO method to further improve computation and memory efficiency?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "AfC66Cjbkv", "reviewer": "Reviewer_M72L", "summary": "This paper introduces ELMO, a low-precision training framework designed to optimize memory and computation for Extreme Multi-label Classification (XMC), where the classification layer dominates memory and compute costs. Key results include: 1) 6× memory reduction compared to Renee (previous SOTA) for 3M-label models, 2) comparable accuracy compared to FP16/FP32 training on XMC tasks, 3) peak memory reduced to 6.6GiB (FP8), compared to 39.7GiB in Renee.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Fused update has already been proposed in early literature\n\nThank you for pointing out these references. While related in motivation, our approach differs in key aspects:\n- LOMO: LOMO performs layer-wise fused updates by materializing gradients, applying optimizer steps, and releasing memory (https://pytorch.org/tutorials/intermediate/optimizer_step_in_backward_tutorial.html) before backward pass through next layer. While this reduces memory relative to full gradient accumulation, it still requires storing the entire gradient of the largest layer at each step. For example, in XMC with 3M labels, the classifier gradient alone requires ~8 GB, compared to an accumulated memory footprint of ~8.7 GB ( 0.7 GB related to non classifier). In contrast, our method applies fused updates using a kernel that avoids materializing the classifier gradient. This reduces peak memory usage to just ~0.7 GB, offering a significantly more efficient solution. We also differ in the choice of optimizers: while both methods use SGD, we adopt a hybrid approach—SGD for the classifier, and AdamW for the encoder achieving better performance-memory trade-offs.\n\n- GaLore: GaLore reduces memory by projecting gradients into a low-rank space, thereby compressing optimizer states. However, this is not directly applicable in our setting as the classifier —responsible for the majority of the memory footprint—is trained with plain SGD, without momentum or adaptive states. Moreover, SGD performs slightly better than AdamW for classifier training (as shown below in the table of the response to reviewer 4), an observation which corroborates the findings in Table 5 of Renee paper. \n\nWe appreciate the reference provided and will include it in the final version.\n\n> Lack of novelty: most of techniques (FP8, fused update, etcs) have already been proposed and widely used in recent works. ...\n\nWhile most of the mentioned techniques are tailored for optimizing large language models (LLMs), these do not directly apply to the setting of large output spaces. For a more detailed explanation based on concrete instances on how our approach differs, please refer to our response to the next question.\n\n> How does ELMO compare to other FP8 training frameworks? such as FP8-LM\n\nThe existing FP8 frameworks primarily retain higher precision for the classification layer. For instance, FP8-LM uses float16 data types. In contrast, in extreme classification, reducing classifier parameter size is crucial. We demonstrate that using FP8 data types for classifiers is viable, achieving competitive performance without the need for tensor scaling, thus eliminating additional overhead or expensive hyperparameter sweeps (in case of loss scaling).\n\nExisting FP8 training frameworks for LLM—including FP8-LM, Transformer Engine, COAT, and torch.ao—focus on reducing memory by quantizing optimizer states or employ mixed-precision strategies that reduce activation memory. However, in XMC, performance is best when using stateless optimizers like SGD. As a result, these methods are not directly applicable, and the dominant memory bottleneck becomes the classifier weights and their gradients, not the optimizer states. For example, on the LF-Paper2Keywords-8.6M dataset, the classifier in FP8-LM would consume approximately 37 GB of memory, whereas our approach requires only 6 GB. Additionally, mixed-precision training—commonly used in LLMs—is particularly harmful in XMC due to increased memory usage and should be avoided in practice.\n\nIn summary, our contributions differ from existing FP8 frameworks in several important ways:\n(1) We apply low precision directly to the classifier, unlike typical LLM-focused methods;\n(2) We show that FP8 classifiers can achieve strong performance in XMC without relying on scaling techniques;\n(3) We demonstrate that mixed-precision strategies are counterproductive in XMC due to memory overhead and are better avoided.\n\n> How much computational overhead does gradient fusion and chunking introduce?\n\nWe have included an ablation study of chunk size versus latency in Table 7. We observe that increasing the chunk size initially reduces epoch time up to an optimal point, after which epoch time begins to increase. We selected a chunk size below this optimal threshold.\n\nGradient fusion introduces no computational overhead and further improves efficiency by reducing gradient-related I/O operations. Without gradient fusion, the number of I/O operations for the classifier update with sgd is: num\\_labels×dim×6 (loading classifier weights for logits and input gradient computation, saving/loading classifier gradient, updating classifier) + num\\_labels×batch\\_size×7 (saving logits into HBM, loading it for input and classifier gradient computation, loading for sigmoid and its gradient computation). With gradient fusion, this reduces to: num\\_labels×dim×4 + num\\_labels×batch\\_size×7.", "claims_and_evidence": "1. ELMO significantly reduces memory usage in XMC models. Supported by empirical experiments, which shows that ELMO reduces peak memory by 75% compared to Renee.\n\n2. FP8 with stochastic rounding and Kahan summation stabilizes training for ELMO. Also, ELMO achieves training efficiency gains without sacrificing performance, as shown in Table 1.", "methods_and_evaluation_criteria": "ELMO introduces several techniques:\n\n1. FP8 Weights and BF16 Gradients\n2. Stochastic Rounding and Kahan Summation\n3. Chunking and Gradient Fusion\n4. Reordered Computation Flow", "theoretical_claims": "There is no theoretical claims included in the paoer", "experimental_designs_or_analyses": "Experiments in the paper covers:\n\n1. Memory efficiency compared to Renne across datasets\n2. Ablation on chunking and fused updates\n3. Trade off between FP8 and FP16", "relation_to_broader_scientific_literature": "related to low-precision training, memory-efficient training", "essential_references_not_discussed": "Fused update has already been proposed in early literature, such as below\n\n1. https://arxiv.org/abs/2306.09782\n2. https://arxiv.org/abs/2403.03507", "other_strengths_and_weaknesses": "Lack of novelty: most of techniques (FP8, fused update, etcs) have already been proposed and widely used in recent works. Applying them to a new learning task (XMC) might not have sufficient novelty.", "other_comments_or_suggestions": "None", "questions_for_authors": "1. How does ELMO compare to other FP8 training frameworks? such as FP8-LM\n2. How much computational overhead does gradient fusion and chunking introduce?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JA5hCFdaiN", "reviewer": "Reviewer_uk61", "summary": "This paper considers the problem of extreme classification where given an input, it is to be categorized into a few categories among a large set of possible categories. Full one-vs-all classifier training is an expensive approach to solve this problem, but has been shown to give best results and be possible to scale to 100M label space by Renee paper. This paper identifies some of the challenges in the Renee implementation, primarily, mixed precision training leading to a big memory footprint. It suggests fp8 based training and relevant optimizations around it to mitigate rounding errors. This yields a much more memory compact one-vs-all classifier training implementation that scales to tens of millions of labels in a reasonable amount of compute.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Scope is limited to extreme classifier based approach which is not very broadly applicable.\n\nBeyond the immediate applications in tagging, search and recommendation systems, large output spaces are becoming largely prevalent in modern LLMs [1,2,3] with increasing vocabulary sizes such as 256K tokens for Gemma2-2B. This presents another scenario where our Float8 training paradigm could be effectively applied.\n\n[1] Wijmans, et al. “Cut Your Losses in Large-Vocabulary Language Models\", arXiv:2411.09009\n\n[2] Tao, Chaofan, et al. \"Scaling laws with vocabulary: Larger models deserve larger vocabularies.\" arXiv preprint arXiv:2407.13623 (2024).\n\n[3] Yu, Da, et al. \"Scaling Embedding Layers in Language Models.\" arXiv preprint arXiv:2502.01637 (2025).\n\n> It would help to visualize Figure 4 in a systematic ablation form i.e. have multiple plots with optimizations being added sequentially (perhaps in sorted order of their impact)\n\nThanks for the suggestion! Section 4.4, along with Figure 3, already provides a detailed analysis of each optimization's contribution to memory usage, so adding an additional ablation study in the main paper might be somewhat repetitive, especially given the strict 8-page limit. However, we would include the systematic ablation visualization you've described as supplementary material in the appendix.\n\n> Table captions should be more verbose, as a reader it's easier if the details to parse a table are given in the caption\n\nThanks for pointing this out! We'll make the table captions more detailed to improve readability.\n\n> Do these training optimizations easily transfer to dual encoder training, specifically to something like the one-vs-all dual-encoder training in DEXML?\n\nOur method focuses on optimizing the classifier in XMC, while DEXML’s main memory bottleneck comes from computing label embeddings via the encoder. One can use an FP8 encoder (e.g., via torch.ao) and store label embeddings in FP8 to enable FP8 matmul for scoring. However, gradient fusion offers limited benefit in this setting, since label embeddings are produced by the encoder and not updated like classifier weights.\n\n> How do baseline (CascadeXML, DEXML, etc) methods perform on the new LF-Paper2Keywords-8.6M dataset?\n\nOn the LF-Paper2Keywords-8.6M dataset, we found that baseline methods like DEXML and CascadeXML did not scale to our multi-GPU training setups, including both 4×A100-80GB and 4×H100-80GB nodes. For example, in CascadeXML, the classifier parameters alone require approximately 25 GB of memory. When accounting for optimizer states (gradients and two momentum buffers), the total memory footprint becomes roughly 100 GB (25 + 25 + 2×25). Under Distributed Data Parallel (DDP), these parameters and states are replicated across GPUs, offering no effective memory savings.  Thank you for pointing this out and we will add these remarks in the appendix of the updated version of the paper.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "Yes. The experiments are done in the standard evaluation setting of extreme classification as well as compared against SoTA baseliens which make the evaluation protocol trustworthy.", "relation_to_broader_scientific_literature": "Extreme classification is a practical problem in recommendation scenarios. Although there has been a lot of work on generic deep learning training optimizations, but because of the special nature of an extreme classification solution (i.e. a wide final classifier) makes it open for specialized training optimizations, which is the core contribution of this paper.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "### Strengths\n- Paper is well presented, easy to read\n- The motivation is clear and straightforward \n- A new dataset is contributed which will help in evaluating larger scale XMC research\n\n### Weakness\n- Scope is limited to extreme classifier based approach which is not very broadly applicable", "other_comments_or_suggestions": "- It would help to visualize Figure 4 in a systematic ablation form i.e. have multiple plots with optimizations being added sequentially (perhaps in sorted order of their impact)\n- Table captions should be more verbose, as a reader it's easier if the details to parse a table are given in the caption", "questions_for_authors": "1. Do these training optimizations easily transfer to dual encoder training, specifically to something like the one-vs-all dual-encoder training in DEXML?\n2. How do baseline (CascadeXML, DEXML, etc) methods perform on the new LF-Paper2Keywords-8.6M dataset?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ubUp93yFn6", "reviewer": "Reviewer_VYeF", "summary": "This paper proposes a collection of quantization techniques to reduce memory usage of extreme classification (where the number of classes is large). Proposed techniques include stochastic rounding, Kahan summation, FP8 weight, FP16 gradient, chunking, etc, leading to several folds of memory usage reduction compared to previous methods.\n\n## Update after rebuttal\n\nThe authors rebuttal have addressed my concerns. I raised my score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Novelty: the paper is largely applying a collection of existing techniques to a specific problem. Key techniques such as stochastic rounding and chunking are quite standard. This paper looks more like an engineering optimization rather than research.\n\nPlease check our response to reviewer M72L on how our approach and setup differs from the similar innovations which are proposed largely in the context of LLMs.\n\n\n\n> While the paper proposes some specific techniques for extreme classification, there might be some more out-of-the-box general methods which can also achieve the goal. For example, 8-bit optimizer and per-block quantization [1].\n\nMost prior work on low-memory optimizers (e.g., Adam-8Bit, GaLore) focuses on reducing the memory footprint of optimizer states. However, when using pure SGD—as we do for the classifier—these methods offer no memory savings and can even increase memory usage due to unnecessary state tracking (see response to Reviewer  M72L on “How does ELMO compare to other FP8 training frameworks?”).\nBeyond memory, these stateful optimizers (e.g., Adam) also underperform in terms of accuracy when applied to the classifier. This is due to the extreme sparsity of tail-label updates in XMC, which makes momentum-based states noisy and less effective. In contrast, stateless SGD is both more memory-efficient and yields better performance in this setting.\nThe table below shows an ablation study comparing optimizers for the classifier on the LF-AmazonTitles-131K and AmazonTitles-670K datasets. The findings in Table 5 of Renee paper also corroborate this.\n\nLF-AmazonTitles-131K\n| Classifier optimizer | P@1 | P@3 | P@5 |\n| --- | --- | --- | --- |\n| SGD+SR (ELMO)| 45.5 | 30.5 | 21.9|\n| Adam+SR | 45.2 |  30.4| 21.8 |\n|  Adam-8bit| 42.9 |28.7  | 20.5 |\n\nAmazonTitles-670K\n| Classifier optimizer | P@1 | P@3 | P@5 |\n| --- | --- | --- | --- |\n| SGD+SR (ELMO)| 44.4 | 39.7 | 36.4 |\n| Adam+SR | 43.6 |  39.2| 36.0 |\n|  Adam-8bit| 42.8 |38.6  | 35.3 |", "claims_and_evidence": "yes", "methods_and_evaluation_criteria": "yes", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "yes", "relation_to_broader_scientific_literature": "See strengths and weaknesses.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strength:\n- The paper has a significant memory reduction compared to previous techniques.\n- The memory profiling is thorough.\n\nWeaknesses:\n- Novelty: the paper is largely applying a collection of existing techniques to a specific problem. Key techniques such as stochastic rounding and chunking are quite standard. This paper looks more like an engineering optimization rather than research.\n- While the paper proposes some specific techniques for extreme classification, there might be some more out-of-the-box general methods which can also achieve the goal. For example, 8-bit optimizer and per-block quantization [1].\n\n[1] Dettmers T, Lewis M, Shleifer S, et al. 8-bit Optimizers via Block-wise Quantization[C]//International Conference on Learning Representations.", "other_comments_or_suggestions": "I think the term \"8-bit training\" and \"16-bit training\" should be usage with caution. By default I suppose 16-bit training is PyTorch amp, and 8-bit training is TransformerEngine, which conducts matrix multiplications in low-precision.", "questions_for_authors": "NA", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The ELMO paper proposes a low-precision training framework for XMC (extreme multi-label classification) models, using BFLOAT16 and FLOAT8 data types. It demonstrates that XMC models can be effectively trained entirely in FLOAT8, without relying on single-precision master weights or tensor scaling, by leveraging Kahan summation and stochastic rounding. The method also incorporates memory optimizations like gradient fusion and chunking.

The paper is well-written and well-motivated. ELMO significantly reduces memory usage in XMC models, with empirical evidence showing a 75% reduction compared to Renee. FP8 training with stochastic rounding and Kahan summation stabilizes training for ELMO, achieving training efficiency gains without sacrificing performance.

On the negative side, the paper lacks novelty, as most of the techniques (FP8, fused update, etc.) have been proposed before. The paper is however solid and open-source release (especially of the Triton kernel) would significantly strengthen the contribution.

Overall, we recommend weak acceptance.
Decision: Accept (poster)