Recently, I found myself becoming a “script kiddie” by following HuggingFace tutorials to try and post-train models. With the advent of coding agents, it has unfortunately become easier than ever to implement something without understanding what is happening under the hood. I am writing these (manually typed) notes to further my understanding.

A quick digression on the Occam’s Razor

While we may not be able to perfectly model the behavior and state of a chaotic system, we should be able to define the set of constraints that govern the system. And, almost all of the chaos in the universe is governed by a set of simple, symetrical laws.

I firmly believe (albeit without proof) that the same must apply to the building blocks of life, consciousness, and intelligence (biological or otherwise). As an extension to the Bitter Lesson, I would bet that the algorithm that would eventually create general intelligence must be simple, easy to understand and easy to implement. Throughout these notes, I will only be discussing the algorithms that follow these principles.

The Elegance of Neural Attention

Fundamentals

Quoting Wikipedia:

Attention is the concentration of awareness directed at some task or phenomenon while mostly excluding others.

Given a collection of input signals, how can a neural network learn to direct its focus on the relevant signals?

Consider that we represent signals in an abstract vector space with $N$ dimensions:

To figure out how related these signals are, we can compute their similarity.

While we could use cosine similarity (calculated as $\frac{S_1 \cdot S_2}{|S_1| |S_2|}$) neural networks typically use the raw dot product for efficiency. Geometrically, the dot product measures how much two vectors align or point in the same direction, essentially capturing the overlapping “footprint” of one signal projected onto another.

Given a set of signals $S_1, S_2, \dots, S_M$, we can efficiently compute the pairwise attention scores between all the signals using matrix multiplication.

If we stack the $M$ signals into an $M \times N$ matrix called $S$, we can compute:

Let’s build an intuition of what this matrix $D$ signifies. $D$ is an $M \times M$ matrix, and each element $D_{ij}$ represents how similar $S_i$ and $S_j$ are. Ultimately, it tells us how much “attention” signal $S_i$ must pay to signal $S_j$.

Next, we normalize these attention scores by performing a softmax operation so that the sum of the attention scores for any particular vector is exactly 1.

Notice that we scale the $D$ matrix by dividing it by $\sqrt{N}$ (the square root of the dimensions). We do this because large dot products can push the softmax function into unstable regions with extremely low gradients, which stalls the learning process.

We can then pull in the contributions of each signal $S_j$ into $S_i$ by computing a new composite signal, $O_i$, which is a weighted sum of all signals based on their attention scores:

Using matrix multiplication once again to calculate this for all signals simultaneously, we arrive at our final output matrix:

If the $S$ matrix represents signals in isolation, and the $A$ matrix represents the instructions on who should listen to whom, the $O$ matrix represents the fully contextualized signals. Each row of O represents a contextualized version of the corresponding row in S.

Elegant, isn’t it?

Queries and Keys

Until now, a signal was comparing itself directly to other raw signals using $S S^T$. There is a teeny tiny issue with this approach: It assumes that “what a word is” is exactly the same as “what a word is looking for.”

If we compared raw signal values, we would always end up with the highest Dot Product for the same signal. Take the following phrase:

The heavy box fell on the fragile table, and it crushed it

If we compared the raw signal values of the second “it”, it would show highest similarity to the first “it”. However, we want each “it” to attend to a corresponding noun (“box”, “table”).

To make the network truly powerful, we want it to actively learn how to direct its attention. We do this by projecting our raw signals into a query dimentsion using learnable weight matrices: $W_Q$ and $W_K$.

When we multiply our input matrix $S$ by these weights, we generate new matrices:

These stand for Queries and Keys. The easiest way to build intuition for this is through a retrieval system:

• The Query ($Q$): This is what a signal is looking for.
• The Key ($K$): This is what the signal represents.

If we do this for our original sentence, the Query for the “it” might map it to a latent vector corresponding to noun. At the same time, Key vector will also map the nouns in the sentence to the same region. This makes the Attention mechanism serve more like a latent retreival than a raw similarity comparison.

Rewriting some of our equations from before:

(Note: We scale by $\sqrt{d_k}$, which represents the dimension of the Key vectors, ensuring our softmax gradients remain stable).

