ELMO：大规模输出空间中基于低精度与峰值内存优化的效率方法
机器学习
2025-10-14 v1 计算与语言
信息检索
摘要
大规模输出空间（亦称极端多标签分类，XMC），在大规模标签和产品到产品推荐等场景中出现，其标签数量范围为数十万到数百万。这意味着线性分类头——通常仅占整体模型的一小部分——成为计算和内存需求的主要驱动力。当前领先的XMC方法主要依赖FP16-FP32混合精度训练，我们展示其在内存使用和计算开销方面可能不稳定且效率低下。同时，现有低精度方法通常为分类层保留更高精度。本文提出ELMO，一种用于XMC模型的纯低精度训练框架，采用BFloat16和Float8数据类型。通过利用Kahan求和和随机舍入，我们展示在不依赖单精度主权重或张量缩放的情况下，XMC模型可有效地完全以Float8进行训练。低精度训练结合我们提出的内存优化措施——梯度融合与分块 ——显著降低GPU内存使用。例如，我们仅使用6.6 GiB的GPU内存即可训练一个300万标签的XMC模型，与优化后的SOTA方法Renee所需的39.7 GiB内存相比， accuracy未受影响。
引用
@article{arxiv.2510.11168,
title = {ELMO: Efficiency via Low-precision and Peak Memory Optimization in Large Output Spaces},
author = {Jinbin Zhang and Nasib Ullah and Erik Schultheis and Rohit Babbar},
journal= {arXiv preprint arXiv:2510.11168},
year = {2025}
}
备注
Accepted to ICML 2025