Head-wise token pruning method and electronic device

Abstract

The present disclosure provides a token pruning method and an electronic device. The method includes: computing a plurality of query matrices and a plurality of key matrices corresponding to a plurality of heads; for each head, computing an attention matrix based on the corresponding query matrix and key matrix, wherein the attention matrix corresponds to multiple tokens and includes multiple attention values; for each head, computing an importance of each token based on the corresponding attention values, and determining at least one pruned token based on the importances, wherein the pruned tokens of a first head are different from those of a second head; and omitting the respective pruned tokens of each head in a subsequent computation.

Images