Transformers (and Attention) are Just Fancy Addition Machines

is a comparatively new sub-field in AI, targeted on understanding how neural networks perform by reverse-engineering their inside mechanisms and representations, aiming to translate them into human-understandable algorithms and ideas. That is in distinction to and additional than conventional explainability methods like SHAP and LIME.

SHAP stands for SHapley Additive exPlanations. It computes the contribution of every function to the prediction of the mannequin, domestically and globally, that’s for a single instance in addition to throughout the entire dataset. This enables SHAP for use to find out function significance on the whole for the use case. LIME, in the meantime, works on a single example-prediction pair the place it perturbs the instance enter and makes use of the perturbations and its outputs to approximate an easier substitute of the black-box mannequin. As such, each of those work at a function degree and provides us some rationalization and heuristic to gauge how every enter into the mannequin impacts its prediction or output.

Then again, mechanistic interpretation understands issues at a extra granular degree in that it’s able to offering a pathway of how the stated function is learnt by completely different neurons in several layers within the neural community, and the way that studying evolves over the layers within the community. This makes it adept at tracing paths contained in the community for a specific function and additionally seeing how that function impacts the result. 

SHAP and LIME, then, reply the query “which function contributes essentially the most to the result?” whereas mechanistic interpretation solutions the query “which neurons activate for which function, and the way does that function evolve and have an effect on the result of the community?“

Since explainability on the whole is an issue with deeper networks, this sub-field majorly works with deeper fashions just like the transformers. There are just a few locations the place mechanistic interpretability appears at transformers in another way than the normal manner, considered one of which is multi-head consideration. As we are going to see, this distinction is in reframing the multiplication and concatenation operations as outlined within the “Consideration is All You Want” paper as addition operations which opens an entire vary of recent potentialities.

However first, a recap of the Transformer structure.

Transformer Structure

These are the sizes we work with:

• batch_size B =1;
• sequence size S = 20;
• vocab_size V = 50,000;
• hidden_dims D = 512;
• heads H = 8

Which means that the variety of dimensions within the Q, Okay, V vectors is 512/8 (L) = 64. (In case you don’t keep in mind, an analogy for understanding question, key and worth: The concept is that for a token at a given place (Okay), primarily based on its context (Q) we need to get alignment (reweighing) to the positions it’s related to (V).)

These are the steps upto the eye computation in a transformer. (The form of tensors is assumed for instance for higher understanding. Numbers in italic symbolize the dimension alongside which the matrix is multiplied.)

The desk defined intimately:

1. We begin with one enter sentence of a sequence size of 20 that’s one-hot encoded to symbolize phrases within the vocabulary current within the sequence. Form (B x S x V): (1 x 20 x 50,000)
2. We multiply this enter with the learnable embedding matrix Wₑ of form (V x D) to get the embeddings. Form (B x S x D): (1 x 20 x 512)
3. Subsequent a learnable positional encoding matrix of the identical form is added to the embeddings
4. The resultant embeddings are then copied to the matrices Q, Okay and V. Q, Okay and V every are cut up and reshaped on the D dimension. Form (B x S x D): (1 x 20 x 512)
5. The matrices for Q, Okay and V are every fed to a linear transformation layer that multiplies them with learnable weight matrices every of form (D x L) Wq, Wₖ and Wᵥ, respectively (one copy for every of the H=8 heads). Form (B x H x S x L): (1 x 1 x 20 x 64) the place H=1, as that is the resultant form for every head.
6. Subsequent, we compute consideration with Scaled Dot Product consideration the place Q and Okay (transposed) are multiplied first in every head. Form (B x H x S x L) x (L x S x H x B) → (B x H x S x S): (1 x 1 x 20 x 20). 
7. There’s a scaling and masking step subsequent that I’ve skipped as that isn’t necessary in understanding what’s the completely different manner of MHA. So, subsequent we multiply QK with V for every head. Form (B x H x S x S) x (B x H x S x L) → (B x H x S x L): (1 x 1 x 20 x 64)
8. Concat: Right here, we concatenate the outcomes of consideration from all of the heads on the L dimension to get again a form of (B x S x D) → (1 x 20 x 512)
9. This output is as soon as extra linearly projected utilizing yet one more learnable weight matrix Wₒ of form (D x D). Remaining form we finish with (B x S x D): (1 x 20 x 512)

Reimagining Multi-Head Consideration

Now, let’s see how the sector of mechanistic interpretation appears at this, and we may also see why it’s mathematically equal. On the appropriate within the picture above, you see the module that reimagines multi-head consideration. 

As a substitute of concatenating the eye output, we proceed with the multiplication “inside” the heads itself the place now the form of Wₒ is (L x D) and multiply with QK’V of form (B x H x S x L) to get the results of form (B x S x H x D): (1 x 20 x 1 x 512). Then, we sum over the H dimension to once more finish with the form (B x S x D): (1 x 20 x 512).

From the desk above, the final two steps are what modifications:

Aspect notice: This “summing over” is paying homage to how summing over completely different channels occurs in CNNs. In CNNs, every filter operates on the enter, after which we sum the outputs throughout channels. Identical right here — every head could be seen as a channel, and the mannequin learns a weight matrix to map every head’s contribution into the ultimate output area.

However why is the venture + sum mathematically equal to concat + venture? In brief, as a result of the projection weights within the mechanistic perspective are simply sliced variations of the weights within the conventional view (sliced throughout the D dimension and cut up to match every head).

Let’s give attention to the H and D dimensions earlier than the multiplication with Wₒ. From picture above, every head now has a vector of measurement 64 that’s being multiplied with the burden matrix of form (64 x 512). Let’s denote the consequence by R and head by h.

To get R₁₁, we’ve got this equation: 

R₁,₁ = h₁,₁ x Wₒ₁,₁ + h₁,₂ x Wₒ₂,₁ + …. + h₁ₓ₆₄ x Wₒ₆₄,₁

Now let’s say we had a concatenated the heads to get an consideration output form of (1 x 512) and the burden matrix of form (512, 512) then the equation would have been:

R₁,₁ = h₁,₁ x Wₒ₁,₁ + h₁,₂ x Wₒ₂,₁ + …. + h₁ₓ₅₁₂ x Wₒ₅₁₂,₁

So, the half h₁ₓ₆₅ x Wₒ₆₅,₁ + … + h₁ₓ₅₁₂ x Wₒ₅₁₂,₁ would have been added. However this half being added is the half that’s current in every of the opposite heads in modulo 64 style. Stated one other manner, if there isn’t any concatenation, Wₒ₆₅,₁ is the worth behind Wₒ₁,₁ within the second head, Wₒ₁₂₉,₁ is the worth behind Wₒ₁,₁ within the third head and so forth if we think about that the values for every head sit behind each other. Therefore, even with out concatenation, the “summing over the heads” operation ends in the identical values being added.

In conclusion, this perception lays the muse of transformers as purely additive fashions in that every one the operations in a transformer take the preliminary embedding and add to it. This view opens up new potentialities like tracing options as they’re learnt by way of additions by the layers (referred to as circuit tracing) which is what mechanistic interpretability is about as I’ll present in my subsequent articles.

We have now proven that this view is mathematically equal to the vastly completely different view that multi-head consideration, by splitting Q,Okay,V parallelizes and optimizes computation of consideration. Learn extra about this on this weblog here and the precise paper that introduces these factors is here.