面向极大型输出空间的内存高效训练——在单块商用 GPU 上学习 50 万标签
机器学习
2023-11-08 v1 人工智能
分布式、并行与集群计算
摘要
在具有大型输出空间（多达数百万标签）的分类问题中，最后一层可能需要极大的内存。使用稀疏连接将大幅降低内存需求，但正如我们下文所示，这会导致模型预测性能显著下降。幸运的是，我们发现可通过引入一个中间尺寸的倒数第二层来缓解此问题。我们进一步证明，可以将稀疏层的连接限制为均匀的，即每个输出神经元具有完全相同的传入连接数。这使得在 GPU 硬件上能够高效实现稀疏矩阵乘法与连接重分配。通过定制的 CUDA 实现，我们展示了所提方法可扩展到在仅 4GB 内存的单块商用 GPU 上处理含 670,000 个标签的数据集。
引用
@article{arxiv.2306.03725,
title = {Towards Memory-Efficient Training for Extremely Large Output Spaces -- Learning with 500k Labels on a Single Commodity GPU},
author = {Erik Schultheis and Rohit Babbar},
journal= {arXiv preprint arXiv:2306.03725},
year = {2023}
}