对比八种注意力机制，找出训练最省电的方案
Attentions Under the Microscope: A Comparative Study of Resource Utilization for Variants of Self-Attention
- 实测八种注意力在GPT-2上的资源消耗，涵盖时耗、显存、功耗
- 优化内核的Flash Attention等三种机制能效最优，显著降低能耗
- 提醒：低功耗不等于省电，训练时间同样关键，适合模型优化者参考
随着大语言模型和视觉语言模型规模增长，注意力机制因高内存与计算复杂度成为主要瓶颈。尽管已有多种高效注意力变体提出，但其训练时的真实能耗与硬件资源需求缺乏严谨评估。本文在GPT-2架构上基准测试了八种注意力机制，测量训练时间、GPU显存占用、浮点运算量（FLOPS）、CPU使用率及功耗等指标。结果表明，采用优化内核实现的Flash Attention、局部敏感哈希（LSH）注意力与多头潜在注意力（MLA）在能效方面表现最佳。研究进一步揭示，仅降低GPU功耗并不意味着总能耗减少，训练时长同样关键。本工作强调了注意力设计中能源感知基准的重要性，并为选择资源高效机制提供了实用依据。所有代码已开源至GitHub。
原文摘要 · Abstract (English)
As large language models (LLMs) and visual language models (VLMs) grow in scale and application, attention mechanisms have become a central computational bottleneck due to their high memory and time complexity. While many efficient attention variants have been proposed, there remains a lack of rigorous evaluation on their actual energy usage and hardware resource demands during training. In this work, we benchmark eight attention mechanisms in training GPT-2 architecture, measuring key metrics including training time, GPU memory usage, FLOPS, CPU usage, and power consumption. Our results reveal that attention mechanisms with optimized kernel implementations, including Flash Attention, Locality-Sensitive Hashing (LSH) Attention, and Multi-Head Latent Attention (MLA), achieve the best energy efficiency. We further show that lower GPU power alone does not guarantee reduced energy use, as training time plays an equally important role. Our study highlights the importance of energy-aware benchmarking in attention design and provides a practical insight for selecting resource-efficient mechanisms. All our codes are available at GitHub.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。