# A Deep Dive into Modern Transformers: Optimizing Attention

#85 Jan 19 Inference

https://www.aibutsimple.com/p/a-deep-dive-into-modern-transformers-optimizing-attention

Artificial intelligence has been around since the 70s, but a team of researchers at Google published a paper in 2017 that transformed the landscape: [Attention Is All You Need](https://arxiv.org/abs/1706.03762?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=a-deep-dive-into-modern-transformers-optimizing-attention).

With the adoption of the transformer architecture came the concept of self-attention, a strategy of predicting tokens that are fed into the transformer.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/f60a6613-b559-4ab7-a733-c5ac7390e9a5/Copy_of_Transformer_Tokens.png)

The main idea of these transformers is very widely studied. Feed in a chain of tokens, process them using attention, then obtain a probability distribution over the next possible tokens to sample from.

However, modern transformers used in LLMs like GPT-5, Gemini 3 Pro, and DeepSeek R1 have undergone significant updates and optimizations. Specifically, the attention mechanism has been heavily updated.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/c56f5ebd-de6f-4b05-9134-e0a8c1378caa/Modern_Transformer_Architectures__2_.png)

In this issue, we’ll dive deep into attention and its modern variants, showing exactly how they improve performance.

### A Review of “Regular” Attention

Let’s trace the way that self-attention may be calculated and utilized. It can be split up into two different parts: the training phase and the generation phase (inference).

#### The Training Phase

Let’s go over the training process of a very basic transformer. We first select a corpus (a body of text, e.g., Wikipedia) to train the model.

The model will simultaneously create vector embeddings (high-dimensional numerical representations) for every word in the corpus, which will initially start out as random numbers.

We will initialize weight matrices, Wq, Wk, and Wv, that all start as random values. This is how we’ll obtain the queries (Q), keys (K), and values (V). The transformer we’re about to train currently “knows” nothing about language.

At the beginning of the model’s training loop, the model will receive a sequence of tokens from somewhere in the corpus.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/6a0da37a-61ca-4e10-b9a9-5e5306496d9e/Transformer_Sentence_Example__1_.png)

The model’s task will be to predict each word, given the words that come before it. Looking at our visualization, it should predict “neighbor” for the token “A” and predict “needs” for “neighbor,” and so on.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/7b129196-5036-4965-a72d-18e808603c34/Transformer_Attention_Example__3_.png)

In one forward pass, the model will go through this process of prediction independently for every word, which will happen in parallel. Read our previous issue on [GPU Optimization](https://www.aibutsimple.com/p/gpu-optimization-for-machine-learning-simply-explained?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=a-deep-dive-into-modern-transformers-optimizing-attention) to learn more about this.

The forward pass can be described using a series of steps:

Step 1: Each word will retrieve its respective vector embedding to start. Remember, these are all random initializations right now.

Step 2: Each of these embeddings is multiplied by the three weight matrices from earlier.

- Wq is the weight matrix to compute query vectors.
- Wk computes key vectors.
- Wv computes value vectors.

In essence, this step will produce queries, keys, and values (vectors/matrices through matrix multiplication) for each word. Let’s assume for this case the queries, keys, and values are vectors.

Step 3: Each word will now compute attention scores to every word that comes before it (a process called causal masking will prevent words from “peeking” and looking ahead).

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/27c06db6-030f-4ef3-942c-687a7990f52d/Causal_Masking_and_Attention_Scores.png)

Each of these scores will then be passed into the softmax function, which normalizes the scores to between 0 and 1 (0 to correspond to the lowest and 1 for the highest).

Think of this as each word expressing the importance of every other word, in its own perspective.

Step 4: With the value vectors we calculated earlier, using Wv, we are able to produce a weighted sum by multiplying our softmax outputs with their respective value vectors. Taking the sum will produce our final output vectors for each word’s calculation.

Step 5: We will again multiply this output vector by Woutput, which will essentially push (project) this vector back into the corpus word space, which will give us values to associate for each word in our original corpus. Each vector entry will correspond to a word.

Step 6: Apply the softmax to the entries along the vectors once again. This time, the normalization will produce values between 0 and 1, which we treat as probabilities.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/1dcc6d0f-e2dd-4852-9fe8-221ba68ef9b9/Transformer_Token_Probabilities.png)

Step 7: The model’s predictions for each word’s successor are then evaluated against the real next words. With the weights being random, we can expect the model to be quite wrong initially.

