Google's New Work Synthesizer: We Still Don't Understand Self-Attention Well Enough
The black box that is deep learning is far darker than we imagine.
Opening remarks #]
The physicist Feynman is said to have once remarked [source]]: "If you think you understand quantum mechanics, you don't understand quantum mechanics." I've increasingly come to feel that in this quote, "quantum mechanics" could just as well be replaced with "deep learning." Although deep learning has proven effective in an ever-growing number of domains, our ability to explain it remains remarkably weak. Of course, in recent years there has been no shortage of work aimed at opening up this black box, but unfortunately most of it is "Monday-morning quarterbacking" — offering explanations that are barely convincing after the fact, based on existing experimental results, without being able to construct and understand model principles from the top down, let alone make any forward-looking predictions.
This post focuses on the self-attention mechanism. Intuitively, self-attention is one of the more interpretable mechanisms out there: it automatically captures the relationships between tokens through attention between the sequence and itself. Indeed, in the original Attention is All You Need] paper, a seemingly quite reasonable visualization was presented:
]
Visualization of attention from the "Attention is All You Need" paper
But does self-attention really work this way? Is this "token-to-token" form of attention really necessary? A recent Google paper, Synthesizer: Rethinking Self-Attention in Transformer Models], carries out some rather "wild" explorations of the self-attention mechanism, and the results may well overturn our understanding of self-attention.
more
Self-Attention #]
The popularity of self-attention models began with Google's 2017 paper Attention is All You Need]. For a gentler introduction, readers can also check out my earlier post A brief reading of "Attention is All You Need" (intro + code)]. Its foundation is Scaled-Dot Attention, defined as follows:
\begin{equation}Attention(\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}) = softmax\left(\frac{\boldsymbol{Q}\boldsymbol{K}^{\top}}{\sqrt{d_k}}\right)\boldsymbol{V}\end{equation}
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 normalizes along the $m$ dimension. Self-attention, then, takes the same $\boldsymbol{X}\in \mathbb{R}^{n\times d}$ and applies different projection matrices $\boldsymbol{W}_q,\boldsymbol{W}_k,\boldsymbol{W}_v\in\mathbb{R}^{d\times d'}$ to obtain $\boldsymbol{Q}=\boldsymbol{X}\boldsymbol{W}_q,\boldsymbol{K}=\boldsymbol{X}\boldsymbol{W}_k,\boldsymbol{V}=\boldsymbol{X}\boldsymbol{W}_v$, and then performs attention on these, i.e.
\begin{equation}\begin{aligned} SelfAttention(\boldsymbol{X}) =&\, Attention(\boldsymbol{X}\boldsymbol{W}_q, \boldsymbol{X}\boldsymbol{W}_k, \boldsymbol{X}\boldsymbol{W}_v)\\ =&\, 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}\end{equation}
As for multi-head attention, it is simply the attention operation repeated multiple times with different parameters, with the outputs concatenated — a fairly straightforward enhancement. For further generalizations of it, see Breaking the Bottleneck: Building a More Powerful Transformer].
Flights of Fancy #]
Fundamentally, self-attention works by using a $n\times n$ matrix $\boldsymbol{A}$ and a $d\times d'$ matrix $\boldsymbol{W}_v$ to transform the original $n\times d$ matrix $\boldsymbol{X}$ into the $n\times d'$ matrix $\boldsymbol{A}\boldsymbol{X}\boldsymbol{W}_v$, where the matrix $\boldsymbol{A}$ is generated dynamically, i.e.
\begin{equation}\boldsymbol{A}=softmax\left(\boldsymbol{B}\right),\quad\boldsymbol{B}=\frac{\boldsymbol{X}\boldsymbol{W}_q \boldsymbol{W}_k^{\top}\boldsymbol{X}^{\top}}{\sqrt{d_k}}\end{equation}
The matrix $\boldsymbol{B}$ is, at its core, simply a collection of pairwise inner products between vectors in $\boldsymbol{X}$, which is why we call it "token-to-token" attention.
]
Comparison between Synthesizer self-attention and standard self-attention
This brings us to the question raised earlier: is "token-to-token" attention really necessary? Could this matrix $\boldsymbol{B}$ be generated some other way? Google's paper does exactly this — it dreams up several new forms and puts them to the test, collectively naming them Synthesizer.
Dense Form #]
The first form, called Dense in the original paper, works as follows. $\boldsymbol{B}$ needs to be of size $n\times n$, and $\boldsymbol{X}$ is of size $n\times d$, so all we need is a $d\times n$ transformation matrix $\boldsymbol{W}_a$ to turn it into $n\times n$, i.e.
\begin{equation}\boldsymbol{B}=\boldsymbol{X}\boldsymbol{W}_a\end{equation}
This is essentially equivalent to fixing $\boldsymbol{K}$ as a constant matrix $\boldsymbol{W}_a^{\top}$. Of course, the original paper takes things a bit further, using two Dense layers:
\begin{equation}\boldsymbol{B}=\text{relu}\left(\boldsymbol{X}\boldsymbol{W}_1 + \boldsymbol{b}_1\right)\boldsymbol{W}_2 + \boldsymbol{b}_2\end{equation}
but conceptually nothing changes.
Random Form #]
We just said the Dense form amounts to fixing $\boldsymbol{K}$ as a constant matrix — can we go even further and just fix $\boldsymbol{Q}$ as a constant matrix too? In that case the entire $\boldsymbol{B}$ becomes a constant matrix, i.e.
\begin{equation}\boldsymbol{B}=\boldsymbol{R}\end{equation}
The original paper actually does experiment with this form, calling it Random — as the name suggests, $\boldsymbol{B}$ is randomly initialized, and one can choose whether or not to update it during training. According to the paper, fixed-form attention first appeared in Fixed Encoder Self-Attention Patterns in Transformer-Based Machine Translation], though the difference is that the attention matrix there is computed by a function, whereas in this Google paper it's completely randomly initialized. In terms of form, Random is essentially equivalent to a depthwise separable convolution operation.
Low-Rank Decomposition #]
Both new forms above tend to face the problem of having too many parameters, so it's natural to think of using low-rank decomposition to reduce the parameter count. For both Dense and Random, the original paper proposes and validates corresponding low-rank decomposition forms, called Factorized Dense and Factorized Random, respectively.
Factorized Dense generates two matrices $n\times a, n\times b$, $\boldsymbol{B}_1,\boldsymbol{B}_2$, where $ab=n$, via a Dense-style transformation; it then repeats $\boldsymbol{B}_1$ $b$ times and $\boldsymbol{B}_2$ $a$ times to obtain corresponding $n\times n$ matrices $\tilde{\boldsymbol{B}}_1,\tilde{\boldsymbol{B}}_2$, and finally multiplies them elementwise (personally I feel $\tilde{\boldsymbol{B}}_2$ should probably be transposed before this multiplication, though the original paper doesn't mention this), yielding a $n\times n$ matrix:
\begin{equation}\boldsymbol{B}=\tilde{\boldsymbol{B}}_1 \otimes \tilde{\boldsymbol{B}}_2\end{equation}
As for Factorized Random, it's easy to understand: instead of one whole $n\times n$ matrix $\boldsymbol{R}$, we now have two $n\times k$ matrices $\boldsymbol{R}_1,\boldsymbol{R}_2$, and then
\begin{equation}\boldsymbol{B}=\boldsymbol{R}_1\boldsymbol{R}_2^{\top} \end{equation}
Mixture Mode #]
At this point, including standard self-attention, we have five different schemes for generating the matrix $\boldsymbol{B}$, and they can also be mixed together, i.e.
\begin{equation}\boldsymbol{B}=\sum_{i=1}^N \alpha_i \boldsymbol{B}_i\end{equation}
where $\boldsymbol{B}_i$ are the different forms of the self-attention matrix, and $\sum\limits_{i=1}^N \alpha_i=1$ are learnable parameters.
Analysis of Results #]
We've now introduced several new forms of self-attention collectively called Synthesizer. What they have in common is that none of them preserve the "token-to-token" form — Random, in particular, entirely abandons the dynamic nature of the original attention mechanism, turning it into a static matrix. So how do these new forms of self-attention actually perform? And how do they challenge our understanding of the self-attention mechanism?
Machine Translation #]
The first evaluation task is machine translation, where the various forms of self-attention are compared in detail:
]
Comparison of Synthesizer performance on machine translation tasks
I don't know what readers make of this, but these Synthesizer results certainly challenged my own understanding of self-attention. The table shows that, apart from the fixed Random variant, essentially all forms of self-attention perform about the same — and even fixed Random achieves a passable result. This suggests that our previous understanding and explanations of self-attention have been far too narrow, and have failed to reveal the true reason why self-attention works.
Summarization and Dialogue #]
Next are the results on summarization and dialogue generation tasks:
]
Comparison of Synthesizer performance on summarization and dialogue tasks
On the summarization task, standard attention performs relatively well; but on dialogue generation, the results are reversed — standard self-attention performs the worst, while Dense (D) and Random (R) perform the best, and when Dense and Random are mixed with standard self-attention (i.e., D+V and R+V), performance actually gets worse. This shows that standard attention has no clear-cut advantage of "reigning supreme," and while several of the Synthesizer variants might look like "degenerate" forms of standard attention, in fact they are independent of one another, each with its own strengths.
Pretraining + Fine-Tuning #]
Finally, for those of us who are more ordinary readers, what we probably care most about is how well this works under the "pretraining + fine-tuning" paradigm — that is, how does performance look if we replace the self-attention in a model like BERT? The original paper does indeed run this experiment, though the baseline used is T5 rather than BERT. The results are as follows:
]
Comparison of Synthesizer performance under "pretraining + fine-tuning"
In this result, Dense and Random come out looking somewhat inferior compared to standard self-attention, suggesting that Dense and Random might perform well on a single task but transfer poorly. That said, we can't dismiss the fact that a form of self-attention like Random, by directly eliminating the $\boldsymbol{Q}\boldsymbol{K}^{\top}$ matrix operation, brings a clear boost in computational efficiency. So if someone can find a way to solve this transferability problem, who knows — the Transformer family of models might well be in for a major shakeup.
Summary #]
This post introduced Google's new work, Synthesizer, which reflects on and explores the currently popular self-attention mechanism. The paper proposes several new forms of self-attention and backs them up with fairly thorough experiments, and the results may well challenge our existing understanding of self-attention. Well worth a read.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.