RTop-K：用于 GPU 上神经网络加速的超快速逐行 Top-K 选择算法
分布式、并行与集群计算
2025-04-03 v4
摘要
Top-k 选择算法是广泛应用中的基础算法，包括高性能计算、信息检索、大数据处理和神经网络模型训练。在本文中，我们提出了 RTop-K，一种专为 GPU 设计的高效并行逐行 Top-K 选择算法。RTop-K 利用基于二分搜索的方法来优化逐行 Top-K 选择，提供了一个可扩展且加速的解决方案。我们对算法中的提前停止进行了详细分析，表明它有效地保持了神经网络模型的测试准确率，同时显著提高了性能。我们的 RTop-K GPU 实现展示了优于最先进的逐行 Top-K GPU 实现的卓越性能，在启用提前停止时平均加速高达 11.49 倍，不启用提前停止时平均加速高达 7.29 倍。此外，RTop-K 加速了 MaxK-GNNs 的整体训练工作流程，在不同模型和数据集上实现了从 11.97% 到 33.29% 不等的加速。
引用
@article{arxiv.2409.00822,
title = {RTop-K: Ultra-Fast Row-Wise Top-K Selection for Neural Network Acceleration on GPUs},
author = {Xi Xie and Yuebo Luo and Hongwu Peng and Caiwen Ding},
journal= {arXiv preprint arXiv:2409.00822},
year = {2025}
}
备注
ICLR 2025 Conference