Last week, I decided to understand transformers. I began this quest because I was actually trying to understand a paper about optimizing transformers. It turns out that it’s difficult to understand an optimization if you don’t pretty deeply understand the mechanism being optimized.

I have actually learned about transformers before. I’ve learned about attention and neural networks and superposition, and yet, I find that each time I reencounter one of these concepts, I have no ability to explain them with any depth. Thus, rather than inevitably forgetting how transformers work again, I thought a more fruitful output of my efforts than would be to capture my intuition for later review the best way I know: in a lesson.

As I was learning about transformers, though, I found myself extrapolating some of the characteristics of the math we see in transformers to the characteristics of my own life (naturally). This first and second parts of this piece are lessons on transformers from first principles, assuming very little prior knowledge on AI (but hopefully helpful even to those with a lot of prior knowledge on AI!). The third part is about my life.

Overview

A transformer is the basic architecture that underlies every modern AI model that we interact with. It was first introduced in 2017 in the seminal paper Attention Is All You Need. There are a lot of incredible resources for learning more about transformers for the similarly curious (I would highly recommend the video series by 3Blue1Brown and a conversation with Claude). Below is my very hand-wavey, metaphorical explanation of what I learned. For simplicity, I refer only to text as input and output (such as in LLMs), but these ideas transfer to every modality (images, audio, etc.).

When interacting with a model, a defining characteristic that differentiates a good experience from a terrible one is the model’s semantic understanding of what you want it to do, often contained in the context of the query that you’ve asked the model and the conversation you’ve had with it thus far. In order to extract this, models break up text into tokens — small clusters of characters — and represent them using vectors that encode their semantic meaning (and their position in the existing sentence).

I picture this is by imagining each word as a point drawn on a really high dimensional graph, where words with similar meanings are clustered together. This initial representation, though, doesn’t encode much about the meaning of that word in the context in which it appears. Suppose we have a model attempt to complete the following two sentences:

A cute, fluffy creature __
A terrifying, animalistic creature __

In the initial vector embedding, if we assume that each word gets its own token, the model’s embedding of creature extracted from the two sentences above is very similar. The question becomes, how does the model learn that these two words have appeared in a vastly different context and adjust the completed sentences accordingly?

Attention

The reason why the concept of attention was so revolutionary was because it, for the first time, allowed a model to look back on its entire context and build a really nuanced understanding of what each token means. Specifically, the model builds an understanding of a given token that is influenced by how all the other tokens interact with it. Previously, neural networks processed sequences step-by-step. For example,

• As recurrent neural networks (RNNs) process each token in sequence, they develop a hidden summary of every token seen so far. The representation of each token is built by the hidden states from the previous token, almost like a long game of telephone. This helps the model encode sequential relationships, like an adjective preceding a noun. It also means that the model is limited by its memory because tokens it saw earlier may not have as strong of a representation — a sequential bottleneck.
• Convolutional neural networks (CNNs) slide a filter over the existing text, learning relationships from words that are clustered together. These filters, though, only allow the model to learn relationships from the words that are close by, meaning that information that is physically far away but should influence the model’s understanding of the meaning of the token being processed don’t get as much of an influence over the constructed probability distribution — a spatial bottleneck.

Attention solves for these bottlenecks because it looks at every single token and compares it with every other token it has seen thus far. What does this mean? Let’s pretend that you are the model. You have been trained on a very large corpus of data, and by using gradient descent, you have gotten much better at predicting what the actual next word will be, given some amount of text. Now, you are running inference (meaning rather than getting trained and tweaking your specific processing methodology, you are just answering a query using that methodology — a given methodology is encoded in the numbers (called weights) that you use to compute something, so running inference just means using your pre-computed weights to compute an output based on a given input). I have asked you a question, and you have taken my text input, broken it up into tokens, and created a vector embedding of each of those tokens. Your job is now to pass the input I have given you through an attention layer.

In order to do so, you create three vectors from each token — Query, Key, and Value (Q, K, V) — by multiplying the vector embedding of each token by three different weight matrices, called $W_Q$, $W_K$, and $W_V$, all of which were learned during training. The resulting matrices perform the following functions:

