Transformer Dissection: A Unified Understanding of Transformer's Attention via the Lens of Kernel (1908.11775v4)

Abstract: Transformer is a powerful architecture that achieves superior performance on various sequence learning tasks, including neural machine translation, language understanding, and sequence prediction. At the core of the Transformer is the attention mechanism, which concurrently processes all inputs in the streams. In this paper, we present a new formulation of attention via the lens of the kernel. To be more precise, we realize that the attention can be seen as applying kernel smoother over the inputs with the kernel scores being the similarities between inputs. This new formulation gives us a better way to understand individual components of the Transformer's attention, such as the better way to integrate the positional embedding. Another important advantage of our kernel-based formulation is that it paves the way to a larger space of composing Transformer's attention. As an example, we propose a new variant of Transformer's attention which models the input as a product of symmetric kernels. This approach achieves competitive performance to the current state of the art model with less computation. In our experiments, we empirically study different kernel construction strategies on two widely used tasks: neural machine translation and sequence prediction.

• The paper reinterprets Transformer attention as a kernel smoothing process, unifying diverse attention formulations.
• It details novel kernel constructions that decouple non-temporal features from positional embeddings, enhancing efficiency.
• Empirical results on NMT and sequence prediction benchmarks demonstrate competitive performance with parameter efficiency.

Transformer Dissection: A Unified Understanding of Transformer's Attention via the Lens of Kernel

The paper authored by Yao-Hung Hubert Tsai et al., explores a new formulation of the attention mechanism in Transformers using the perspective of kernel methods. This approach provides an insightful understanding of various components of the Transformer's attention mechanism and proposes a novel variant that efficiently combines performance with computational efficiency.

Introduction and Motivation

Transformers have established themselves as the prominent architecture for sequence modeling tasks across multiple domains, such as neural machine translation, language understanding, and image generation. Central to their success is the attention mechanism, which captures the dependencies between elements in a sequence. Given the order-agnostic nature of attention, the authors propose a new formulation by interpreting it through the lens of kernel smoothers.

Formulation through Kernel Methods

The attention mechanism traditionally operates as a weighted summation over inputs, where the weights are determined by the pairwise similarities between input elements. Recognizing this, the paper establishes an equivalence between this process and kernel smoothing:

$$\text{Attention}(x_q \, ; \, M(x_q, S_{\mathbf{x}_k})) = \sum_{x_k \in M(x_q, S_{\mathbf{x}_k})} \frac{k(x_q, x_k)}{\sum_{x_k' \in M(x_q, S_{\mathbf{x}_k})} k(x_q, x_k')} \, v(x_k)$$

Here, the authors utilize a kernel function $k(\cdot, \cdot)$ to measure similarity and a function $v(\cdot)$ to produce the values. This kernel-based perspective enables a broader exploration and categorization of different attention mechanisms.

Kernel Construction and Positional Embeddings

The core insight involves constructing the kernel on a joint feature space consisting of both non-temporal features and positional embeddings. The paper discusses several methods for incorporating positional embeddings, such as:

1. Direct Sum in Feature Space: $$k((f_q, t_q), (f_k, t_k)) = k_{\text{exp}}(f_q + t_q, f_k + t_k)$$
2. Look-up Table Integration: $$k((f_q, t_q), (f_k, t_k)) = L_{t_q-t_k} \cdot k_{\text{exp}}(f_q, f_k)$$
3. Product Kernel with Asymmetric Kernel: $$k((f_q, t_q), (f_k, t_k)) = k_{\text{exp}}(f_q, f_k) \cdot k_{f_q}(t_q, t_k)$$

Through experimentation, they propose a novel form:

$$k((f_q, t_q), (f_k, t_k)) = k_F(f_q, f_k) \cdot k_T(t_q, t_k)$$

where both kernel functions $k_F$ and $k_T$ are exponential kernels, which decomposes the similarity measurement into independent components for non-temporal and positional features.

Value Function and Attention Mask Discussions

The authors discuss the role of the value function $v(\cdot)$ and argue that positional information need not be included within it. Their experimental results substantiate that avoiding positional embeddings in the value function achieves better performance.

Moreover, they address the misconception of the attention mechanism's order-invariance due to the presence of masks in decoder self-attention. Their argument highlights that decoder self-attention inherently contains order information, thus challenging previous notions about the necessity of positional embeddings in such contexts.

Experimental Results

The empirical studies are conducted on two benchmarks: IWSLT'14 German-English for neural machine translation (NMT) and WikiText-103 for sequence prediction (SP). Key findings include:

• The proposed product kernel form shows competitive performance, with a BLEU score of 34.71 for NMT and a perplexity of 24.28 for SP.
• Symmetric kernels are as effective as asymmetric ones while being more parameter-efficient.
• Exponential and RBF kernels outperform polynomial kernels, with RBF achieving the highest BLEU score in NMT.

Implications and Future Work

The formulation enhances the understanding of attention, providing a unified view that can harmonize various existing approaches. This approach could potentially guide the design of more sophisticated attention mechanisms tailored to specific applications. Future research may explore the benefits of this kernel-based formulation in other domains like multimodal learning or discover new ways of integrating long-range dependencies in sequence models.

Conclusion

By dissecting and reinterpreting Transformer's attention through kernel methods, this paper offers a robust framework for understanding and innovating in attention mechanisms. This formulation has both theoretical and practical implications, broadening the scope of potential improvements in Transformer architectures.

Acknowledgements

The paper reflects the collaborative efforts at Carnegie Mellon University, Kyoto University, and RIKEN AIP, with support from grants by DARPA, ONR, AFRL CogDeCON, NSF, National Institutes of Health, JST PRESTO program, and NVIDIA’s GPU support. The insights derived from Zhilin Yang's discussions were particularly valuable in refining the positional encoding strategies.