Improving Multi-Head Attention with Capsule Networks

TL;DR

This paper introduces a novel approach using capsule networks to enhance multi-head attention in neural machine translation, effectively clustering similar information and preserving unique features to improve translation quality.

Contribution

The paper proposes integrating capsule networks with multi-head attention to better handle semantic overlaps and improve translation performance.

Findings

Consistent improvements on Chinese-English translation

Enhanced clustering of similar attention information

Preserved unique features in attention heads

Abstract

Multi-head attention advances neural machine translation by working out multiple versions of attention in different subspaces, but the neglect of semantic overlapping between subspaces increases the difficulty of translation and consequently hinders the further improvement of translation performance. In this paper, we employ capsule networks to comb the information from the multiple heads of the attention so that similar information can be clustered and unique information can be reserved. To this end, we adopt two routing mechanisms of Dynamic Routing and EM Routing, to fulfill the clustering and separating. We conducted experiments on Chinese-to-English and English-to-German translation tasks and got consistent improvements over the strong Transformer baseline.