• The Query vector, Q, asks a question on behalf of the token — I am this word, and encoded in me is some understanding of what I mean. Based on that understanding, what else do I need to know?
• The Key vector, K, uses the understanding of the token to advertise information about itself to answer all of the other Q matrices from every other token.
• The Value vector, V, encodes the information about the token itself that other tokens can learn from it.

In the world of vectors, what matters is closeness. In order to extract the closeness of what is essentially an ordered list of numbers, we use a dot product, which is the sum of the element wise multiplication of every number in two vectors. If the dot product is high, the two vectors are close. Thus, to determine which Ks answer which Qs well, we take the dot product between the Q of every token and the K of every other token (or in GPT-style models, every token that came before it). This produces a bunch of numbers which we squish and scale so that some are very high, most are very close to zero. These are our attention scores.

Intuitively, these dot products tell our token of choice (the token that produced Q) which of the other tokens are most important. As the model, you want to update the token you’re working on with all the information that is relevant to it. This allows you to improve your understanding of that token from being a word in isolation to a word that has some meaning informed by, for example, the adjectives directly before it. In order to do this, we add the Value of every preceding token to our existing vector embedding, weighting each V by the attention score we just computed. Thus, the tokens that deserve the most attention have the highest impact on your updated understanding of a given token.

A key aspect of attention is the fact that it is multi-headed, which is important for two reasons: one, it means that each “block” of attention calculations can index for a different type of relationship; and two, it means that each of these relationships can be calculated in parallel. To picture this, imagine that you have a vector embedding with 100 dimensions (meaning 100 numbers) and two attention heads. Each attention head compresses the entire embedding down into its own little workspace of 50 numbers. Let’s say, for example, one head attends to synonyms of the word and the other head attends to whatever is relevant to pronoun relationships. The first head of attention has its own three $W_Q$, $W_K$, and $W_V$ matrices and the second head has three different $W_Q$, $W_K$, and $W_V$ matrices. Within each head, these matrices have been tuned and tweaked to extract information specifically based on the relationship they are trying to understand and project it into their own workspace.

Thus, at the end of the updates from both attention blocks, the vector embedding for a given token contains not just updates influenced by words with similar semantic meanings, but also some information about syntax. GPT-3 has 96 attention heads in each attention layer, encoding relationships like dependencies, the scope of a negation (like what exactly a not might apply to), and the beginning and end of a sequence. That second piece about being able to compute these relationships in parallel is also essential; it means that the updates to these vectors are fast, scalable, and better utilize the chips they run on, which is a really interesting and relevant field of research that is dramatically out of the scope of this explanation. The point is, you want these things to be fast; think about how annoying it is when Claude or ChatGPT buffers while trying to compute the answer to your question.

In short, attention allows you, the model, to build a more nuanced understanding of a given token because you update that understanding based on all the information that came before it. Attention, however, for all its power, can only mix and average. Though it is not strictly linear in a mathematical sense, the fundamental mechanism can be extrapolated to a phenomenon that mimics linearity: the most important things get the most attention.

Parameters

Remember that you are a model. Let’s pretend that you are GPT-3. This means that you have 96 attention heads in each of 96 layers of attention (more on the purpose of layers later). Each of those attention heads has unique weight matrices to produce unique Q, K, and V.

GPT-3 specifically uses word embeddings that are 12,288 dimensions each, which means that in order to actually multiply each weight matrix by a word embedding, each of the weight matrices are 12,288 x 128 in dimension (the second number comes from 12,288/96, think of every attention head compressing the full 12,288-number embedding down to its own 128-number workspace that it gets to independently update). All of this, along with a final matrix per layer which stitches all the heads’ updates back together, contributes to a total of 58 billion parameters, 58 billion specific numbers that help to encode the understanding of a word. But you are GPT-3. You do not have 58 billion parameters, you have 175 billion parameters. What are the rest of those numbers doing?

A multi-layered approach

The output of each attention layer is passed through a small neural network, sometimes called a Feed Forward Network (FFN), sometimes called a Multi-Layer Perceptron (MLP). The remaining parameters, all 117 billion of them, construct these networks. In short, a FFN is a small neural network that processes each token independently. A crucial part of that processing occurs when the output of a FFN’s linear layer, which is just a matrix multiplication, is passed through an activation function, because that function adds nonlinearity.

