Contributions of Transformer Attention Heads in Multi- and Cross-lingual Tasks

TL;DR

This study investigates the importance of attention heads in Transformer models for cross-lingual and multi-lingual tasks, finding that pruning heads can improve performance and identify key heads using gradient-based methods.

Contribution

It demonstrates that pruning attention heads in multi-lingual Transformers can enhance performance and introduces a gradient-based ranking method to identify heads to prune.

Findings

Pruning attention heads generally improves cross-lingual and multi-lingual task performance.

Gradient-based ranking effectively identifies heads for pruning.

Results are consistent across mBERT and XLM-R models on multiple languages.

Abstract

This paper studies the relative importance of attention heads in Transformer-based models to aid their interpretability in cross-lingual and multi-lingual tasks. Prior research has found that only a few attention heads are important in each mono-lingual Natural Language Processing (NLP) task and pruning the remaining heads leads to comparable or improved performance of the model. However, the impact of pruning attention heads is not yet clear in cross-lingual and multi-lingual tasks. Through extensive experiments, we show that (1) pruning a number of attention heads in a multi-lingual Transformer-based model has, in general, positive effects on its performance in cross-lingual and multi-lingual tasks and (2) the attention heads to be pruned can be ranked using gradients and identified with a few trial experiments. Our experiments focus on sequence labeling tasks, with potential applicability on other cross-lingual and multi-lingual tasks. For comprehensiveness, we examine two pre-trained multi-lingual models, namely multi-lingual BERT (mBERT) and XLM-R, on three tasks across 9 languages each. We also discuss the validity of our findings and their extensibility to truly resource-scarce languages and other task settings.