The black box of deep learning is far darker than we imagine.
Introduction
It is said that the physicist Richard Feynman once remarked: "If you think you understand quantum mechanics, you don’t understand quantum mechanics." I increasingly feel that "quantum mechanics" in this sentence could be replaced with "deep learning." Although deep learning has proven its effectiveness in more and more fields, our interpretability of it remains quite weak. Of course, in recent years, many efforts have been dedicated to opening this black box of deep learning. Unfortunately, most of these works are "post-hoc"—they propose explanations that are barely convincing based on existing experimental results, failing to construct and understand model principles from the top down, let alone make forward-looking predictions.
This article focuses on the self-attention mechanism. Intuitively, self-attention is considered one of the more interpretable models. It automatically captures the correlation between tokens through attention between itself and itself. In fact, the paper "Attention is All You Need" provided the following visualization, which seems quite reasonable:
But does the self-attention mechanism really work this way? Is this "token-to-token" attention necessary? Recently, Google’s new paper "Synthesizer: Rethinking Self-Attention in Transformer Models" conducted some "imaginative" explorations of the self-attention mechanism, and the results therein might subvert our understanding of self-attention.
Self-Attention
The popularity of self-attention models began with the paper "Attention is All You Need" published by Google in 2017. Its foundation is Scaled-Dot Attention, defined as follows: Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}) = \text{softmax}\left(\frac{\boldsymbol{Q}\boldsymbol{K}^{\top}}{\sqrt{d_k}}\right)\boldsymbol{V} Where \boldsymbol{Q}\in\mathbb{R}^{n\times d_k}, \boldsymbol{K}\in\mathbb{R}^{m\times d_k}, \boldsymbol{V}\in\mathbb{R}^{m\times d_v}, and softmax is normalized along the m dimension. Self-attention, for the same \boldsymbol{X}\in \mathbb{R}^{n\times d}, obtains \boldsymbol{Q}=\boldsymbol{X}\boldsymbol{W}_q, \boldsymbol{K}=\boldsymbol{X}\boldsymbol{W}_k, \boldsymbol{V}=\boldsymbol{X}\boldsymbol{W}_v through different projection matrices \boldsymbol{W}_q, \boldsymbol{W}_k, \boldsymbol{W}_v \in \mathbb{R}^{d\times d'}, and then performs Attention: \begin{aligned} SelfAttention(\boldsymbol{X}) &= Attention(\boldsymbol{X}\boldsymbol{W}_q, \boldsymbol{X}\boldsymbol{W}_k, \boldsymbol{X}\boldsymbol{W}_v) \\ &= \text{softmax}\left(\frac{\boldsymbol{X}\boldsymbol{W}_q \boldsymbol{W}_k^{\top}\boldsymbol{X}^{\top}}{\sqrt{d_k}}\right)\boldsymbol{X}\boldsymbol{W}_v \end{aligned} As for Multi-Head Attention, it is simply the Attention operation repeated multiple times under different parameters, with the outputs concatenated—a relatively simple enhancement.
Imaginative Explorations
Essentially, self-attention transforms an n\times d matrix \boldsymbol{X} into an n\times d' matrix \boldsymbol{A}\boldsymbol{X}\boldsymbol{W}_v using an n\times n matrix \boldsymbol{A} and a d\times d' matrix \boldsymbol{W}_v. The matrix \boldsymbol{A} is generated dynamically: \boldsymbol{A}=\text{softmax}\left(\boldsymbol{B}\right),\quad\boldsymbol{B}=\frac{\boldsymbol{X}\boldsymbol{W}_q \boldsymbol{W}_k^{\top}\boldsymbol{X}^{\top}}{\sqrt{d_k}} For matrix \boldsymbol{B}, it is essentially a combination of dot products between pairs of vectors in \boldsymbol{X}, which is why we call it "token-to-token" Attention.
So, we arrive at the question posed earlier: Is "token-to-token" necessary? Can we generate this matrix \boldsymbol{B} through other means? Google’s paper explores several new forms, collectively referred to as Synthesizer.
Dense Form
The first form is called Dense in the original paper: \boldsymbol{B} needs to be of size n\times n, and \boldsymbol{X} is n\times d, so only a d\times n transformation matrix \boldsymbol{W}_a is needed to turn it into n\times n: \boldsymbol{B}=\boldsymbol{X}\boldsymbol{W}_a This is actually equivalent to fixing \boldsymbol{K} as a constant matrix \boldsymbol{W}_a^{\top}. Of course, the original paper makes it slightly more complex by using two Dense layers: \boldsymbol{B}=\text{relu}\left(\boldsymbol{X}\boldsymbol{W}_1 + \boldsymbol{b}_1\right)\boldsymbol{W}_2 + \boldsymbol{b}_2 But the underlying idea remains the same.
Random Form
If the Dense form is equivalent to fixing \boldsymbol{K} as a constant matrix, can we be even more "imaginative": fix \boldsymbol{Q} as a constant matrix? In this case, the entire \boldsymbol{B} is equivalent to a constant matrix: \boldsymbol{B}=\boldsymbol{R} The original paper indeed experimented with this form, calling it Random. As the name suggests, \boldsymbol{B} is randomly initialized and can be chosen to be updated during training or not. According to the paper, fixed-pattern Attention first appeared in the paper "Fixed Encoder Self-Attention Patterns in Transformer-Based Machine Translation", with the difference being that the Attention matrix there was calculated by a function, whereas in this Google paper, it is completely randomly initialized. Formally, Random is actually equivalent to a Depthwise Separable Convolution operation.
Low-rank Factorization
The two new forms above often face the problem of having too many parameters, so it is natural to think of reducing the number of parameters through low-rank factorization. For Dense and Random, the original paper also proposed and verified corresponding low-rank factorization forms, called Factorized Dense and Factorized Random, respectively.
Factorized Dense generates two matrices \boldsymbol{B}_1, \boldsymbol{B}_2 of size n\times a and n\times b via the Dense method, where ab=n. Then \boldsymbol{B}_1 is repeated b times and \boldsymbol{B}_2 is repeated a times to obtain corresponding n\times n matrices \tilde{\boldsymbol{B}}_1, \tilde{\boldsymbol{B}}_2. Finally, they are multiplied element-wise (personally, I feel \tilde{\boldsymbol{B}}_2 should be transposed before multiplication, but the original paper does not mention this) to synthesize an n\times n matrix: \boldsymbol{B}=\tilde{\boldsymbol{B}}_1 \otimes \tilde{\boldsymbol{B}}_2
As for Factorized Random, it is easy to understand. Instead of a single n\times n matrix \boldsymbol{R}, it becomes two n\times k matrices \boldsymbol{R}_1, \boldsymbol{R}_2, and then: \boldsymbol{B}=\boldsymbol{R}_1\boldsymbol{R}_2^{\top}
Mixed Mode
So far, including standard self-attention, we have five different schemes for generating matrix \boldsymbol{B}. They can also be mixed: \boldsymbol{B}=\sum_{i=1}^N \alpha_i \boldsymbol{B}_i where \boldsymbol{B}_i represents different forms of self-attention matrices, and \sum_{i=1}^N \alpha_i=1 are learnable parameters.
Analysis of Results
We have introduced several new forms of self-attention collectively known as Synthesizer. Their common feature is that they do not maintain the "token-to-token" form; Random, in particular, completely abandons the dynamic characteristics of original attention and becomes a static matrix.
So, how effective are these new self-attention forms? And how do they challenge our understanding of the self-attention mechanism?
Machine Translation
The first evaluation task is machine translation, which compares the effects of various self-attention forms in detail:
I don’t know what readers think, but these Synthesizer results certainly challenged my understanding of self-attention. The table shows that, except for the fixed Random, all self-attention forms perform basically the same. Even the fixed Random has respectable results. This indicates that our previous understanding and explanation of self-attention were too one-sided and did not reveal the true reason why self-attention is effective.
Summarization and Dialogue
Next are the results on summarization and dialogue generation tasks:
In the automatic summarization task, standard attention performs better. However, in the dialogue generation task, the result is the opposite: standard self-attention is the worst, while Dense (D) and Random (R) are the best. When Dense and Random are mixed with standard self-attention (i.e., D+V and R+V), the performance actually degrades. This shows that standard attention has no "monopolistic" advantage. While several Synthesizers look like "degenerations" of standard attention, they are in fact independent and have their own advantages.
Pre-training + Fine-tuning
Finally, for most readers, the concern might be the effect of "pre-training + fine-tuning"—that is, how does performance change after replacing the self-attention in models like BERT? The original paper did conduct this experiment, using T5 as the baseline instead of BERT. The results are as follows:
In these results, Dense and Random appear inferior compared to standard self-attention. This suggests that Dense and Random might perform well on single tasks but have weaker transferability. However, it cannot be denied that self-attention like Random significantly improves computational efficiency by skipping the \boldsymbol{Q}\boldsymbol{K}^{\top} matrix operation. Therefore, if a way can be found to solve this transferability issue, the Transformer model family might undergo a major overhaul.
Summary
This article introduced Google’s new work, Synthesizer, which is a reflection and exploration of the currently popular self-attention mechanism. The paper proposes several new self-attention mechanisms and conducts quite sufficient experiments. The experimental results are likely to challenge our existing understanding of the self-attention mechanism and are well worth reading.
Original Address: https://kexue.fm/archives/7430