Building Your Own Foundation Model (Part 4) : Self-Attention — How AI Decides What Matters

Building Your Own Foundation Model

Part 4 — Self-Attention

How AI Decides What Matters

The most important word in a sentence is not always the closest word.

Self-Attention is how a Transformer learns which pieces of information should influence each other.

---

⚡ Before You Continue...

Try this.

Read:

The animal didn't cross the road because it was tired.

What does "it" refer to?

The road?

Or the animal?

Humans resolve this almost instantly.

But a neural network doesn't receive the sentence as a human-readable sentence.

It receives vectors.

So here's the big question:

How does a Transformer figure out which words should pay attention to which other words?

That question leads us to one of the most important ideas in modern AI.

Self-Attention.

🎮 Attention Detective

You are now the Transformer.

Click a word.

Watch what happens.

The model doesn't simply look at the previous word.

It compares the current token against other tokens and calculates how strongly they should influence one another.

This is the beginning of Attention.

---

From Position to Attention

In Part 3 we solved one problem:

How does the model know where a token is?

Now we have another:

How does the model know which tokens are relevant to each other?

Consider:

The cat drank the milk because it was thirsty.

The important relationship is:

it
 ↓
cat

The word "it" needs information from somewhere else in the sequence.

Self-Attention provides a mechanism for making that connection.

---

The Big Idea

Imagine every token asking:

"Who should I pay attention to?"

For every token, the Transformer creates three representations:

Q = Query
K = Key
V = Value

Don't worry about the mathematics yet.

Think of them as:

Query = What am I looking for?

Key = What information do I contain?

Value = What information should I provide?

That simple mental model is enough to start.

---

🧠 A Human Analogy

Imagine entering a library.

You ask:

"Where can I find books about neural networks?"

Your question is the:

QUERY

Every book has a description:

KEY

Once you find the relevant books, you read their contents:

VALUE

So:

Query
   ↓
Find matching Keys
   ↓
Retrieve Values

That's the basic intuition behind attention.

---

The Three Characters

🔎 Query

The Query represents what the current token is looking for.

For example:

"it"

might effectively ask:

"Which previous token helps me understand what I refer to?"

---

🏷 Key

Every token has a Key.

The Key describes what kind of information that token can provide.

For example:

cat → animal/entity
road → location/object
tired → state

---

📦 Value

The Value contains the information that can actually be passed forward when a token receives attention.

So the simplified process becomes:

Query
   ↓
Compare with Keys
   ↓
Calculate Scores
   ↓
Convert Scores to Weights
   ↓
Combine Values

---

The First Mathematical Step

Now we can introduce the first important equation.

Attention begins by comparing:

Query × Key

More precisely:

Q × Kᵀ

This produces an attention score.

A higher score means:

"These two tokens are more relevant to each other."

---

🎮 PLAYGROUND 2 — Attention Heatmap

🔥 Watch Attention Move

The heatmap should immediately show something fascinating.

Instead of reading the sentence linearly, you can visualize relationships between tokens.

For example:

             The Cat drank the milk because it was thirsty

it            ░░  ███  ░░   ░░    ░░      ░░   ░░   ░░

The brighter the cell:

Higher attention

The darker the cell:

Lower attention

This is one of the most important visualizations in the entire series.

---

Why Is It Called "Self"-Attention?

Because the sequence attends to itself.

The tokens are both:

The information source

and

The information seeker

For example:

The → looks at → cat
cat → looks at → drank
drank → looks at → milk
it → looks at → cat

The sequence builds relationships internally.

That's why it is called:

Self-Attention

---

The Complete Attention Formula

Now we can reveal the famous equation.

Attention(Q,K,V)
=
softmax(QKᵀ / √dₖ)V

Don't memorize it yet.

Let's break it apart.

---

Step 1 — QKᵀ

Q × Kᵀ

This compares Queries against Keys.

Result:

Attention Scores

---