Another advantage of the projection is that they allow us to compute the Attention Score in a latent space with a different (lower) cardinality than the original Signals. This can make our computation faster!

Values

Currently, we calculated the Output matrix by multiplying the Attention Scores with the raw signals. While this would work in theory, we can do better. Consider the following sentence:

She matched her new paint to the sweet apple she was eating.

Assuming the attention scores between “paint” and “apple” are high, what would we want “apple” to contribute to “paint”? If our Key vector mapped “apple” to “color”, we would want that aspect of “apple” to contribute to the Output for “paint”. That is we would want a corresponding mapping from “apple” to “red”.

We achieve this corresponding mapping using another set of trainable weights $W_V$ and compute the Value Matrix ($V$) as

Substituting $Q$, $K$, and $V$ into our previous formula, we arrive at:

(Note: $V$ needs to be the same dimensions as $S$ to satisfy the equation above, but this is not always a requirement).

By separating the inputs into Queries, Keys, and Values, the network is no longer just calculating blind similarity. It has learnable parameters that allow it to actively ask questions, advertise its own contents, and selectively share information across the sequence.

This exact mechanism, formally known as Scaled Dot-Product Attention, was introduced in the now ubiquitous “Attention Is All You Need” paper. By allowing networks to efficiently weigh the importance of different inputs, this mathematical operation became the foundation of the modern Transformer architecture.

Multi Head Attention

Consider the same sentence as before

She matched her new paint to the sweet apple she was eating.

The Scaled Dot-Product Attention we described is perfectly capable of computing the Contextualized Output for the word “paint” by computing a vector corresponding to the color “red”.

But a single word can convey more information! In the same sentence, what if we want to compute the output score for the word “eating”? In this case, we can imagine that the Key for “apple” must correspond to something like “food” and the corresponding value must represent the “apple” itself.

We achieve this by having different sets of trainable weights $W_Q^i$, $W_K^i$, $W_V^i$ that “specialize” in extracting a different aspect of the input signal.

Subsituting into the earlier equations,

$\text{head}_i$ is the output of a single Attention Head.

Once all the heads have done their specialized extractions, their outputs are concatenated (glued together) and passed through one final weight matrix ($W_O$) to blend the insights of the entire committee back together into a single, incredibly rich contextual signal.

Usually, to keep the amount of compute roughly constant, we make the dimension of each head:

(Note: Recall that $V_i$ need not be the same dimension as $S$)

By using multiple heads, the network doesn’t have to compromise. It can look at the word “apple” from 96 different perspectives simultaneously, extracting every nuance needed to understand the full picture.

The gemma-2-2b-it model achieves a 56% score (5-shot)¹ on the MMLU. What would we expect to happen if we asked it a batch of N questions at a time? In the ideal case, each of the questions would be answered correctly 56% of the time and this probability would simply be a independent variable that depends on the “intelligence” of the model.

However, this is not what we see!

When given a list of questions, the model gets the last question correct twice as often as the first question. The first few questions are pretty much guessed correctly with a random chance (25%).

Note: I ran these experiments using gemma-2-2b-it due to hardware limitations. Please try to run them on bigger models if you have the resources - I would love to see the results :)

Example Prompt

Read the following multiple choice questions and answer them one by one. Your answer for each question should be in following format: . <the index of the correct option (1, 2, 3, or 4)>

1. Question: A family therapist using the structural approach of Salvador Minuchin would most likely:
   Choices:
   1. clarify boundaries between family members in order to reduce enmeshment.
   2. work initially with the most differentiated family member.
   3. “use a multiple-therapist team to prevent any one therapist from becoming ““triangulated”” into the family system.”
   4. “issue specific ““directives”” designed to counteract dysfunctional processes.”

2. Question: Which of these countries was not a member of the Axis alliance during World War II?
   Choices:
   1. Germany
   2. Italy
   3. Spain
   4. Japan

(..the other questions, in a similar format)

Here are the answers to the multiple-choice questions, one index per line:

Pseudocode

The evaluation loop is fairly straightforward:

N_q = 50      # number of questions in a batch
N_iters = 100 # number of iterations

