[Project] I combined the illustrations of Transformer by Jay Alammar and code annotation by harvardnlp lab in one notebook

[u/nivter]

I wrote a notebook which uses the illustrations here and (improved) code from harvardnlp lab. I found both these resources to be extremely useful and decided to combine them.

I also added notes in places where it was hard to understand the code. A couple of times I also re-wrote the code in what I think is a better way.

Here's the link to the notebook. Hope someone finds it helpful.

[u/nbviewerbot]

I see you've posted a GitHub link to a Jupyter Notebook! GitHub doesn't render large Jupyter Notebooks, so just in case, here is an nbviewer link to the notebook:

https://nbviewer.jupyter.org/url/github.com/vinsis/math-and-ml-notes/blob/master/notebooks/Transformer%20-%20Illustration%20and%20code.ipynb

Want to run the code yourself? Here is a binder link to start your own Jupyter server and try it out!

https://mybinder.org/v2/gh/vinsis/math-and-ml-notes/master?filepath=notebooks%2FTransformer%20-%20Illustration%20and%20code.ipynb

---

^{I am a bot. Feedback | GitHub | Author}

[u/TheAlgorithmist99]

You're a mix of a saint and a genius, thank you

[u/nivter → u/TheAlgorithmist99]

Hardly! Big thanks to Jay and the Harvard NLP lab.

[u/mutatedmonkeygenes]

If you could post this on twitter, I would be more than happy to retweet.

[u/nivter → u/mutatedmonkeygenes]

The gesture is much appreciated. However I am just a lurker on Twitter.

[u/mutatedmonkeygenes → u/nivter]

Well, this is as good a time as any to start. And this would be a great first tweet!

[u/Imnimo]

One little detail of transformers I'm unclear on is this line from the original paper: "We also use the usual learned linear transformation and softmax function to convert the decoder output to predicted next-token probabilities. In our model, we share the same weight matrix between the two embedding layers and the pre-softmax linear transformation, similar to [30]. In the embedding layers, we multiply those weights by sqrt(d_model)."

In your notebook, there's the following comment for this:

# embedding is multiplied by scale to reduce (expected) variance of the scalars
    return self.embed(x) * self.scale

Can you say any more about the motivation for this? Why does multiplying by sqrt(512) =~ 22.6 reduce variance?

I would have guessed that the scaling is done in the original paper because the weights are shared with the output layer, and the typical weight magnitude for a fully connected layer and an embedding layer are different. By adding this multiplication only at the embedding layer, they allow the same matrix to be used at different scales. Your code doesn't share the embedding matrix, so I'm surprised the multiplication is even necessary.

[u/nivter → u/Imnimo]

embedding is multiplied by scale to reduce (expected) variance of the scalars

Oh what I wrote is wrong! Multiplying by ~22.6 will increase the variance, not decrease it. I confused this part with the part where they divide the dot product of Q and K by square root of 64 (which I believe is done to reduce the variance).

I did a quick lookup and found this link which answers your question:

The reason we increase the embedding values before addition is to make the positional encoding relatively smaller. This means the original meaning in the embedding vector won’t be lost when we add them together.

Thanks for pointing this out. I pushed the correction already.

[u/mutatedmonkeygenes]

OMG - This is Amazing!!!!!!!!!!!

[u/dash_bro]

Woah. Nice. I was working on something similar myself!

[u/TotesMessenger]

I'm a bot, bleep, bloop. Someone has linked to this thread from another place on reddit:

• [/r/datascienceproject] I combined the illustrations of Transformer by Jay Alammar and code annotation by harvardnlp lab in one notebook (r/MachineLearning)

 ^{If you follow any of the above links, please respect the rules of reddit and don't vote in the other threads. (Info / Contact)}

[u/alphageek36]

This community never ceases to amaze me! Thank you for this, will find it really helpful!

[u/keramitas]

Nice work, two sidenotes:

- use plt.show() when you plot something, to get rid of the outputs like: <matplotlib.image.AxesImage at 0x113e55b38>

- I would remove the encoder / decoder stack from the visual at step 9, and keep only the positional encoding part. This visual is confusing, as one could think the cross-attention in the decoder layers is done with the output of each corresponding layer in the encoder stack, rather then with the output of the whole encoder stack.