Step 2 — Scale

Why divide by:

√dₖ

Because as vector dimensions become larger, raw dot products can become large.

Large values can make the softmax distribution extremely sharp.

Scaling keeps the values in a more useful range.

---

Step 3 — Softmax

Now we convert scores into probabilities.

Imagine:

cat       0.70
road      0.05
animal    0.20
tired     0.05

The values add up to approximately:

1.0

Now the model has a weighted view of the sequence.

---

Step 4 — Multiply by V

The attention weights determine how much information to take from each Value.

Attention Weights
        ×
      Values
        ↓
New Representation

This becomes the information passed deeper into the Transformer.

---

🎮 PLAYGROUND 3 — Build Attention Yourself

Try It Yourself

Start with:

Tokens:

Cat
Drank
Milk

The playground generates tiny vectors.

Change one value.

Watch:

Q
 ↓
K comparison
 ↓
Attention score
 ↓
Softmax
 ↓
Weighted Value

change in real time.

You have just implemented the central idea of Self-Attention.

---

A Tiny Example

Suppose:

Q = [1, 0]

and

K₁ = [1, 0]
K₂ = [0, 1]

The dot products are:

Q · K₁ = 1

Q · K₂ = 0

So the first Key is more relevant.

The model can therefore assign more attention to the first token.

This is simplified.

Real Transformers perform these operations across large matrices and many dimensions.

But the principle is the same.

---

Multi-Head Attention

Now comes the next surprise.

Transformers don't normally use just one attention operation.

They use multiple attention heads.

This is called:

Multi-Head Attention

Imagine eight people reading the same sentence.

One person focuses on:

Subject

Another:

Object

Another:

Grammar

Another:

Long-distance relationship

Another:

Context

The model can learn different relationship patterns in different heads.

---

🎮 PLAYGROUND 4 — Multi-Head Attention Explorer

One Sentence. Many Perspectives.

Imagine:

Head 1
"The cat"

Head 2
"drank → milk"

Head 3
"because → thirsty"

Head 4
"it → cat"

The actual behavior of attention heads is more complex than these labels.

But the analogy gives us the intuition:

Different attention heads can specialize in different patterns.

---

🚨 The Problem That Changes Everything

There is a problem.

If every token can attend to every other token...

How many relationships do we calculate?

For:

10 tokens

roughly:

10 × 10 = 100

For:

1,000 tokens

we get:

1,000 × 1,000
=
1,000,000

For:

100,000 tokens

we get:

10,000,000,000

That's 10 billion pairwise interactions before considering the rest of the model.

This is why attention becomes expensive as context grows.

---

📈 PLAYGROUND 5 — The Attention Cost Simulator

Try 10 Tokens.

Now:

100

Now:

1,000

Now:

10,000

Move the slider.

Watch the matrix explode.

This is the moment you should ask:

If attention is so powerful, why does it become so expensive?

That question leads directly to modern Transformer optimization.

---

Causal Attention

There is another fascinating trick.

When generating text, the model cannot look into the future.

Suppose we're generating:

The cat

The model can use:

The
cat

But it cannot use the token that comes after the answer.

Otherwise the model would be cheating.

So we apply a:

Causal Mask

---

🎮 PLAYGROUND 6 — Break the Model

The Future Is Invisible

For:

The cat sat on the mat

the token:

cat

can see:

The
cat

but not:

sat
on
the
mat

Visualize it as:

        The Cat Sat On The Mat

The     ✓   ✗   ✗  ✗  ✗  ✗
Cat     ✓   ✓   ✗  ✗  ✗  ✗
Sat     ✓   ✓   ✓  ✗  ✗  ✗
On      ✓   ✓   ✓  ✓  ✗  ✗
The     ✓   ✓   ✓  ✓  ✓  ✗
Mat     ✓   ✓   ✓  ✓  ✓  ✓

That triangle is one of the simplest and most important pictures in an autoregressive LLM.

