OpenAI, Week 5-6 // Implementing Transformer in PyTorch

These past two weeks, I've been studying Transformers and coding up my own implementation in PyTorch. There are quite a few excellent step-by-step guides to understanding the Transformer architecture, with gorgeous visualizations, so I'm not going to try to compete! Instead, I want to achieve two things with this post: 1) share my curated curriculum recommendations, and 2) describe a few subtleties that were not addressed in any of the tutorials I came across.

1) Curriculum

As I said, there are a lot of options when it comes to learning material - a potentially overwhelming number of them. Here are my favorites, in the order I recommend accessing them:

• The Illustrated Transformer by Jay Alammar
• Attention Is All You Need: The Transformer by Lennart Van der Goten
• Walkthrough: The Transformer Architecture by Matthew Barnett
• Positional Encoding by Amirhossein Kazemnejad
• The Annotated Transformer by Alexander Rush

2) Subtleties

Here are a couple of things I scratched my head over:

The dimension of the embedding does not have to match the dimension of the model.

In every implementation I came across, the dimension of the model defaulted to the dimension of the embedded input element. That is, if each word in a sequence is represented as an $m$-dimensional vector, and if $W^Q$, $W^K$, $W^V$ are ($m \times d_k$), ($m \times d_k$), ($m \times d_v$) (with $d_v$ often set to be equal to $d_k$), most implementations set $d_k = m$ as well (or more often $d_k = m / N_{\rm heads}$, for multi-headed attention). As far as I can tell, there is no obvious reason to do this.

At first, I thought it was necessary to make the residual connections work without having to project or pad the input vectors. But then I realized that isn't true; the input has dimensions ($L_{\rm seq} \times m$) and so does the final self-attention vector $Z$. The subtlety is that the matrices $W^Q$, $W^K$, $W^V$  are ($m \times d_k$) but the matrix $W^O$ that takes the concatenated multi-head attention and turns it into the final self-attention vector is ($d_k \times m$). So the residual connections work out even without $d_k = m$. 

The decoder.

The encoder part of the Transformer made perfect sense to me. The decoder mystified me. The autoregressive nature of language models - wherein one tries to guess the probability of the next word given the previous words - was hard for me to wrap my head around. I understand that it is a convenient way to frame language problems as it allows for self-supervised training on massive unstructured datasets, but it's a very odd way to think about language. In particular, the Transformer implementations I saw used "teacher forcing," such that the decoder predicts the next word at each time step based on the previous words in the "true" sentence rather than the previous words it has already generated using the model. That seemed particularly strange to me.