Multi-head attention, a collection of several attention mechanisms that independently attend to different parts of the input, is the key ingredient in the Transformer. Recent work has shown, however, that a large proportion of the heads in a Transformer's multi-head attention mechanism can be safely pruned away without significantly harming the performance of the model; such pruning leads to models that are noticeably smaller and faster in practice. Our work introduces a new head pruning technique that we term differentiable subset pruning. Intuitively, our method learns per-head importance variables and then enforces a user-specified hard constraint on the number of unpruned heads. The importance variables are learned via stochastic gradient descent. We conduct experiments on natural language inference and machine translation; we show that differentiable subset pruning performs comparably or better than previous works while offering precise control of the sparsity level.

翻译：多头注意力机制是Transformer的核心组件，它包含多个独立的注意力机制以并行关注输入的不同部分。然而近期研究表明，Transformer多头注意力机制中大部分注意力头可以在不显著影响模型性能的情况下被安全剪枝，这种剪枝技术能在实际应用中显著缩小模型规模并提升计算效率。本文提出一种名为可微子集剪枝的新型注意力头剪枝技术。其核心思路是通过学习每个注意力头的重要性变量，进而对未剪枝头数量施加用户指定的硬约束。这些重要性变量通过随机梯度下降进行学习。我们在自然语言推理和机器翻译任务上开展实验，结果表明可微子集剪枝在实现精确稀疏度控制的同时，其性能与既往方法相当或更优。