---

💻 Build a Tiny Attention Layer

Now let's move from theory to engineering.

A simplified PyTorch implementation:

import torch
import torch.nn.functional as F

def attention(Q, K, V):
    d_k = Q.size(-1)

    scores = Q @ K.transpose(-2, -1)

    scores = scores / (d_k ** 0.5)

    weights = F.softmax(scores, dim=-1)

    output = weights @ V

    return output, weights

That's it.

The complete Transformer attention mechanism contains many additional details, but this tiny function captures the central mathematical operation.

---

🧪 Run the Algorithm

Walk the same pipeline as the tiny attention(Q, K, V) function. Edit a vector, then press Step or Run.

---

flowchart LR

    A["🔤 Token Embeddings"] --> B["🧩 Input Representations"]

    B --> C["🔎 Query Q"]
    B --> D["🏷 Key K"]
    B --> E["📦 Value V"]

    C --> F["Q × Kᵀ"]
    D --> F

    F --> G["📊 Attention Scores"]

    G --> H["÷ √dₖ"]

    H --> I["🌡 Softmax"]

    I --> J["🎯 Attention Weights"]

    J --> K["Weights × V"]
    E --> K

    K --> L["🧠 Attention Output"]

    classDef input fill:#2563eb,color:#fff,stroke:#1e40af,stroke-width:2px
    classDef qkv fill:#7c3aed,color:#fff,stroke:#5b21b6,stroke-width:2px
    classDef math fill:#f59e0b,color:#111827,stroke:#b45309,stroke-width:2px
    classDef output fill:#10b981,color:#fff,stroke:#047857,stroke-width:2px

    class A,B input
    class C,D,E qkv
    class F,G,H,I,J,K math
    class L output

---

The Transformer Block

Self-Attention is not the entire Transformer.

A simplified Transformer block looks like:

flowchart TD

    A["Input"] --> B["Multi-Head Self-Attention"]

    B --> C["Add & Normalize"]

    C --> D["Feed Forward Network"]

    D --> E["Add & Normalize"]

    E --> F["Output"]

    F --> G["Next Transformer Block"]

    classDef input fill:#2563eb,color:#fff,stroke:#1e40af
    classDef attention fill:#8b5cf6,color:#fff,stroke:#5b21b6
    classDef process fill:#f59e0b,color:#111827,stroke:#b45309
    classDef output fill:#10b981,color:#fff,stroke:#047857

    class A input
    class B attention
    class C,D,E process
    class F,G output

We will break down every box in future episodes.

---

🖼️ — Transformer Attention

---

🧠 The Most Important Mental Model

Don't memorize:

Q K V

Remember this:

QUERY
"What am I looking for?"

        ↓

KEY
"Do I contain something relevant?"

        ↓

VALUE
"Here's the information I can provide."

        ↓

ATTENTION
"How much should I use it?"

Once this becomes intuitive, Transformer architecture becomes dramatically easier to understand.

---

🎯 Become the Attention Mechanism

Let's play one final game.

Sentence:

The cat drank the milk because it was thirsty.

Your job:

What should "it" pay attention to?

A. The
B. cat
C. milk
D. thirsty

Don't look at the answer yet.

Think.

...

...

...

The interesting part isn't simply getting the answer.

The interesting part is:

How could a mathematical system discover the relationship without anyone explicitly programming "it = cat"?

That's the magic of learned representations and attention.

---

🔬 Research Frontier: What Comes After Transformers?

Transformers dominate modern language modeling.

But researchers are asking a dangerous question:

Do we really need to perform dense numerical computation for every token?

What if a neural network only computed when something important happened?

That idea leads us toward:

Spiking Neural Networks

---

⚡ What Is a Spiking Neural Network?

Traditional neural networks usually produce continuous numerical activations.

A simplified neuron might output:

0.731

A spiking neuron behaves more like:

No spike
No spike
SPIKE
No spike
SPIKE