Step 8: Using backpropagation, the model will essentially ask itself, “How can I adjust Wq, Wk, Wv, and Woutput (embedding space) to reduce error, or in other words, loss?”

These steps are then repeated millions, if not billions, of times. This should clearly highlight that training a transformer is a compute-intensive task and largely inefficient if a GPU, enabling parallelism, is not used.

However, this parallelism comes with problems. To compute the standard attention pattern between tokens, there is a runtime complexity of O(n2).

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/6ce16fae-16db-4fa1-91e4-64f9b7a45bbe/Transformer_Attention_Grid.png)

In our example, computing attention for 7 words requires 49 computations. For our entire corpus, say 10,000 words, that is 100 million computations. Quadratic scaling is the primary bottleneck that modern optimizations have been developed to address.

---

### [Level Up Your Career With AI, But Simple PRO](https://www.aibutsimple.com/upgrade?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=a-deep-dive-into-modern-transformers-optimizing-attention)

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/281697a8-c915-435c-a68b-43374b671698/aibutsimple_academy.png)

Supporting us by purchasing a PRO membership will get you:

- Exclusive Hands-on Tutorials With Code
- Full Math Explanations and Beautiful Visualizations
- Ad-Free Reading for an Uninterrupted Learning Experience

[Break into AI or sharpen your edge by joining today!](https://www.aibutsimple.com/upgrade?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=a-deep-dive-into-modern-transformers-optimizing-attention)

---

#### The Generating (Inference) Process

Now that training is complete, the model’s embeddings and weights are no longer random.

The model’s embeddings can capture semantic meaning, and the weight matrices have proper values that relate tokens by attention. At this stage, the model is no longer focused on learning/updating its weights. Now, we would like to generate useful outputs.

In the generation phase (also called inference), we can pass any phrase through. The difference between the learning and the generation phase is that the generation phase is unbounded and can “think on its own.”

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/058a209f-1a17-4afe-be33-31718811b5b2/Transformer_Inference_vs_Training.png)

The steps of generation are similar to the training steps:

Step 1: Attention scores are calculated for each word, using the process we described earlier, and an output vector is created, passed through softmax, and a word is selected.

Step 2: Instead of going and repeating the process, like we would have previously, the model will now append the word it just created and redo the entire process.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/7364faa1-68fb-48d6-9196-dae54833a061/Transformer_Inference_Example.png)