After being passed through the attention layer, the token is updated by a lot of information from surrounding tokens. The FFN uses those updates to transform the given token in isolation, applying operations to it that cannot be accomplished with only linear updates. I find this tricky conceptually, so this is how I think about it. Imagine that every number in your vector corresponds to some potential feature about that token, like whether or not the given token is a proper noun or whether it is inside of a quoted phrase. The activation function acts somewhat like a gate, using the updates from the attention mechanism to amplify or subtract or combine features together.

It is this added complexity, these nonlinear transformations, that make using multiple layers of attention meaningful. Without the added nonlinearity, there would be no transformation of the existing information, no new meaning that is not directly contained by the surrounding text. It is nonlinearity that allows the model to encode complexity, to build an understanding that extends beyond the confines of context. Each time you process an input nonlinearly, you twist and mold and bend a representation of a word beyond the simple understanding that “if we are similar, we should influence each other”. This means that each subsequent layer works on a better shaped, more complex embedding. Attention then uses every more nuanced token to subsequently update every other more nuanced token — the underlying mechanism stays the same, but the updates get more nuanced.

On nonlinearity in life

Every year, around this time, a couple of my family friends or friends’ siblings or high schoolers on LinkedIn will reach out to me and ask me about my Stanford application. They’ll ask me what classes I took, what extracurriculars I did, what I wrote about in my essays. It’s an intricate investigation skirting around a thinly veiled, singular question: why you?

I have answers to most of the questions they ask me. I am liberal about sharing my essays, my extracurriculars, the awards I received as a high schooler. My great love in high school was my high school publication. To this day, my adviser — Julia Satterthwaite — remains one the most influential figures in my life. This was not just because she taught us to produce excellence from the moment we stepped foot in her classroom, not only because she was beyond dedicated to every aspect of her work, but also because of all the tiny, silent ways she believed in and uplifted me. But that is a different story for a different day.

The thing about journalism, though, is that I didn’t actually win any awards for it until after I got into Stanford. When I applied to Stanford, I’d written a couple of great articles, but I didn’t have any sort of extraordinary recognition for them. There was no evidence to any admissions officer that I stood out anymore than the thousands of other high school seniors that were the Editors-in-Chief of their high school newspaper. I went to a really competitive high school in the Bay Area. Kids my age had, at that point, already done research at Stanford and Harvard and MIT with renowned professors, already won national math olympiads and speech tournaments, had published books, had played instruments at an international level, had four or five summers worth of prestigious summer programs in comparison to my one.

I worked reasonably hard in high school. In part, this was because I have always loved hard work and gotten a lot of satisfaction from doing things well. And this was in part because, especially during the peak of the COVID-19 pandemic, I didn’t have a lot of friends. I had a big falling out with my high school friend group the week that schools closed in my sophomore year of high school. As a result, I was not invited to the group discord calls or the group Zoom meetings — to this day, I have still never played Among Us. I was 15, bored, and had approximately 2.5 friends. There was not much for me to do other than to read and write and draw (I must add here that one of my 2.5 friends was Riya Ranjan, who remains one of my best friends today, and is the singlehandedly most intelligent, humble, and driven person I have ever met. It would not be an exaggeration to say that I owe much of my Stanford acceptance purely to her. Our high school and college graduation photos are pictured).

As a result of my wholly unwanted friendless state, while my classmates were playing basketball at our local middle school every evening and going on drives to the beach over the weekend, I worked. Riya and I crafted, advertised, and taught a biotechnology curriculum to hundreds of high schoolers; first across California, and then, across the world (including Egypt and India). I spent hours and hours on Blender, painstakingly following an Argentianian man’s YouTube tutorial on how to make a physically realistic net for a science competition we did together.

