Welcome to The Nonlinear Library, where we use Text-to-Speech software to convert the best writing from the Rationalist and EA communities into audio. This is: Building a transformer from scratch - AI safety up-skilling challenge, published by Marius Hobbhahn on October 12, 2022 on The AI Alignment Forum.
It is not always obvious whether your skills are sufficiently good to work for one of the various AI safety and alignment organizations. There are many options to calibrate and improve your skills including just applying to an org or talking with other people within the alignment community.
One additional option is to test your skills by working on projects that are closely related to or a building block of the work being done in alignment orgs. By now, there are multiple curricula out there, e.g. the one by Jacob Hilton or the one by Gabriel Mukobi.
One core building block of these curricula is to understand transformers in detail and a common recommendation is to check if you can build one from scratch. Thus, my girlfriend and I have recently set ourselves the challenge to build various transformers from scratch in PyTorch. We think this was a useful exercise and want to present the challenge in more detail and share some tips and tricks. You can find our code here.
Building a transformer from scratch
The following is a suggestion on how to build a transformer from scratch and train it. There are, of course, many details we omit but I think it covers the most important basics.
Goals
From the ground up we want to
Build the attention mechanism
Build a single-head attention mechanism
Build a multi-head attention mechanism
Build an attention block
Build one or multiple of a text classification transformer, BERT or GPT. The quality of the final model doesn’t have to be great, just clearly better than random.
Train the model on a small dataset.
We used the polarity dataset for binary text sentiment classification.
We used the AG_NEWS dataset (PyTorch built-in) for BERT and GPT.
Test that the model actually learned something
We looked at the first batch of the test data to see if the model predicted something plausible.
We compared the test loss of a random network with the test loss of the trained network to see if our model is better.
Bonus goals
Visualize one attention head
Visualize how multiple attention heads attend to the words of an arbitrary sentence
Reproduce the grokking phenomenon (see e.g. Neel’s and Tom’s piece).
Answer some of the questions in Jacob Hilton's post.
Soft rules
For this calibration challenge, we used the following rules. Note, that these are “soft rules” and nobody is going to enforce them but it’s in your interest to make some rules before you start. We were
allowed to read papers such as Attention is all you need or the GPT-3 paper.
allowed to read tutorials on attention such as The illustrated transformer (as long as they don’t contain code snippets).
allowed to look at tutorials to build generic models in PyTorch as long as they don’t contain NLP architectures.
allowed to watch videos such as the ones from Yannic Kilcher on NLP
not allowed to look at the source code of any transformer or attention mechanism before you have implemented it ourselves. In case we struggle a lot, we can take a peek after we tried and failed to implement one building block ourselves.
We found Andrej Karpathy’s code helpful for the GPT implementation.
allowed to replace a part with a PyTorch implementation once we have demonstrated that it is equivalent. For example, once we have shown that our attention mechanism produces the same output for the same input as the PyTorch attention mechanism, we can use the PyTorch code block.
allowed to use generic PyTorch functions that are not directly related to the task. For example, we don’t have to write the embedding layer, linear layer or layer-norm from scratch.
Things to look out for
Here are some suggestions on what to look out for during the project
Do I understand the tut...