Differentiable Subset Pruning of Transformer Heads

TL;DR

This paper introduces a differentiable subset pruning method for Transformer multi-head attention, which learns head importance and enforces a fixed number of active heads, resulting in smaller, faster models with maintained or improved performance.

Contribution

It proposes a novel differentiable head pruning technique that allows precise control over the number of unpruned heads in Transformer models.

Findings

Performs comparably or better than previous pruning methods.

Offers precise control over sparsity levels.

Reduces model size and increases speed without sacrificing accuracy.

Abstract

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.