动态分配注意力头,让模型更省力地处理不同复杂度任务。
Adaptive Head Budgeting for Efficient Multi-Head Attention

- 根据输入自动选择最相关的注意力头,避免全量计算。
- 在文本分类任务中减少计算量,同时保持甚至超越原模型性能。
- 适合需要高效推理的部署场景,尤其对资源受限设备友好。
多头注意力使Transformer能够捕捉多样化的表征,但所有注意力头通常对每个输入都全量激活,无论任务复杂度如何。对于文本分类等粗粒度任务,相关信息往往是全局性的,这种固定分配会引入不必要的计算开销。本文提出BudgetFormer,一种Transformer架构,可基于每个输入动态分配注意力头。该模型学习头预算和头相关性分布,以选择最具信息量的头部。为支持有效的头部选择,我们设计了一种平衡探索与利用的训练策略。在文本分类任务上的实验表明,BudgetFormer在减少浮点运算次数(FLOPs)和内存使用的同时,性能与标准多头注意力相当或更优。结果表明,自适应头分配是提升Transformer效率与性能的有效方法。
原文摘要 · Abstract (English)
Multi-head attention enables Transformers to capture diverse representations, but all attention heads are typically activated for every input, regardless of task complexity. For coarse-grained tasks such as text classification, where relevant information is often global, this fixed allocation can introduce unnecessary computation. We propose BudgetFormer, a Transformer architecture that dynamically allocates attention heads on a per-input basis. The model learns both a head budget and a relevance distribution to select the most informative heads. To support effective head selection, we introduce a training strategy that balances exploration and exploitation. Experiments on text classification tasks show that BudgetFormer reduces FLOPs and memory usage while matching or surpassing the performance of standard multi-head attention. These results highlight adaptive head allocation as an effective approach to improving Transformer efficiency and performance.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。