Research map / Architectures and efficiency / RNNs and normalization

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

ICML 2025 · Architectures and efficiency › RNNs and normalization · approach: Dataset distillation and data-efficient training

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes ELMO, a pure BF16/FP8 low-precision training framework for extreme multilabel classification that uses Kahan summation, stochastic rounding, gradient fusion and chunking to cut peak GPU memory several-fold versus Renee, and releases a new 8.6M-label dataset.

Problem

Training extreme multilabel classification (XMC) models with millions of labels is memory- and compute-bound by the classifier layer, and existing SOTA end-to-end training (Renee) uses FP16-FP32 mixed precision that is unstable and memory-hungry.

Approach

Train the whole XMC model, including the huge classifier, in pure BF16 or FP8 without FP32 master weights or tensor scaling, using stochastic rounding (classifier) and Kahan summation (encoder), dropping momentum, fusing classifier gradient computation with the SGD update in a Triton kernel, and chunking the label dimension to reduce transient memory.

Key results

• On Amazon-3M, ELMO (FP8) reaches 6.6 GiB peak training memory versus 39.7 GiB for Renee, with P@1 52.73 vs 52.6 for Renee.
• On LF-Paper2Keywords-8.6M, ELMO (BF16) achieves P@1 45.4 versus 43.60 for Float32 and 17.65 for Renee, using 18.8 GiB versus 105.64 GiB.
• On Wiki-500K, ELMO (FP8) uses 5.01 GiB versus 12.69 GiB for Renee with P@1 of 78.39 vs 78.69.

Measured effects reported in the paper

Surprising findings

• Renee's FP16-FP32 mixed precision training is unstable for large label spaces due to gradient overflow in the classifier input, and Renee underperforms badly on LF-Paper2Keywords-8.6M.
• Pure FP8 classifier training works without tensor scaling, since E4M3 native range suffices for classifier weights, and momentum can be removed for the classifier with plain large-learning-rate SGD.

Limitations

• FP8 encoder via torchao is not fully supported on commodity hardware such as RTX 4060 Ti, requiring a BF16 encoder there.
• Some FP8 results trail BF16 or Renee slightly on certain datasets (e.g., LF-AmazonTitles-1.3M P@1 54.97 vs 56.04 for Renee), and the paper notes a recipe gap for bit-widths like FP6/FP4 without micro-tensor scaling.
• Gradients of the transformer backbone remain in higher precision, so the pipeline is only nearly pure FP8.

Experimental setup

Datasets: Wiki-500K AmazonTitles-670K Amazon-670K Amazon-3M LF-AmazonTitles-131K LF-WikiSeeAlso-320K LF-AmazonTitles-1.3M DBLP-Citation-network V14

Benchmarks: Wiki-500K AmazonTitles-670K Amazon-670K Amazon-3M LF-AmazonTitles-131K LF-WikiSeeAlso-320K LF-AmazonTitles-1.3M LF-Paper2Keywords-8.6M

Metrics: P@1 P@3 P@5 PSP@k peak training memory (GiB) epoch time

Baselines: Renee LightXML CascadeXML NGAME DEXML Float32

Builds on: Renee mixed precision training Kahan summation stochastic rounding torchao

Closest papers by meaning

• HASTE: Hardware-Aware Dynamic Sparse Training for Large Output Spaces (ICML 2026 · similarity 0.88)
• A no-regret generalization of hierarchical softmax to extreme multi-label classification (NeurIPS 2018 · similarity 0.88)
• Large Language Models Meet Extreme Multi-label Classification: Scaling and Multi-modal Framework (AAAI 2026 · similarity 0.87)
• ELIAS: End-to-End Learning to Index and Search in Large Output Spaces (NeurIPS 2022 · similarity 0.86)
• Extreme Multi-label Classification from Aggregated Labels (ICML 2020 · similarity 0.86)
• Taming the Sigmoid Bottleneck: Provably Argmaxable Sparse Multi-Label Classification (AAAI 2024 · similarity 0.85)