To speed up the attention computations, modern transformers employ a technique called [KV Caching](https://www.aibutsimple.com/p/llm-kv-caching-simply-explained?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=a-deep-dive-into-modern-transformers-optimizing-attention), which stores previous attention calculations.

This means that for each new word, it will only have to compute the attention calculation for only that word and all others, while retaining the scores for every other word and all of its targets.

Step 3: The model will continue generating new words until it predicts a sentence stop or the user stops generation.

The primary point of interest in optimizing the generation process is in how generation tackles the process of calculating attention.

Instead of being able to parallelize these steps of predicting new words for every word, the model must go sequentially, generating each word before starting a new calculation.

### Attention Optimizations

Now that we have reviewed standard attention, we can dive into key optimization techniques.

In both training and inference, we are mostly concerned with the bottleneck caused by the O(n2) complexity and the large memory footprint of regular attention.

#### Sliding Window Attention

As mentioned before, the problem lies in compute complexity. Certain databases, books, and even codebases can exceed tens of thousands of tokens, making this an important issue to address.

Let’s look back at our process of computing initial attention scores. We can see that certain words get very minuscule values compared to other words.

A trend that you may notice is that proximity matters. Researchers noted this and speculated that attention is local to an extent. In understanding this, sliding window attention (introduced in the 2020 [Longformer](https://arxiv.org/abs/2004.05150?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=a-deep-dive-into-modern-transformers-optimizing-attention) paper) addresses the observation that nearby context matters most.

So, when the model computes attention for the entire corpus, we can set window sizes for each pass.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/554e9212-fa20-49cf-a66d-672d0569713b/Sliding_Window_Attention__1_.png)

The commonly used number for this is a window size of 4096 tokens, where each token only attends to the previous 4096 tokens.

This reduces the time complexity from O(n2 ) to O(4096*n) complexity, as now it depends on the window size and the input size rather than solely on the input size.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/7169eed0-081f-4591-b004-624f16b36826/Copy_of_Sliding_Window_Attention.png)

As you can expect, the drawback of this method is that we lose the full long-range dependencies, where tokens outside of the window will be voided in terms of importance.

However, since the sliding window incorporates adjacent tokens, over many iterations, attention can slowly transfer from window to window, maintaining decent long-range context.

We can also introduce “global” tokens, which can attend to all windows, learned through training. This way, important tokens can stay relevant for long sequences.

Sliding window attention was popularized by Mistral’s 7B model, which led to widespread adoption in modern LLMs.

#### FlashAttention

Attention calculations in the generation phase are limited by how much memory is available. The problem here is the transfer of data between different types of GPU memory.

Standard attention is normally written and read from the global memory, which is further away from the GPU than its shared memory. If you read our previous issue on [GPU Optimization](https://www.aibutsimple.com/p/gpu-optimization-for-machine-learning-simply-explained?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=a-deep-dive-into-modern-transformers-optimizing-attention), there’s a useful strategy: tiling.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/db923ddb-c6a6-4bac-9e2b-f72069dc47cb/GPU_optimization_ML.png)

This is the exact strategy implemented in FlashAttention. The core idea is that instead of storing the n-by-n matrix of values to keep building on, we only keep a certain size.

In the same way as utilizing specific ranges of tokens, we can think of tiling as a way to break up certain parts of our attention matrix.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/1ff8ae9e-11dd-482b-96c2-9ebca294ffc3/FlashAttention.png)

In FlashAttention, instead of writing new n-by-n matrices, the global memory would only keep a certain portion, while others are sent to faster shared memory (SRAM).

Observe that this does not decrease our overall complexity (O(n2 )), as we still compute more and more as n grows. The difference is in speed, where it no longer takes hundreds of GPU clock ticks to retrieve cached memory.

This has been adopted in the [PyTorch](https://pytorch.org/?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=a-deep-dive-into-modern-transformers-optimizing-attention) attention implementation and was used in developing GPT-4 and other modern LLMs.

#### Grouped-Query Attention (GQA)

Grouped-Query Attention (GQA) is also key for inference. While FlashAttention focuses more on speed, GQA is aimed at reduce memory usage, especially when used in combination with KV caching.

In a model with 32 attention heads, or collections of queries, keys, and value vectors, we must keep 32 of each kind across all layers.

In practice, with typical query dimensions of 512 and 40 processing layers, the storage capacity can balloon to upwards of many gigabytes for one conversation.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/a8dd6a6b-5763-4da6-83c4-fa6c55f682c1/Grouped_Query_Attention__GQA___1_.png)

The solution GQA provides is that many different queries will share the same key and value vector, which downsizes the amount of memory (specifically KV memory) being used.

Instead of 32 attention heads, imagine if every eight queries were handled by just one KV head. This decreases the amount of KV parameters by 8x.

This was actually the ratio used by Llama 2, which meant serving 8x more uses with the same GPU utilization and 8x longer chats with a minimal hit on quality.

### The Bottom Line

In this issue, we’ve explored the base process for calculating attention and three different optimizations that tackle scaling transformer deployments in the real world.

Sliding window attention changes how we compute, while FlashAttention and Grouped-Query Attention change what we store and how we store it.

These combined innovations have enabled the leap from the very first transformers to modern transformers, now able to handle many millions of tokens.

---

Here’s a special thanks to our [biggest supporters](https://www.aibutsimple.com/upgrade?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=a-deep-dive-into-modern-transformers-optimizing-attention):

- Sushant Waidande
- Sai Krishna Pashapu
- Bob Boncella
- Demet Yazici

---

If you enjoy our content, consider supporting us so we can keep doing what we do. Please share this with a friend!

Want to reach 50000+ ML engineers? Let’s work together:

[WORK WITH US 📈](https://www.passionfroot.me/aibutsimple?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=a-deep-dive-into-modern-transformers-optimizing-attention)

Feedback, inquiries? Send us an email at [edwin@aibutsimple.com](mailto:edwin@aibutsimple.com).

If you like, you can also donate to help our team push out better newsletters every week!

[Support us ☕](https://ko-fi.com/aibutsimple?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=a-deep-dive-into-modern-transformers-optimizing-attention)

---

That’s it for this week’s issue of AI, but simple. See you next week!

—AI, but simple team