Information is represented through events over time.

Think:

Traditional AI

████████████████
continuous numbers

Spiking AI

· · ⚡ · · ⚡ · · · ⚡
events

---

Why Is This Interesting?

Brains do not continuously perform enormous dense matrix multiplications for every possible connection.

Neurons communicate through sparse events.

That makes researchers interested in:

• energy efficiency
• event-driven computation
• neuromorphic hardware
• low-power inference
• edge AI
• always-on systems

IBM's NorthPole research is an important example of brain-inspired AI hardware: its architecture places memory close to computation, targeting the data-movement bottleneck that can dominate AI workloads. IBM reported sub-millisecond/token latency for a 3-billion-parameter LLM in a specific 16-chip research configuration, along with substantially higher energy efficiency than the GPUs used in its comparison. These are research results for a specialized prototype, not evidence that neuromorphic hardware has replaced GPUs for general LLM workloads.

---

🧬 Can We Combine Attention + Spikes?

Researchers are already exploring this.

For example, Spikformer combines spiking neural networks with Transformer-style attention ideas, while SpikeGPT explored generative language modeling using spiking neurons.

More recent research is pushing toward SNN-based LLM inference. A 2025 preprint, NeurTransformer, explored converting GPT-2-style models toward spike-based inference and reported estimated energy reductions for its attention mechanism on digital hardware. The authors also highlight the trade-offs: conversion can lose accuracy, and spike time steps can increase latency.

This is exactly why this field is interesting.

We don't yet have a simple answer:

GPU Transformer
        vs
Spiking Transformer

The real research question is:

Can we redesign the algorithm AND hardware together so that AI only spends energy when computation is actually useful?

---

🎮 RESEARCH PLAYGROUND — Dense vs Spiking

Try It

Set:

Tokens: 100

Dense mode:

Many interactions

Spiking mode:

Only selected events

Now increase:

100
→
1,000
→
10,000
→
100,000

Watch the difference in activity, not just token count.

This playground should teach an important research principle:

Fewer operations do not automatically mean better AI.

We also need:

Accuracy
Latency
Memory
Energy
Hardware compatibility
Training stability

That is the real engineering problem.

---

⚠️ Don't Fall for the "Brain = Free AI" Myth

Spiking Neural Networks are exciting.

But they are not magic.

Current challenges include:

• difficult training
• surrogate-gradient methods
• temporal credit assignment
• accuracy trade-offs
• hardware availability
• software ecosystem maturity
• conversion overhead
• latency from multiple spike time steps

SpikeGPT itself described SNN language generation as an early research direction and reported results on relatively small 45M and 216M parameter models—not frontier-scale LLMs.

So the right conclusion is not:

"SNNs will replace Transformers."

The better conclusion is:

SNNs are an important research direction for making future AI more computationally and energetically efficient.

---

🌌 And What About Quantum Computing?

This question becomes even more interesting.

Could:

Transformer
+
Quantum Computing

produce faster AI?

Research into Quantum Transformers is actively exploring approaches such as quantum parameterized circuits, quantum attention mechanisms, and quantum linear-algebra techniques. But the field is still experimental: current research highlights challenges including scalability, benchmarking, training difficulties, and hardware limitations.

So don't expect:

Quantum Computer
        ↓
Instant ChatGPT

Instead, think:

Classical AI
      +
Quantum Algorithms
      +
Specialized Hardware
      ↓
Potential Future Acceleration

We will return to this later in the series when we have learned enough mathematics to understand where quantum computing could actually fit.

---

🧪 Researcher's Question

Here's the question I want you to keep in your mind:

What if the next generation of AI isn't simply a bigger Transformer?

Maybe it becomes:

Transformer
      +
Sparse Attention
      +
Mixture of Experts
      +
Spiking Computation
      +
New Memory Architecture
      +
Neuromorphic Hardware
      +
Quantum Algorithms

