Attention before Transformers

The encoder–decoder of the previous note compresses the source into one fixed-size state before the decoder writes anything, and the outputs show the cost: fluent sentences about the wrong things. The decoder should be able to look back at the source as it writes, and to look at different parts of it for different target words. Bahdanau, Cho and Bengio (2015) added exactly this and called it attention. This note builds the mechanism, works one step by hand, and shows an experiment that measures what it buys.

Notation. Encoder states $h_1, \dots, h_S \in \mathbb R^{H}$ (one per source token; $2H$ if bidirectional). Decoder state $s_t \in \mathbb R^{H}$ at step $t$. Scores $e_{t,i}$, weights $\alpha_{t,i}$, context vector $c_t$.

Baseline: the final state as the only channel#

The decoder is initialized with $s_0 = f(h_S)$ and then runs on its own. There is no operation by which step $t$ of the decoder can read $h_3$ specifically. The states $h_1, \dots, h_{S-1}$ are computed and discarded.

The mechanism#

At decoder step $t$:

1. Score each encoder state against the decoder state. The simplest score is a dot product, $e_{t,i} = s_t \cdot h_i$; it is large when the two vectors point the same way. Other choices exist (below), and whichever is used, the output is one number per source position.
2. Normalize with a softmax over source positions: $$\alpha_{t,i} = \frac{\exp(e_{t,i})}{\sum_{j=1}^S \exp(e_{t,j})}, \qquad \sum_i \alpha_{t,i} = 1,\ \alpha_{t,i} > 0.$$ The weights form a distribution over the source: “how much of each source token to use now”.
3. Aggregate into a context vector, the weighted sum of encoder states: $$c_t = \sum_{i=1}^S \alpha_{t,i}\, h_i \in \mathbb R^{H}.$$
4. Predict from the decoder state and the context together, for example by concatenating them before the output layer: $P(y_t \mid y_{<t}, x) = \mathrm{softmax}(W_o [s_t; c_t])$, and feeding $c_t$ into the next decoder step as well.

Every operation is differentiable, so the scoring is learned by backpropagation from the translation loss; nothing tells the model which source word corresponds to which target word. The weights are recomputed at every step: $t = 1$ can attend to the first source word and $t = 5$ to the last.

Shapes for one step and a batch of $B$: encoder states $[B, S, H]$, decoder state $[B, H]$, scores $[B, S]$, weights $[B, S]$, context $[B, H]$. Over a whole target of length $T$ the weights form a $[B, T, S]$ tensor, which is what the heatmaps in the translation note display.

Worked example#

Three encoder states in two dimensions, $h_1 = (1, 0)$, $h_2 = (0, 1)$, $h_3 = (1, 1)$, and a decoder state $s = (0.9, 0.2)$:

$i$	$e_i = s \cdot h_i$	$\exp(e_i)$	$\alpha_i$
1	$0.9$	$2.460$	$0.368$
2	$0.2$	$1.221$	$0.183$
3	$1.1$	$3.004$	$0.449$

The denominator is $2.460 + 1.221 + 3.004 = 6.685$. The context is $0.368\,(1, 0) + 0.183\,(0, 1) + 0.449\,(1, 1) = (0.817, 0.632)$. It is a point inside the triangle spanned by the encoder states, closest to the ones the decoder state resembled; nothing was selected, everything was mixed in proportion.

Figure

One attention step: what goes where

The figure runs the four steps of the mechanism on these numbers; the stepper under it shows them one at a time. Compared with the encoder–decoder figure, nothing about the encoder changes. What changes is the line from the encoder to the decoder: instead of one vector handed over once, the decoder reaches all the encoder states at every step, through weights it computes itself.

Scoring functions#

The dot product assumes $s_t$ and $h_i$ live in the same space and have comparable scale. Two generalizations are common (Luong, Pham and Manning, 2015):

Name	$e_{t,i}$	Parameters
dot	$s_t^\top h_i$	none
general (bilinear)	$s_t^\top W h_i$	$W \in \mathbb R^{H \times H}$
additive (Bahdanau et al., 2015)	$v^\top \tanh(W_1 s_t + W_2 h_i)$	$W_1, W_2, v$

The additive form was the original; the dot form is the cheapest and, with a learned projection of both sides, is what the self-attention note uses. Cosine similarity, which divides by the norms, is a possible score too but discards magnitude information that the model can use; the standard choice is the unnormalized dot product, rescaled by a constant in the Transformer.

Evidence: removing the bottleneck from a Transformer#

The clean comparison is not LSTM-with-attention versus LSTM-without, because the two models differ in several ways. A sharper experiment uses the one-layer Transformer of the translation note, whose decoder reads the encoder through exactly the mechanism above (there called cross-attention). Two versions are trained identically for 10 epochs on Multi30k:

Decoder may attend to	Test BLEU
all $S$ encoder states	0.320
only the first encoder state (an artificial bottleneck)	0.190
LSTM encoder–decoder, no attention (previous note)	0.152

Restricting the decoder to one encoder vector costs 0.13 BLEU and brings the Transformer down near the LSTM baseline. Everything else (self-attention inside the decoder, the feed-forward layers, the embeddings) is unchanged, so the difference is attributable to the decoder’s access to the full set of source states. The bottleneck is the problem, and attention is its removal.

The learned weights are interpretable in this task. Plotted as a matrix with target tokens as rows and source tokens as columns, the one-layer model’s cross-attention is close to diagonal, with each English token attending to the German token after the one it translates, because the model is predicting the next token and the aligned source word for the next output is the most useful thing to read.

What attention gives#

• No fixed-size summary. The decoder reads $S$ vectors, however long the source is. Long sentences are no longer compressed harder than short ones.
• Alignment for free. The weights $\alpha_{t,i}$ are a soft word alignment, learned from translation pairs alone.
• A short gradient path. The loss at target step $t$ reaches encoder state $h_i$ through one weighted sum instead of through $S + t$ recurrent steps, which eases the vanishing problem of the RNN note.

What changed#

	Encoder–decoder	With attention
Source information per step	$s_0$, fixed	$c_t$, recomputed
Encoder states used	$h_S$ only	all $h_1, \dots, h_S$
Extra computation per step	—	$S$ scores, one softmax, one weighted sum
Extra parameters	—	none (dot) to $O(H^2)$ (general, additive)
BLEU, Multi30k, 1-layer Transformer decoder	0.190 (bottlenecked)	0.320

Limitations#

• The encoder is still recurrent. The states $h_i$ are computed sequentially, so the encoder cannot be parallelized along the source, and each $h_i$ still carries the source only through a chain of recurrences.
• Attention is only between decoder and encoder. Within the source, token $i$ learns about token $j$ only through the recurrence; within the target, likewise. There is no reason the same scoring could not relate positions of one sequence to each other.
• Cost per step is $O(S)$. Fine for sentences; a concern for documents, and quadratic once every position attends to every other.
• Soft weights are not explanations. A near-diagonal alignment on a caption dataset does not generalize to a claim that attention weights show what the model “uses” (Jain and Wallace, 2019).

What you should remember#

• Attention at step $t$: scores $e_{t,i}$ between the decoder state and each encoder state, softmax over $i$, context $c_t = \sum_i \alpha_{t,i} h_i$, prediction from $[s_t; c_t]$.
• The weights are a distribution over source positions, recomputed per output token and learned from the task loss.
• Dot-product scoring is the cheapest; additive and bilinear variants add parameters.
• Removing the decoder’s access to all encoder states costs 0.13 BLEU on Multi30k in an otherwise identical model.
• Attention removes the bottleneck but not the recurrence.

Notebook#

week04_lstm_attention_transformer.ipynb. Cell numbers count every cell, markdown included, from 1. Quoted numbers come from the run saved with the course; cells changed since then for current library versions have no saved output.

• Cells 58–64 define and train the one-layer Seq2SeqTransformer (hidden size 128, 8 heads); the training cell has no saved output, and the markdown cell 64 states the test BLEU, 0.32. The translation note describes the model.
• Cells 72–79 extract the decoder’s cross-attention weights with a forward hook (the [32, 31, 32] tensor of batch × target × source weights, each row summing to 1) and plot them for six short test pairs; the near-diagonal, one-ahead pattern described above is visible in the saved figure.
• Cells 81–86 define Seq2SeqTransformerNoCA, whose decode passes only encoded_input[:, :1] (the state of the first source token) as the encoder memory, train it the same way; the markdown cell 86 states its BLEU, 0.19.

Next#

Attention compared a decoder state with encoder states. Nothing in the mechanism requires the two sides to come from different networks: a sequence can attend to itself, each position scoring every other position and aggregating their values. Doing that removes the recurrence entirely, because every position’s new representation is a weighted sum computed in one matrix product, in parallel over the sequence. The next note writes self-attention in matrix form, with queries, keys and values, and the Transformer is built on it.

References#

• D. Bahdanau, K. Cho, Y. Bengio. Neural Machine Translation by Jointly Learning to Align and Translate. ICLR, 2015. Additive attention.
• M.-T. Luong, H. Pham, C. D. Manning. Effective Approaches to Attention-based Neural Machine Translation. EMNLP, 2015. Dot, general and concat scores; global versus local attention.
• S. Jain, B. C. Wallace. Attention is not Explanation. NAACL, 2019.
• D. Jurafsky, J. H. Martin. Speech and Language Processing, 3rd ed. draft. Chapter 13 (machine translation; the encoder–decoder with attention).