scores = [0 for i in 1..N_iters]
for i in 1..N_iters:
    # batch is an array of size N_q containing (Q, A) tuples
    batch = get_random_batch()

    # concatenate questions and add instructions for the model
    prompt = format_prompt(batch)

    # generate output from the LLM
    out = generate(prompt)

    # returns a boolean array
    is_correct = check_answers(batch, out)
    
    # scores[i] += 1 if is_correct[i]
    update_scores(scores, is_correct)

Why might this be happening?

I believe that there are two optimizations here that are bottlenecking LLM performance.

1. Local Attention

Standard self-attention is powerful, letting every token “attend” to every other token in the context. However, this global view becomes computationally expensive (O(n^2)) as context length grows. Enter Local Attention². It’s an efficiency optimization technique. Instead of looking everywhere, each token only attends to a fixed-size window of its neighboring tokens. For decoder-only transformers like current LLMs, each token attends to a fixed-size window of preceeding tokens. This drastically reduces the computational cost (O(n)) and memory required, especially for long documents.

Gemma2 alternates between a 4096-token sliding window local attention layer and a global attention layer (8192 tokens, since that is the full context window). This interleaving allows the model to attend to the global state while being slightly computationally efficient.

This means that Gemma2 should be able to answer all the questions within the 4096-token window with the same accuracy. This seems to hold for the most part - when batches of <20 questions are supplied, Gemma2 answers each question correctly with roughly the same probability.

However, I would imagine the graph to look more like a step-function where questions within the 4096 context window are answered with probability P and the others are answered with a lower probability P’. This doesn’t seem to be happening - something else is causing the ~linear decrease in probability with distance within the 4096 sliding window.

2. RoPE: Rotary Positional Embedding

The attention mechanism is order-invariant, so it can’t distinguish “boy eats sandwich” from “sandwich eats boy” without positional information. This is generally performed by adding a positional embedding to the embedded tokens. Rotary Positional Embedding (RoPE)³ is used in pretty much all the SOTA open-source models including Gemma3, Llama4, and Qwen3. It is also used in Gemma2 - the model we’re testing. I found this youtube video by Efficient NLP a great explanation of the different positional embeddings used in transformer architectures, including RoPE. RoPE also has advantages like context-window expansion, which makes it attractive for LLMs.

RoPE is applies a rotation to the Query (Q) and Key (K) vectors within the self-attention mechanism. The angle of rotation is determined by the token’s absolute position. The crucial insight is that when these rotated Q and K vectors interact via the dot product during attention calculation, the result naturally depends only on their relative distance.

However, since the attention between two tokens in a sequence is scaled by their dot product, RoPE causes a decay in attention between distant tokens. This is explained in the original paper that introduces RoPE as well.

For a simpler derivation, consider that tokens are represented in a 2d embedding space. To apply positional embedding, the token embeddings are rotated by some angle \(\theta = i\delta\) where \(i\) denotes the positon of the token in the sequence. It is clear that the dot-product between two similar embeddings spaced N tokens apart will scale by approximately \(\cos(N\delta)\).

This \(\cos(N\delta)\) decrease might explain the decrease within the 4096 sliding window.

What does this mean for LLM performance?

A core assumption behind both of these optimizations is that most relevant context for predicting a token lies nearby. I imagine this is true for learning linguistic representations while pre-training LLMs to do next-token prediction, but does it still apply for the agentic tasks that we want these models to perform? I don’t think so. Agentic tasks generally do not differentiate between different parts of the context.

More recent architectures like Gemma3⁴ have both decreased the sliding token window to 1024 and decreased the frequency of the global attention layers to 1 in 5. While the Gemma3 paper reports no loss in perplexity due to this, I suspect the model will perform worse on this test. I wasn’t able to test on gemma3-1b-it because it wasn’t able to format the answers correctly according to the prompt (even when supplying just 1-2 questions) and the 4b+ models were too big to run.

While the open source world has not found any new alternatives to this issue, large labs might have. The Gemini models especially seem to support a 2M context window while simultaneously performing well on the needle in the haystack test. Regardless, this is certainly an important bottleneck that needs to be solved on the path to AGI.