Bahdanau's attention (additive)

A model with attention needs one weight for each input position at every output step. The previous lesson, on the idea of attention, set up everything around those weights: they come out of a softmax over the positions of the source, they mix the encoder states into a context vector of each step's own, and that vector enters the decoder's recurrence through the matrix $\mathbf{W}_{ch}$. What that lesson didn't write down is what produces the weights. Before the softmax there's a function, $a$, that takes a pair (where the decoder is, and what's at one input position) and returns a score. So far it has only had to do two things: return a number, and be differentiable. This lesson fills it in with the form Bahdanau gave it in 2014.

The question is small and concrete: what function takes two vectors and returns a number saying how well they correspond? Bahdanau's answer is a tiny neural network with a single hidden layer. But the network isn't what teaches the most; the constraint that decides its shape is. The score has to be a function of the pair. A function that gives an opinion on each vector separately and adds the two opinions up doesn't score a correspondence, only two unrelated things, and further down you'll see that this destroys the whole mechanism.

Go back to the previous lesson's map and look at three rows. The rows for i and read take nearly everything from the same position, the one holding leí, which says i read in a single word, with weights of $0.74$ and $0.81$. The next row, for a, gives that same position $0.05$. The state $\bar{\mathbf{h}}_2$ hasn't changed between those three rows (the encoder computed it once, and there it stays), and yet it receives two high scores and one low one. The only thing that differs is which decoder state it was compared with each time.

A small network that scores pairs

The additive score projects the two states into a space of its own, with $d_a$ coordinates, and adds them there:

with $\mathbf{W}_{sa}, \mathbf{W}_{ha} \in \mathbb{R}^{d_a \times d_h}$, one for each side. Both have $d_h$ columns because the encoder and the decoder are the same width in this block, but the sum doesn't require that. All it asks is that the two projections land in the same space, that is, that the two matrices have $d_a$ rows; the number of columns in each is its own side's business. If the decoder states had $d_s$ coordinates, $\mathbf{W}_{sa}$ would be $d_a \times d_s$ and the rest of the score would stay as it is. In fact, in Bahdanau's paper the two widths do differ: its encoder reads the sentence in both directions and concatenates the two states, so its encoder states are twice as wide as the decoder's.

That sum is the pre-activation of a hidden layer, and what comes next is what any layer already does: squash it and read it out.

with $\mathbf{v}_a \in \mathbb{R}^{d_a}$. This is a multilayer perceptron (MLP) from Block 2 at its smallest size: one hidden layer of $d_a$ units with $\tanh$, and a linear output of a single neuron. That neuron is the one from the lesson on the artificial neuronBlock 2 · Lesson 1The artificial neuronA weighted sum and a threshold are all a single unit needs to turn a vector into a decision. Its weights sit perpendicular to a boundary that splits the space in two, and the number it computes before deciding measures how far the input has landed from that boundary., $z = \mathbf{w}^{\top}\mathbf{x} + b$: its input is the hidden layer, $\tanh(\mathbf{p}_{ij})$, and its weights are $\mathbf{v}_a$, one per hidden unit. We call $\mathbf{v}_a$ the read-out vector, because it turns the $d_a$ numbers of the hidden layer into the single number that is the score.

Neither layer has a bias, for different reasons. The output's would be redundant, because it would add the same number to the whole row, and a constant per row cancels out in the softmax. The hidden layer's wouldn't be: inside the $\tanh$ it shifts the point where each unit operates, closer to saturation or further from it, and that does change the scores. Bahdanau leaves it out and many implementations put it in; here I follow the paper's formula.

Count the weights: $2 d_a d_h + d_a$. Neither $T_x$ nor $T_y$ appears, and neither does the pair. There's a single set of weights, and it scores the $T_x \cdot T_y$ pairs of this sentence and of every other, which is the same economy as the previous block's recurrence (one $\mathbf{W}_{hh}$ for every position), applied to a grid instead of a row.

Without the non-linearity, every step would look at the same place

Remove the $\tanh$ and look at what's left. The read-out vector distributes over the sum:

two numbers added together without either having looked at the other. The first, $\mathbf{v}_a^{\top}\mathbf{W}_{sa}\mathbf{s}_{i-1}$, depends on step $i$ but not on position $j$. So in row $i$, that is, among the $T_x$ scores $e_{i1}, \dots, e_{iT_x}$ that step $i$'s softmax turns into weights, it's the same in every score: a constant added to the whole row. And a constant like that doesn't change the softmax. It's the shift invariance from the lesson on loss functionsBlock 2 · Lesson 5Loss functions: MSE and cross-entropyA loss sums up in one number how badly the network is doing, and between two candidates its derivative is what decides. Mean squared error hands over its smallest correction exactly where the network is wrong with the most confidence; cross-entropy, which comes from reading the output as a probability, leaves that derivative at the plain difference between what was predicted and the label., the same property that lets you subtract the maximum before exponentiating, and here it makes the first number cancel between numerator and denominator:

The index $i$ has disappeared from the right-hand side. All $T_y$ steps receive the same row of weights, and therefore the same context vector, so the model is back to being the fixed summary of the earlier lessons: a learned alignment, yes, but still frozen, only now it sits wherever training puts it instead of on the last token. The $\tanh$ is what prevents that split. It squashes the sum before the read-out, and a non-linear function of a sum doesn't distribute over its terms, so the score can't come apart again into one opinion about the decoder plus another about the encoder.

How the gradient reaches it

Now there are parameters to train, and the question is whether the loss reaches them. The chain starts at the mixture, where $\alpha_{ij}$ multiplies a known vector:

and continues through the softmax, which makes each score affect every weight in its row. Out the other side comes the derivative with respect to a score:

Read it slowly, because it says how attention learns where to look. It has two parts. The first is what happens to the mixture when the score goes up: if $e_{ij}$ grows a little, position $j$ gains weight at the expense of the others, and $\mathbf{c}_i$ shifts towards $\bar{\mathbf{h}}_j$. Precisely, $\partial \mathbf{c}_i / \partial e_{ij} = \alpha_{ij}\left(\bar{\mathbf{h}}_j - \mathbf{c}_i\right)$, which is the same calculation as in the drop-down, with $\bar{\mathbf{h}}_k$ in place of $\partial \ell / \partial \alpha_{ik}$. The second part is $\nabla_{\mathbf{c}_i}\ell$, whose opposite points to where the loss goes down, that is, where gradient descent would want to take the mixture. The derivative is the dot product of the two, so it measures whether moving $\mathbf{c}_i$ towards $\bar{\mathbf{h}}_j$ takes it where it ought to go.

That tells us what training does with the score. If moving the mixture towards $\bar{\mathbf{h}}_j$ makes the loss go down, the score goes up; if it makes the loss go up, the score goes down; and if $\bar{\mathbf{h}}_j$ coincides with $\mathbf{c}_i$, raising the score doesn't move the mixture, so the score gets no correction at all. A score, then, isn't corrected according to what its state is worth on its own, but according to what that state adds to the mixture the row has already built. It makes sense for this to be a comparison and not a grade: the weights sum to $1$, so raising one means lowering the others, and there's no way to reward one position without saying which one pays for it.

From there back, it's the chain rule as you already know it. The read-out hands back $\tanh(\mathbf{p}_{ij})$, the $\tanh$ contributes its derivative $1 - \tanh^{2}$, and each matrix collects the outer product with the vector it multiplied:

All $T_x \cdot T_y$ pairs add into the same three arrays, which is what sharing weights means.

There's one more consequence, and it isn't in the parameters. Look at where $\bar{\mathbf{h}}_j$ appears: in the mixture, as one of the terms, and in the $T_y$ scores, as an argument of $a$. That's two places, so the gradient reaches it along two routes, and the two contributions have to be added:

On top of those two comes the one the previous block already had, the gradient that comes down from $\bar{\mathbf{h}}_{j+1}$ through the encoder's recurrence. And here is where it gets interesting: the two routes above do not go through that recurrence. They run from step $i$ of the decoder straight to position $j$ of the input, without travelling the stretch of recurrence that separates that position from the end of the sentence. They don't multiply by $\mathbf{W}^{\text{enc}}_{hh}$ even once, and their length doesn't grow as $j$ moves further away. The product of matrices that the lesson on the vanishing gradientBlock 3 · Lesson 4The vanishing gradientMultiplying by the same recurrent matrix at every step makes the gradient exponential in the distance it travels. What guarantees that it vanishes is γ·σ_max < 1, the steepest tanh slope times the largest singular value of W_hh; the spectral radius only sets the pace of the matrix on its own. Gradient clipping stops the explosion but can't rescue a gradient that has already vanished, and that's what will force a change of architecture. identified as the cause of the problem doesn't appear on either of them.

The score, in NumPy

The first cell computes the whole grid. It stores the states by rows (one row per step, one row per position), which is why the matrices appear transposed relative to the equations. There are two things to look at: how closely the sharing-out that freshly initialised weights produce resembles a map, and what happens when you ask for the version without the $\tanh$.

The freshly initialised map says nothing: the largest weight in each row is around $0.18$ when an even share would be $0.167$, so the eight rows are practically the mean of the sentence. The explorable's interpretable alignment doesn't come from the architecture but from training; what the architecture guarantees is a route along which to train it. And the output after the blank line is the previous section checked by hand: without the $\tanh$, row $1$ and row $8$ come out equal to the last decimal.

The second cell follows the two routes separately. It needs a loss, so it uses the cheapest one possible (one whose gradient with respect to each $\mathbf{c}_i$ is a given vector), which means the mechanism is the only thing being measured.

The last line is the check: added together, the two routes reproduce the numerical gradient to the order of $10^{-10}$, while on its own each one is wrong in the first digit. The norms tell you something more, and it's worth being honest about it: with untrained weights, the value route is eleven times larger than the score route. At the start, almost all of the gradient arrives along the route it would take even if the alignment were uniform, and the small route, the one through $a$, is the only one that can pull the map out of that uniformity. Change the scale of va to $2.0$ and run both cells again: the scores spread apart, the weights sharpen and the two norms draw closer.

Test your intuition

Five questions: what breaks without the $\tanh$, how many weights the function $a$ carries, how the gradient reaches an encoder state, what the derivative of a score says, and which part of the computation can come out of the loop.

The last question leaves a bill unsettled. The two projections come out of the loop because each depends on a single index, but between the sum and the read-out vector there's a $\tanh$, and it forces the grid of $T_x \cdot T_y$ vectors of $d_a$ coordinates to be built in full. A sixty-token sentence translated into sixty tokens asks for $3\,600$ of those vectors, and none of the operations that produce them is a large matrix product: they're sums and $\tanh$s over a three-dimensional block. The hardware that trains these networks is far better at one thing than at anything else, and this isn't it.

There's another way to score a pair that avoids the whole problem, and its formula fits in five symbols: Luong's attention, the multiplicative kind, the subject of the next lesson. Instead of projecting, adding and squashing, it multiplies the two states through a matrix, so the complete grid of scores becomes one matrix product. That lesson compares the two scores (what each costs, what each can do), and that's where the shape of Block 5 starts to be decided.