[Paper Review] Capsule-Transformer for Neural Machine Translation

Sufeng Duan, Juncheng Cao+1 more|arXiv (Cornell University)|Apr 30, 2020

Topic Modeling29 references4 citations

TL;DR

This paper proposes capsule-Transformer, a novel architecture that generalizes the multi-head self-attention mechanism in Transformers by replacing linear transformations with a capsule routing algorithm, enabling richer contextual representation through hierarchical attention aggregation. Experiments on neural machine translation benchmarks show significant performance gains over strong Transformer baselines without adding substantial parameters.

ABSTRACT

Transformer hugely benefits from its key design of the multi-head self-attention network (SAN), which extracts information from various perspectives through transforming the given input into different subspaces. However, its simple linear transformation aggregation strategy may still potentially fail to fully capture deeper contextualized information. In this paper, we thus propose the capsule-Transformer, which extends the linear transformation into a more general capsule routing algorithm by taking SAN as a special case of capsule network. So that the resulted capsule-Transformer is capable of obtaining a better attention distribution representation of the input sequence via information aggregation among different heads and words. Specifically, we see groups of attention weights in SAN as low layer capsules. By applying the iterative capsule routing algorithm they can be further aggregated into high layer capsules which contain deeper contextualized information. Experimental results on the widely-used machine translation datasets show our proposed capsule-Transformer outperforms strong Transformer baseline significantly.

Motivation & Objective

• To address the limitation of standard multi-head self-attention in capturing deeper contextualized representations due to its reliance on simple linear transformations.
• To integrate capsule network principles into the self-attention mechanism to enable dynamic, hierarchical information aggregation across attention heads and tokens.
• To enhance the representational capacity of Transformers without significantly increasing model complexity or parameters.
• To explore whether capsule routing can improve attention distribution learning in sequence modeling tasks like neural machine translation.

Proposed method

• Groups of attention weights from different heads are treated as low-level capsules, each encoding preliminary linguistic features.
• The capsule routing algorithm is applied iteratively to aggregate these low-level capsules into high-level capsules that capture more contextualized, compositional representations.
• The routing mechanism dynamically determines the importance of each attention head in contributing to the final representation, replacing fixed linear projections in standard self-attention.
• The proposed capsule-attention mechanism is integrated into the Transformer encoder and decoder, replacing the standard multi-head attention module.
• The routing process uses dynamic routing with attention-based routing weights, allowing the model to learn which attention patterns are most relevant for higher-level representations.
• The architecture preserves the original Transformer’s efficiency and is trained end-to-end with minimal additional parameters.

Experimental results

Research questions

• RQ1 Can capsule routing improve the representational quality of self-attention by enabling hierarchical aggregation of attention heads?
• RQ2 Does replacing linear transformations in self-attention with capsule routing lead to better contextualized representations in sequence modeling?
• RQ3 Can the capsule-Transformer achieve superior performance on neural machine translation tasks compared to standard Transformers?
• RQ4 Is the performance gain due to improved attention distribution learning rather than increased model capacity?

Key findings

• The capsule-Transformer outperforms the strong baseline Transformer on standard neural machine translation benchmarks, demonstrating consistent and significant improvements.
• The model achieves better performance with minimal parameter increase, indicating that the improvement stems from architectural enhancement rather than capacity scaling.
• The capsule routing mechanism enables more coherent and contextually aware attention distributions by aggregating information across heads and tokens.
• The ablation study confirms that the routing mechanism contributes meaningfully to performance, and removing it degrades results significantly.
• The proposed method generalizes well, suggesting potential for broader application in other NLP tasks using Transformer-based models.