The future may not be about one architecture winning.

It may be about co-designing algorithms, memory, hardware and learning systems together.

---

🧩 What We Have Built So Far

We can now see the complete journey:

flowchart LR

    A["👤 Human Prompt"]
    --> B["🔤 Tokenization"]

    B --> C["🔢 Token IDs"]

    C --> D["🧩 Embeddings"]

    D --> E["📍 Positional Information"]

    E --> F["🔎 Query"]

    E --> G["🏷 Key"]

    E --> H["📦 Value"]

    F --> I["Q × Kᵀ"]
    G --> I

    I --> J["📊 Attention Scores"]

    J --> K["🌡 Softmax"]

    K --> L["🎯 Attention Weights"]

    L --> M["Weights × V"]
    H --> M

    M --> N["🧠 Transformer Block"]

    N --> O["➡️ Next Block"]

    O --> P["📈 Next Token Probability"]

    classDef user fill:#2563eb,color:#fff,stroke:#1e40af,stroke-width:2px
    classDef token fill:#10b981,color:#fff,stroke:#047857,stroke-width:2px
    classDef vector fill:#7c3aed,color:#fff,stroke:#5b21b6,stroke-width:2px
    classDef attention fill:#f59e0b,color:#111827,stroke:#b45309,stroke-width:2px
    classDef output fill:#ec4899,color:#fff,stroke:#9d174d,stroke-width:2px

    class A user
    class B,C token
    class D,E vector
    class F,G,H,I,J,K,L,M,N attention
    class O,P output

---

🚀 You Just Crossed an Important Line

Until Part 3, you were learning what goes into a Transformer.

Now you understand:

How tokens interact.

That's a major milestone.

But we have only opened the door.

Because Self-Attention creates the relationships...

Something else must transform those relationships into increasingly powerful representations.

That component is:

The Feed-Forward Network.

And there is a surprising question waiting for us:

If Attention decides what information matters, what actually transforms that information into something useful?

---

🔥 Part 5 — The Neural Network Inside the Transformer

Next we'll open the Transformer block.

You'll see:

Attention
   ↓
Residual Connection
   ↓
Layer Normalization
   ↓
Feed Forward Network
   ↓
Activation Function
   ↓
Residual Connection
   ↓
Normalization

Then we'll build a tiny Transformer block ourselves.

And for the first time in this series, we'll take:

"The cat sat."

and push it through an actual miniature model.

Not a diagram.

Not an animation.

A working neural network.

---

🧠 Your Challenge Before Part 5

Don't continue yet.

Try to answer these five questions:

1. Why do we need Query, Key and Value?
2. What does QKᵀ calculate?
3. Why do we divide by √dₖ?
4. Why do we use Softmax?
5. Why can't an autoregressive model look at future tokens?

If you can answer these without looking back...

you understand the core idea of Self-Attention.

If not, go back to the playgrounds.

Change the numbers.

Break the attention matrix.

Try again.

That's how you learn this—not by memorizing the equation.

---

📚 Research References

This chapter builds on the Transformer architecture introduced in Attention Is All You Need, which replaced recurrence/convolution with attention mechanisms and made the architecture substantially more parallelizable.

For positional information discussed in the previous episode, RoFormer introduced Rotary Position Embeddings (RoPE), which encode position through rotations and expose relative-position information inside attention.

For the research frontier discussed here, see work on Spikformer, SpikeGPT, and more recent SNN-based Transformer/LLM inference research.

---

📈 Series Roadmap

Follow the winding road from Part 1 through designing your own foundation model. This episode is the highlighted pin.

---

The Goal

By the end of this series, you shouldn't merely be able to say:

"I know how ChatGPT works."

You should be able to say:

"I can build a small language model, understand why each component exists, measure its performance, optimize its inference, experiment with alternative architectures, and reason about where future computing paradigms such as neuromorphic and quantum computing might fit."

That's the difference between using AI and engineering AI.