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

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.

Please 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.

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].

Most 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?”). Beyond 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. The 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.

LF-AmazonTitles-131K

AmazonTitles-670K

Scope is limited to extreme classifier based approach which is not very broadly applicable.

Beyond 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.

[1] Wijmans, et al. “Cut Your Losses in Large-Vocabulary Language Models", arXiv:2411.09009

[2] Tao, Chaofan, et al. "Scaling laws with vocabulary: Larger models deserve larger vocabularies." arXiv preprint arXiv:2407.13623 (2024).

[3] Yu, Da, et al. "Scaling Embedding Layers in Language Models." arXiv preprint arXiv:2502.01637 (2025).

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)

Thanks 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.

Table captions should be more verbose, as a reader it's easier if the details to parse a table are given in the caption

Thanks for pointing this out! We'll make the table captions more detailed to improve readability.

Do these training optimizations easily transfer to dual encoder training, specifically to something like the one-vs-all dual-encoder training in DEXML?

Our 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.

How do baseline (CascadeXML, DEXML, etc) methods perform on the new LF-Paper2Keywords-8.6M dataset?

On 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.

Fused update has already been proposed in early literature

Thank you for pointing out these references. While related in motivation, our approach differs in key aspects:

• 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.

• 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.

We appreciate the reference provided and will include it in the final version.

Lack of novelty: most of techniques (FP8, fused update, etcs) have already been proposed and widely used in recent works. ...

While 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.

How does ELMO compare to other FP8 training frameworks? such as FP8-LM

The 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).

Existing 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.

In summary, our contributions differ from existing FP8 frameworks in several important ways: (1) We apply low precision directly to the classifier, unlike typical LLM-focused methods; (2) We show that FP8 classifiers can achieve strong performance in XMC without relying on scaling techniques; (3) We demonstrate that mixed-precision strategies are counterproductive in XMC due to memory overhead and are better avoided.

How much computational overhead does gradient fusion and chunking introduce?

We 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.

Gradient 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.

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.

On 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.

On 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.

For precision-sensitive applications, it is also not clear how to mitigate the precision/accuracy drop caused by the ELMO method.

For such applications where recovering the last bit of accuracy is critical, two potential practical mitigation strategies that still operate within similar memory budgets are:

(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.

(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).

The results of the above mentioned mitigation strategies are shown in the table below.

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.

FP8 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.

For 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.

Is there a way to write a fused CUDA kernel for the proposed ELMO method to further improve computation and memory efficiency?

A 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.