When school started, by some mysterious luck, I was placed in the first iteration of every class (there was no AP Biology before my AP Biology period, no AP Calculus before my AP Calculus period — I was in the first Zoom meeting for every subject I was taking), so I finished every quiz and submitted every assignment before everyone else, which meant no shortcuts (a side effect of this was that the only reason many of my would-be friends texted me was to ask me for quiz or homework answers, and I, being 16 and friendless, obliged). I poured into my schoolwork, especially my humanities classes. The essays I wrote for my Honors American Literature class in high school often took me over thirty hours. I would dedicate two weeks out of the month purely to journalism. When I worked on pieces for my high school publication, I did nothing else.

There was no immediate reward in any of this. My teachers liked me and I got good grades, but everyone at my high school got good grades. But because I worked a lot, I learned a lot. I wanted each piece of work I produced to be wholly better than the previous, purely for my own satisfaction, so I read more and became more ambitious in what I pursued. Some of this, I’m sure, helped my acceptances into summer programs, which in turn helped my acceptances into university. But this essay is not a lineage on how every piece of effort I put into working in high school allowed me to get into Stanford; it is, in fact, the opposite. If I mapped out the benefits of each my pursuits in high school, I think it’s safe to say that a vast majority of that work was not even remotely associated with why I got into Stanford (for example, I spent a ton of time in high school drawing graphics on my iPad and getting really proficient at Photoshop, which has yet to serve me in the slightest).

But there was also beauty in the nonlinearity. By middle school, I already believed that the kids around me were more accomplished, intelligent, and better positioned than I was to get into a good school. As a result, I accepted from a really young age that geting into a top university was a fruitless goal. I worked hard because I wanted to, and also because I didn’t have much else to do. And then, 647 days after I lost a lot of my friends and wrote a bunch of articles about activism at my high school, after 647 days of being molded and shaped and transformed by the things I planned for as much as the things I didn’t, I got into Stanford.

I marvel at this nonlinearity because from my first day of undergrad, I began tuning myself into being an optimization machine. Every activity felt like a tradeoff; every class I took meant a class I would have to push off until later, every club I joined was an evening that was occupied by meetings instead of schoolwork, dinner with one set of potential friends meant passing on a different set. I wanted to engineer perfection. I wanted to ensure that every minute of my time would pay off in an impressive degree and wonderful friends and most importantly, in an exceptional career. I studied the subject that I believed had the highest ROI, I joined preprofessional clubs, I leetcoded, I took project-based classes, I recruited for better and better summer internships every year. I am lucky: much of my effort did pay off. Though I’ve gotten far better acquainted with rejection than perhaps I would have liked, my freshman year self’s hypothesis in linearity was actually largely proven: decide on a goal, work backwards, and reap the reward.

There is benefit in linearity. There was a lot of benefit in identifying what subjects and what skills mattered, in weighting the classes that made me smarter and the friends that made me better more heavily by devoting my more time and energy to them. There is benefit in attention.

But attention is perhaps not all you need. Many of the best parts of my undergrad were borne purely from nonlinearity; from reading every word of an essay or book when only one chapter was assigned just because I was interested; from joining the Bhangra team, not because I got better at dancing (I did not), but because I met three of the people who have influenced my values and goals most in undergrad from that singular group. Some of my favorite classes at Stanford were within computer science and math and equally as many were within philosophy and creative writing.

My most rewarding experience in undergrad was becoming a TA for our flagship computer science courses (a program known as section leading). A tremendous amount of investment goes into making new section leaders prepared to teach, and when I was attending multiple hours of weekly training at the beginning of my senior year, I seriously doubted whether or not the time investment was worthwhile. I didn’t see myself getting a job from the program, becoming a better programmer, or even making a bunch of friends. Furthermore, in order to keep up with the time commitment of being a section leader, I was going to have to drop a math class I wanted to take for my minor. Nevertheless, I saw it through, and for the most part, I was right: being a section leader has not led directly to an internship, nor has it made me a much better programmer. All it did was remind me how much I love to teach. That part, for me, has been pure joy.

To conclude

It is difficult, I suppose, to know when you are in the midst of a twist or a loop, an unwelcome transformation when all you would really like is a clear path forward, a clean victory. When I feel restless and uncertain, as I do now, I try to remember this, try to remember my 16 year old self, who did things for the joy of doing them and not much else.