Cloth folding is one of those tasks that seems trivially easy for humans but is brutally hard for robots. Fabric has effectively infinite degrees of freedom. It crumples, slides, and deforms unpredictably. You can’t plan a rigid trajectory the way you would for picking up a block — the cloth reshapes itself with every touch.

In this post, I’ll walk through how I fine-tuned SmolVLA — a compact 450M-parameter Vision-Language-Action model — on 50 teleoperated demonstrations to teach two robot arms to fold a cloth. I’ll compare it head-to-head against ACT (Action Chunking with Transformers), a popular lightweight imitation learning baseline, and share what worked, what didn’t, and what surprised me.

SmolVLA successfully initiates the fold, completes it tightly, and pushes the cloth aside — scoring 2.80/3.00.*

The Setup: Two Arms, Three Cameras, One Cloth

The bimanual SO-101 setup with three camera views

The hardware is deliberately simple and accessible:

-Two SO-101 robot arms — open-source, low-cost, 6-DOF manipulators (~$300 each)

- Three cameras— left wrist, right wrist, and an overhead webcam (all 480x640, synced at 30 FPS)

- Action space — 12 dimensions (6 joint positions per arm)

I collected 50 teleoperation demonstrations of myself folding a cloth, totaling ~48,000 synchronized frames stored in LeRobot v3.0 format. Each demonstration captures the full pipeline: reach for the cloth edge, fold it over, press down, push it aside.

Why 50? That’s enough to cover the natural variation in cloth placement and draping, while staying practical for a single-person data collection session (~2 hours).

How SmolVLA Works

SmolVLA is a Vision-Language-Action model — it takes in camera images and a language instruction, and directly outputs motor commands. Here’s what makes its architecture interesting:

SmolVLA architecture: visual tokens from SigLIP are compressed via PixelShuffle and fused with language tokens in a VLM backbone. The action expert generates chunked trajectories via flow matching.

Vision: See the Cloth

A SigLIP vision encoder processes each camera frame into tokens. The key trick is PixelShuffle compression — 1,024 visual tokens per frame get compressed down to just 64, a 94% reduction. With three cameras, that’s 192 visual tokens instead of 3,072. This makes the model tractable to train and run in real-time.

Language: Understand the Task

The instruction “Fold the cloth” gets tokenized and processed alongside the visual tokens through a SmolVLM2 language backbone. In a single-task setting like ours, this might seem unnecessary — but it means the same model architecture could handle “Fold the cloth in half” vs. “Roll the cloth up” without retraining the vision pipeline.

Action: Move the Arms

This is where SmolVLA diverges from standard VLMs. Instead of predicting text tokens, it uses a dedicated action expert that generates “action chunks” — sequences of future joint positions — using flow matching.

Flow matching training: noise is added to ground-truth trajectories, and the model learns to predict the velocity field that denoises back to the clean trajectory.

Think of flow matching like this: during training, you take a real trajectory and add noise to it. The model learns to predict the “velocity field” — the direction to push each noisy point to recover the clean trajectory. At inference, you start from pure noise and iteratively apply the learned velocity field to generate a smooth, coherent action sequence.

Why action chunks instead of single-step predictions? Because predicting one joint angle at a time compounds errors quickly — the robot drifts. By predicting a chunk of ~50 future timesteps at once, the motion stays coherent and smooth.

The action expert uses interleaved cross-attention (grounding in visual context) and self-attention (temporal coherence across the action chunk).

Training: Pretraining Is the Secret Weapon

Training pipeline — pretrained on 30K community episodes, then fine-tuned on 50 cloth-folding demonstrations.

SmolVLA’s base model (`lerobot/smolvla_base`) was pretrained on ~30,000 episodes from LeRobot community datasets — various robots doing various manipulation tasks. This broad exposure gives the model prior knowledge about how objects move, how grippers interact with surfaces, and how to coordinate motor commands.

Fine-tuning on our 50 cloth-folding episodes took ~2 hours 49 minutes on a single A100 GPU** for 10,000 steps. The final training loss was 0.021.

How much does pretraining matter? According to the SmolVLA paper, the gap is dramatic:

Without pretraining: 51.7% success. With pretraining: 78.3% success. A +26.6% absolute improvement from pretraining alone.

The model doesn’t just learn faster — it learns better, because it already has a foundation of sensorimotor knowledge to build on.

The Benchmark: SmolVLA vs. ACT

To put SmolVLA’s performance in context, I trained ACT (Action Chunking with Transformers) on the exact same 50 demonstrations and evaluated both models side-by-side.

ACT is a fundamentally different philosophy:

SmolVLA: the “Generalist adapted to a specialty” — 450M parameters, pretrained on 30K episodes, language-conditioned, flow matching. ACT: the “Specialist” — 52M parameters, trained from scratch, vision-only, Conditional VAE.

The question: does all that extra capacity and pretraining actually help on a single concrete task?

Results: The Numbers Tell a Clear Story

I evaluated SmolVLA on 7 episodes and ACT on 10 episodes, scoring each on three criteria (0-1 each, max 3.0 total):

- R1 — Fold Initiation: One arm clamps, the other folds over

- R2 — Fold Completion: Press down, tighten the fold

- R3 — Push Aside:** Move the folded cloth to clear the workspace

The Headline Numbers

Head-to-head comparison: SmolVLA scores 2.27/3.0 mean total reward vs ACT’s 0.63/3.0 — a 260% improvement. SmolVLA achieves 86% partial success rate vs ACT’s 20%.

SmolVLA outperforms ACT by 260% on mean total reward.

ACT’s best episode (1.80/3.00) — it manages a fold and push, but this level of performance was the exception, not the norm.*

What ACT Gets Wrong

ACT’s biggest failure mode is inaction — in 5 out of 10 episodes, the robot simply didn’t move. The arms positioned near the cloth but never executed a folding motion. When ACT does act, the fold initiation is actually decent (R1 of 0.70-0.80), but it almost never completes the full pipeline.

What SmolVLA Gets Right

SmolVLA is remarkably consistent. 6 out of 7 episodes initiate the fold, complete it to a reasonable degree, and push the cloth aside. The scores cluster tightly between 1.80 and 2.80 — it’s not that SmolVLA occasionally nails it, it’s that it reliably executes the full task pipeline.

Per-Episode Breakdown

Per-episode performance heatmap. Left: ACT shows sparse activation — 50% of episodes had zero motion (white rows). Right: SmolVLA shows dense, consistent high scores across all episodes. The contrast visually captures SmolVLA’s reliability vs ACT’s bimodal behavior.

Consistency across episodes — SmolVLA follows the same fold-and-push strategy with minor variations.

Why the Gap Is So Large

Three factors explain the 260% performance difference:

1. Pretraining Provides Motor Priors

SmolVLA doesn’t learn bimanual coordination from scratch. Its base model has seen 30,000 episodes of various manipulation tasks. It already “knows” how to coordinate two end-effectors, how to approach objects, and how to execute smooth trajectories. Fine-tuning on 50 episodes refines this knowledge for cloth specifically.

ACT starts from random initialization. 50 episodes and 100,000 gradient steps simply aren’t enough for a from-scratch model to reliably learn when to move and when not to. Hence the 50% inaction rate.

2. Flow Matching Produces Smoother Trajectories

SmolVLA’s flow matching generates trajectories by iterative denoising — the output is inherently smooth and temporally coherent. ACT uses a Conditional VAE, which can produce jerkier, less coordinated motions, especially in the bimanual setting where both arms need to synchronize.

3. Multi-Camera Fusion

Both models use the same three cameras, but SmolVLA’s SigLIP encoder (pretrained on internet-scale image data) extracts richer features from each view. ACT’s ResNet-18 (trained from scratch) has to learn visual representations simultaneously with motor policies — a much harder optimization problem.

Architectural comparison: ACT’s simpler pipeline (ResNet-18 + CVAE) vs SmolVLA’s richer pipeline (SigLIP + VLM backbone + Flow Matching action expert). Key differences highlighted: pretrained vs from-scratch vision, language conditioning, and flow matching vs CVAE.

Evaluation Method: Using an LLM as a Judge

One practical contribution worth highlighting: I used Claude Sonnet 4 as an automated evaluation judge. For each episode, the judge received 8-12 sampled frames and scored R1, R2, R3 on the defined rubric.

This mostly worked well, with one systematic bias: the VLM judge underscored R3 (push aside). From the overhead webcam angle, lateral displacement of the cloth was hard for the vision model to detect. Manual frame-by-frame verification corrected this, bumping SmolVLA’s R3 from 0.30 to 0.81.

Lesson: VLM-based evaluation is a useful automation, but always validate on a sample, especially for spatial reasoning tasks where camera perspective matters.

What I’d Do Differently

More demonstrations. 50 episodes is the minimum viable dataset. The fold completion scores (R2) are the weakest link — the model struggles with the fine-grained pressing motion. More demos with variation in fold tightness would likely help.

Better fabric variety. All 50 demos used the same cloth. Thin, slippery fabrics would likely cause failures. Mixing fabric types during data collection would improve robustness.

Recovery behavior. Neither model recovers after a missed grasp. The training data contains only successful demonstrations, so the policy has no “what to do when things go wrong” knowledge. Adding a few recovery demonstrations or using RL fine-tuning could address this.

Episode-level evaluation at scale. 7 and 10 episodes give directional signal but not statistical significance. A proper evaluation would run 50+ episodes per model.

Key Takeaways

1. Pretraining dominates. The single biggest factor in SmolVLA’s success is its pretrained base. 30,000 episodes of prior experience make 50 task-specific demonstrations go much further.

2. Compact VLAs are practical. 450M parameters runs on consumer hardware. You don’t need billion-parameter models for tabletop manipulation.

3. 50 demonstrations is enough to get started. Not enough for production reliability, but enough to demonstrate clear task competence with a pretrained model.

4. ACT is not the right baseline for bimanual deformable manipulation. ACT shines on rigid-object tasks with clear visual cues. Bimanual cloth folding exposes its limitations — particularly the lack of pretraining and the weaker visual backbone.

5. The LeRobot ecosystem is a force multiplier. Standardized data formats, community pretrained models, and plug-and-play training scripts dramatically reduce the engineering burden. This entire project — data collection, training, evaluation — was done by one person.

Resources and Reproducibility

Everything needed to reproduce this work is open:

- Training dataset: Dataset Link (50 episodes)

- SmolVLA policy: Finetuned SmolVLA Model

- ACT policy: Trained ACT Model

- Evaluation data (SmolVLA) SmolVLA Evaluation Set

- Evaluation data (ACT): ACT Evaluation Set

- Base model: https://huggingface.co/lerobot/smolvla_base

- Framework: https://github.com/huggingface/lerobot

If you’re working on manipulation with open-source robots or experimenting with VLAs, I’d love to hear about it. Drop a comment or reach out — this space moves fast and I think we’re just getting started.

Our live bootcamp on VLA and World Models: https://robotlearningmastery.vizuara.ai

Minor in Robotics:

https://minor-robotics.vizuara.ai

This chapter covers

• Understanding inductive biases in CNNs and why they matter for image understanding

• Why Vision Transformers are data-hungry and computationally expensive without these biases

• How DeiT (Data-efficient Image Transformers) achieves competitive accuracy using only ImageNet-1K

• The mechanics of knowledge distillation, from soft labels to dark knowledge

• The mathematics behind temperature scaling, KL divergence, and DeiT’s loss functions

• Building a complete DeiT implementation from scratch in PyTorch

Data Efficient Image Transformer Code is available below

https://github.com/VizuaraAI/Transformers-for-vision-BOOK

[Don't forget to star the code repo!]

In previous blog, we explored how Vision Transformers (ViTs) brought the power of self-attention to computer vision by splitting images into patches and processing them as sequences. The results were remarkable: ViT achieved state-of-the-art accuracy on ImageNet. But there was a catch. That performance came at an enormous cost. The original ViT required pre-training on JFT-300M, a private Google dataset containing 300 million labeled images, and demanded thousands of TPU-days of compute. For most researchers and practitioners, this was simply out of reach.

In this blog, we will explore how DeiT (Data-efficient Image Transformers) solved this problem. Published by Touvron et al. from Facebook AI Research and Sorbonne University in 2021, DeiT achieved competitive accuracy using only ImageNet-1K (1.2 million images), which is 250 times less data than the original ViT required. The key insight was combining a clever training recipe with knowledge distillation, a technique where a smaller student model learns from a larger, pre-trained teacher model. But before we dive into DeiT itself, we need to understand why Vision Transformers struggle with limited data in the first place. The answer lies in a concept called inductive bias.

Prerequisites blogs

1.1 Inductive bias: the built-in assumptions of neural networks

When we train a neural network, we are asking it to learn patterns from data. But no model starts from a completely blank slate. Every architecture carries certain built-in assumptions about how data is structured, and these assumptions shape what the model finds easy or difficult to learn. These built-in assumptions are called inductive biases.

Think of it this way: if you were asked to find a lost cat in a neighborhood, you would naturally check under porches, in gardens, and near food sources. You would not start by searching the sky or underwater. Your prior knowledge about where cats tend to be gives you a huge advantage. You are not starting from scratch. Inductive bias works the same way for neural networks: it encodes structural assumptions that guide learning in the right direction, reducing the amount of data needed to reach good performance.

What are the inductive biases in CNNs?

Convolutional Neural Networks (CNNs) have two powerful inductive biases built directly into their architecture: locality and translation equivariance. Let us examine each in detail.

Locality bias. CNNs assume that nearby pixels are more related to each other than distant pixels. This assumption is enforced by the small convolutional filters (typically 3x3 or 5x5) that examine only a local patch of the image at each step. The network processes the image by looking at small neighborhoods first, detecting low-level features like edges and textures, and then gradually combining these local features into higher-level concepts through deeper layers. An edge detector in the first layer might find horizontal lines. The next layer combines several edges into a corner or a curve. Deeper layers assemble these parts into recognizable shapes like eyes, wheels, or letters.

This hierarchical, local-to-global processing mirrors how natural images are actually structured. In a photograph, a pixel showing part of a dog’s ear is highly correlated with neighboring pixels that also show the ear, but has little direct relationship with a pixel in the far corner showing a blade of grass. CNNs exploit this statistical property by design.

Figure 1.1: The effect of inductive bias on translation handling. Top: a cat is placed at four different positions within the input grid. A translation invariant model (CNN) produces the same logits every time, correctly recognizing the cat regardless of where it appears. This is because weight sharing and pooling are built into the architecture. Bottom: the same cat at different positions is passed through a model sensitive to translation (such as a ViT trained on limited data), which produces different logits for each position. The CNN does not need to learn that position should not affect the classification; this property is baked into its design. A ViT must discover this from data, which is why it requires far more training examples.

Translation equivariance.

The second key inductive bias is translation equivariance, a CNN detects the same feature regardless of where it appears in the image. This property comes from weight sharing , where the same convolutional filter is applied identically at every spatial position. Mathematically, f is a convolution operation and T is a spatial translation, then:

This means that if an object shifts position in the input, the corresponding feature map shifts by exactly the same amount. A cat-ear detector will fire whether the ear appears in the top-left corner or the bottom-right corner of the image. Figure 1.2 shows two concrete examples: the same convolutional kernel slides over the entire input and detects the same pattern wherever it occurs, whether that is a left edge in a grid of circles or a car on a road versus on a building rooftop.

Figure 1.2: Translation equivariance in convolutions Top: a convolutional kernel (shown on the left) is applied across an input grid containing circles at different positions. The filter identifies left edges irrespective of where they appear in the input, producing the same edge-detection response in the output. Bottom: the same principle applied to a real-world scene. The kernel identifies “car” features regardless of whether the car is on the road or on a building rooftop. Because the same filter weights are shared across all spatial positions, the network does not need to learn separate detectors for every possible location of a feature. (conceptual illustration)

When we add pooling layers (max pooling or global average pooling) on top of convolutions, we achieve translation invariance: the final classification output remains the same regardless of where the object appears. The pooling operation discards precise positional information, summarizing each region with a single statistic (maximum or average value). This is why a CNN classifies an image as "cat" whether the cat is centered, shifted left, or shifted right.

NOTE

Translation equivariance and translation invariance are related but distinct properties. Equivariance means "the output shifts with the input" (convolution layers). Invariance means "the output stays the same regardless of position" (achieved after pooling). A CNN is equivariant through its convolutional layers and becomes invariant at its classification output through pooling.

Why are these biases so powerful?

These inductive biases are powerful precisely because they match the statistical structure of natural images. Consider what they give us:

1. Fewer parameters. Weight sharing means one small filter is reused across the entire image, rather than learning separate weights for every position. A 3x3 filter on 3 channels has only 27 parameters, yet it processes every location in the image.

2. Better generalization with less data. Because the model already “knows” that patterns can appear anywhere and that local neighborhoods matter, it does not need millions of examples to discover these facts.

3. Natural feature hierarchies. The local-to-global processing pipeline naturally builds the kind of part-based representations that are useful for recognition: edges combine into textures, textures into parts, parts into objects.

This is exactly why CNNs dominated computer vision for nearly a decade, from AlexNet in 2012 through the ResNet era. The built-in assumptions aligned beautifully with the task.

When do these biases become a limitation?

However, the same rigid assumptions that make CNNs data-efficient can become a ceiling on performance. Consider the image in figure 1.3: a person walking near a traffic signal. To understand this scene, the model needs to relate the person (in one part of the image) to the traffic signal (in a completely different part). These two elements are spatially far apart, but semantically they are deeply connected: the traffic signal determines whether the person should walk or stop.

Figure 1.3 Long-range dependencies in images. The person and the traffic signal are spatially far apart in pixel space, but they are semantically related. A CNN must stack many layers to connect these distant regions through its limited local receptive field. A transformer, with global self-attention, can relate these elements directly in a single layer.

A CNN's local filters can only see a small neighborhood at each layer. To connect the person and the traffic signal, information must propagate through many layers, each expanding the receptive field slightly. By the time the model can "see" both elements simultaneously, the information has passed through many transformations and may be diluted or lost.

This is where transformers shine. Self-attention computes relationships between all patches simultaneously, regardless of spatial distance. A transformer can directly attend from the "person" patch to the "traffic signal" patch in a single layer. As we discussed in the context of natural language processing, this is analogous to how attention connects distant words in a sentence.

Figure 1.4 Attention mechanisms connect distant elements directly. Just as attention in NLP allows the word "smiled" to attend strongly to "teacher" despite the intervening words, visual attention allows distant image patches to interact directly. The attention weights (shown below the sentence) indicate how strongly each word attends to others, capturing long-range dependencies that local convolutions would struggle with.

So, inductive bias in CNNs is a double-edged sword. When the assumptions match the data (local features, translation-invariant patterns, moderate image complexity), CNNs are remarkably efficient. When the task demands flexible, global reasoning across distant image regions, those same assumptions become a constraint. This trade-off is precisely what motivated the development of Vision Transformers, and subsequently, DeiT.

1.2 Why Vision Transformers are data-hungry and computationally heavy

Now that we understand what inductive biases are and why CNNs benefit from them, we can understand the fundamental challenge that Vision Transformers face.

The ViT paper’s stark finding

The original Vision Transformer paper by Dosovitskiy et al. (2020) contained an important admission that is easy to overlook:

The paper continued with a crucial finding:

Let us unpack exactly why. A Vision Transformer processes images very differently from a CNN:

1. No locality constraint. In every transformer layer, each patch attends to every other patch through self-attention. There is no architectural constraint forcing the model to prioritize nearby patches. The attention mechanism computes pairwise relationships between all patches simultaneously, whereas CNNs restrict each layer to local neighborhoods.

2. No weight sharing across positions. Unlike CNNs where the same filter slides across the entire image, in ViTs the attention mechanism can learn different relationships for different positions. Positional information comes only from learned positional embeddings, not from the architecture itself.

3. No hierarchical feature extraction by default. CNNs naturally build hierarchical features (edges, then textures, then parts, then objects) through stacked local convolutions with increasing receptive fields. Standard ViTs have uniform global attention at every layer, so this hierarchy must be learned entirely from data.

The data requirements

The consequence is stark. Since a ViT must learn locality, translation patterns, and hierarchical feature extraction all from scratch, it needs vastly more data.

The original ViT paper systematically studied this relationship with a dataset scaling experiment:

• ImageNet-1K (~1.2 million images): ViT-Large underperforms comparable CNNs. Worse, larger ViT variants perform worse than smaller ones due to overfitting. The model memorizes training data rather than generalizing.

• ImageNet-21K (~14 million images): ViT-Large and comparable CNNs perform similarly. The ViT begins to show its potential.

• JFT-300M (~303 million images): ViT-Large significantly outperforms comparable CNNs. The pattern reverses completely: larger models perform better, not worse.

This reveals a fundamental trade-off in machine learning: the less prior knowledge (inductive bias) you build into an architecture, the more data it needs to learn those patterns from experience. CNNs “know” about locality and translation equivariance before seeing a single training image. ViTs must discover these properties purely from data, and discovering them requires seeing hundreds of millions of examples.

Scaling laws and the bitter lesson

This phenomenon connects to a broader principle in deep learning captured by scaling laws. Research by Kaplan et al. (2020) established that neural network performance follows predictable power-law relationships with three variables: model size (N parameters), dataset size (D tokens or images), and compute budget (C FLOPs):

Figure 1.5 Scaling laws demonstrate that model performance improves as a smooth power law of model size, dataset size, and compute budget. These relationships, first established for language models, were later confirmed for Vision Transformers as well. The declining curves show that each doubling of resources yields a predictable (though diminishing) improvement in performance.

The Scaling Vision Transformers paper by Zhai et al. (2022) confirmed that these same power-law relationships hold for ViTs, scaling from 5 million to 2 billion parameters. The key insight is that architectural details matter less than scale. With enough data and compute, the flexibility of transformers (their lack of restrictive inductive biases) becomes an advantage rather than a limitation. The ViT paper itself summarized this as: “Large-scale training trumps inductive bias.”

But this raised an uncomfortable question: what if you do not have 300 million images and thousands of TPU-days? What if you have a single machine, a standard dataset, and a few days of training time? This is precisely the problem that DeiT set out to solve.

The small context window insight

Consider how a small local patch of an image can be ambiguous without global context. Figure 1.6 shows an interesting property: small context windows of different digits can look remarkably similar. The bottom halves of the digits 8, 0, and 6 share nearly identical local features. A local convolutional filter might struggle to distinguish them, while global attention can consider the entire digit.

Figure 1.6 Small context windows can be ambiguous. The bottom portions of the digits 8, 0, and 6 appear nearly identical when viewed in isolation. This illustrates both the strength and weakness of local processing: CNNs capture these shared local features efficiently, but transformers with global attention can disambiguate by considering the full spatial context simultaneously.

This is a microcosm of the broader tension between CNNs and ViTs. Local processing is efficient but can miss the forest for the trees. Global processing is powerful but expensive to learn. DeiT’s genius was finding a way to get the best of both worlds.

1.3 DeiT: data-efficient training through distillation

With the problem clearly defined (ViTs need too much data and compute), let us explore how DeiT solved it. The paper “Training data-efficient image transformers & distillation through attention” by Touvron et al. (2021) introduced a remarkably elegant solution that combines two key ingredients: a carefully designed training recipe and a novel form of knowledge distillation.

The core idea

DeiT’s central insight is this: instead of requiring the Vision Transformer to learn everything about images from scratch, we can transfer knowledge from a pre-trained CNN teacher. The CNN already understands locality, translation equivariance, and hierarchical features because these properties are built into its architecture. Through knowledge distillation, the ViT student can inherit these implicit biases without having them hardcoded into its architecture.

The result was striking: DeiT-B (86 million parameters) achieved 83.1% top-1 accuracy on ImageNet using only ImageNet-1K for training, in approximately 53 hours on a single 8-GPU node. Compare this with the original ViT, which required JFT-300M (300 million images) and thousands of TPU-days to achieve similar performance.

DeiT architecture

DeiT retains the standard Vision Transformer architecture with one crucial addition: a distillation token. Let us walk through the full architecture as shown in figure 1.7.

Figure 1.7 The DeiT architecture extends the standard Vision Transformer with a distillation token. The input image is divided into fixed-size patches, which are linearly projected into embeddings. Two special tokens are prepended: the [CLS] token (supervised by ground-truth labels) and the [DIST] token (supervised by the teacher model's predictions). Both tokens pass through all transformer encoder layers, interacting with patch tokens and with each other through self-attention. At the output, separate classification heads produce predictions from each token.

The architecture works as follows:

1. Patch embedding. The input image (e.g., 224x224 pixels) is divided into non-overlapping patches (e.g., 16x16 pixels each), yielding a sequence of

   patches. Each patch is linearly projected into a D-dimensional embedding vector using a learnable projection matrix. In practice, this is implemented as a single convolution with kernel size and stride equal to the patch size.

2. Special tokens. Two learnable embedding vectors are prepended to the sequence:

   • The [CLS] token: a standard classification token (as in the original ViT) that aggregates information for predicting the ground-truth label.

   • The [DIST] token: a novel distillation token that aggregates information for mimicking the teacher model’s predictions.

3. Positional embeddings. Learnable 1D position embeddings are added to all tokens (patches + CLS + DIST), giving the model information about spatial arrangement. The total sequence length is N+2.

4. Transformer encoder. The sequence passes through L standard transformer encoder layers, each consisting of multi-head self-attention (MSA) and a feed-forward network (FFN). Both special tokens interact with all patch tokens and with each other through self-attention.

5. Dual classification heads. At the output, two separate linear heads produce predictions:

6. Inference. At test time, the predictions from both heads are combined by averaging the softmax outputs:

Now let us look at the original architecture diagram from the paper for additional perspective.

Figure 1.8 The DeiT architecture as presented in the original paper, showing the complete training pipeline. The student (Vision Transformer) receives supervision from two sources: the ground-truth labels (via cross-entropy on the CLS token) and the pre-trained CNN teacher (via the distillation loss on the DIST token). The teacher model is frozen during training and its weights are never updated. The distillation token enables the transformer to learn the CNN's implicit understanding of image structure without architecturally constraining itself.

Why a separate distillation token?

You might wonder: why not simply add the teacher’s loss to the existing CLS token? Why introduce a separate token? The DeiT paper provides a compelling empirical answer.

The authors experimented with using two identical CLS tokens (both trained on ground-truth labels) and found that they converge to nearly identical representations, with cosine similarity approaching 0.999. They carry redundant information.

In contrast, the CLS token and distillation token, trained on different objectives, develop meaningfully different representations. Their cosine similarity is approximately 0.06 in early layers and rises to only about 0.93 by the final layer. This means the distillation token captures complementary information that the CLS token alone would miss. The two tokens learn different “perspectives” on the input, and combining them at inference yields better predictions than either alone.

The teacher model

DeiT uses a pre-trained CNN as the teacher. In the paper, the primary teacher is RegNetY-16GF (84 million parameters, 82.9% top-1 accuracy on ImageNet), though the authors also experimented with other architectures. Critically, the teacher is frozen during training: its weights are never updated. It simply provides predictions that the student learns to mimic.

A surprising finding from the paper is that a CNN teacher produces dramatically better student performance than a transformer teacher. The authors note that the transformer student “learned least from a transformer-teacher but learned most from a big convolution-teacher.” This supports the hypothesis that distillation effectively transfers the CNN’s inductive biases (locality, translation equivariance) to the transformer student, giving it the benefits of both architectures.

But how exactly does knowledge transfer from teacher to student? To understand this, we need to dive deep into the mechanics of knowledge distillation.

1.4 Knowledge distillation: teaching a student network

Knowledge distillation is a model compression technique where a small student model is trained to mimic the behavior of a larger, more capable teacher model. The goal is to produce a compact model that retains much of the teacher’s performance while being faster and cheaper at inference. Let us trace the evolution of this idea and then build up the mathematics step by step.

A brief history

The idea of compressing knowledge from large models into small ones dates back to 2006, when researchers at Cornell University proposed model compression. At the time, the best-performing models were not single networks but ensembles of hundreds or thousands of models whose predictions were averaged. These ensembles were accurate but far too large for deployment on devices like PDAs (personal digital assistants, the predecessors to smartphones).

Figure 1.9 The two-step process of model compression, as introduced in 2006. In Step 1, an ensemble of multiple models is trained, with each model making independent predictions. In Step 2, a single small model is trained to directly predict the averaged output of the ensemble, compressing the knowledge of many models into one.

The insight was elegantly simple: instead of shipping all the ensemble models to the device, train a single small model to predict the ensemble's averaged output. This small model captures the collective wisdom of the ensemble while being compact enough for deployment.

Figure 1.10 Ensemble soft voting. Multiple models (Model 1 through Model 5) each output a probability distribution over classes. These distributions are averaged to produce a final classification. The key insight is that this averaged distribution carries more information than any single model’s hard prediction.

In 2015, Geoffrey Hinton and Jeff Dean (at Google) revisited this idea and gave it the name we use today: knowledge distillation. Their crucial insight was that distillation is valuable even without an ensemble. A single large model’s soft probability outputs carry richer information than hard labels alone, and this information can be transferred to a smaller student.

Figure 1.11 Knowledge distillation does not require an ensemble. A single large teacher model generates soft labels (probability distributions) that carry more information than one-hot hard labels. The student learns from these enriched targets, inheriting the teacher’s nuanced understanding of inter-class relationships.

Hard labels vs. soft labels: the concept of dark knowledge

To understand why soft labels are so valuable, consider a concrete example. Suppose we have a teacher model classifying animal images into three categories: dog, cat, and mouse.

Given an image of a husky, the hard label (ground truth) is simply:

This tells the student: "This is a dog. Period." But the teacher's soft prediction might be:

This soft distribution tells the student something much richer: “This is most likely a dog, but it looks somewhat like a cat (perhaps because of the fur texture), and it looks very little like a mouse.” This additional information about what the input is not (and how much it is “not” each class) is what Hinton poetically called dark knowledge.

Figure 1.12 Dark knowledge revealed through soft labels. While hard labels simply say "this is a husky," the teacher's soft probability distribution reveals richer inter-class relationships. A husky image gets high probability for "husky" but also notable probability for "wolf" and some for "dog," reflecting visual similarities between these animals that the teacher has learned. This relational information (the dark knowledge) helps the student learn more efficiently.

Dark knowledge encodes the teacher’s learned understanding of inter-class similarities and relationships. A husky looks more like a wolf than like a goldfish. A handwritten “7” looks more like a “1” than like a “0.” These relational insights are completely absent from one-hot hard labels but are naturally present in soft probability distributions.

The mathematics of knowledge distillation

Now let us formalize the distillation process mathematically. The framework involves three key components: temperature scaling, the distillation loss function, and the combined training objective.

Temperature scaling

When a well-trained teacher model makes predictions, its output distribution is often very “peaked”: the correct class gets probability close to 1.0, and all other classes get near-zero probabilities. For example:

These near-zero probabilities contain the dark knowledge we want to transfer, but they are so small that they provide almost no gradient signal during training. The student cannot learn from values like 0.005 because the gradients are essentially zero.

Temperature scaling solves this by softening the probability distribution. Given a vector of logits

(the raw, unnormalized outputs of the network before softmax), the standard softmax function is:

With a temperature parameter τ>0 , we define the temperature-scaled softmax:

The temperature controls how "spread out" the distribution is:

Standard softmax. The distribution reflects the model's learned confidence.

Higher temperature softens the distribution, making all probabilities more uniform and revealing the relative differences between logits. This amplifies dark knowledge.

The distribution approaches a uniform distribution 1/K for K classes.

The distribution collapses to a one-hot vector concentrated on the largest logit (argmax).

Figure 1.13 The effect of temperature on softmax distributions. The formula shows how temperature controls the “peakiness” of the output. Left: the raw logit values before temperature-scaled softmax. Right: after applying softmax with high temperature, the distribution becomes more uniform, revealing the relative magnitudes of all logits. This smoothing is essential for transferring dark knowledge from teacher to student.

Let us work through a concrete numerical example. Suppose a teacher produces logits z=[5.0,2.0,1.0] for three classes. Here is how temperature affects the resulting probabilities:

At τ=1, the distribution is peaked: dog gets 93.6% and the dark knowledge (cat and mouse probabilities) is barely visible. At τ=5, the distribution is much flatter: we can clearly see that the teacher considers this image more cat-like than mouse-like (27.5% vs. 22.5%). This relational information is the dark knowledge that distillation transfers.

Figure 1.14 A direct comparison of probability distributions before and after temperature scaling. The left distribution (low temperature) shows sharp peaks where the correct class dominates. The right distribution (high temperature) shows a smoothed version where the relative probabilities of non-dominant classes become visible. This smoothing allows the student to learn not just which class is correct, but which incorrect classes are most similar to the correct one.

Cross-entropy loss

The standard classification loss in deep learning is cross-entropy. For a ground-truth one-hot label y=[y1,y2,…,yK] and predicted probabilities p=[p1,p2,…,pK], the cross-entropy is:

Since y is one-hot with y_c=1 for the true class cc and y_i=0 for all other classes, this simplifies to:

Let us verify with an example. If the true class is "cat" (c=1) and the model predicts p=[0.80,0.15,0.05]:

A confident correct prediction gives a low loss. If instead pc=0.1:

An incorrect or uncertain prediction gives a high loss, pushing the model to adjust its weights. Note that −log⁡(x) produces a positive value when x∈(0,1), which is always the case for probabilities.

KL divergence

While cross-entropy measures how well predictions match a target distribution, Kullback-Leibler (KL) divergence measures how different two probability distributions are from each other. For distributions p (teacher) and q (student):

KL divergence has several important properties:

Let us work through a concrete example. Suppose the teacher outputs

and the student outputs

A small KL divergence indicates the student is already quite close to the teacher. Now consider an overconfident student that outputs

The KL divergence is larger because the student is more confident than the teacher. KL divergence heavily penalizes cases where the student assigns very low probability to classes that the teacher considers plausible.

NOTE

There is an important mathematical subtlety about KL divergence when pi=0. Since

(which can be verified using L'Hôpital's rule), terms where pi=0 contribute zero to the divergence. This means the KL divergence is well-defined even when the teacher assigns zero probability to some classes.

The relationship between KL divergence and cross-entropy is:

where

is the cross-entropy and

is the entropy of the teacher distribution. Since the teacher is frozen, H(p) is a constant, and minimizing KL divergence is equivalent to minimizing the cross-entropy between teacher and student distributions.

The complete knowledge distillation loss

Putting it all together, the knowledge distillation loss from Hinton et al. (2015) combines two terms:

The first term is the standard cross-entropy between the student’s predictions (at τ=1) and the ground-truth labels. This keeps the student honest: it must still learn to classify correctly.

The second term is the KL divergence between the teacher’s and student’s softened predictions (both at temperature τ). This transfers the dark knowledge from teacher to student.

Figure 1.15 The complete knowledge distillation framework. Training data feeds into both the teacher (large network, frozen) and the student (small network, being trained). The teacher produces soft labels via temperature-scaled softmax, while the ground truth provides hard labels. The student's total loss combines two components: the cross-entropy loss L_CE against hard labels, and the KL divergence loss L_KL against the teacher's soft labels. The formula

weights these components, and gradients flow back through the student to update its weights.

You may wonder why the KL divergence term is multiplied by τ^2. This is not arbitrary; it is a mathematical necessity for keeping the gradient magnitudes balanced.

When we compute the gradient of the soft cross-entropy loss with respect to the student's logits zs,i:

The temperature scaling introduces a factor of 1/τ in the gradient. Additionally, the softened probabilities themselves compress the differences between logits by another factor of approximately 1/τ, yielding an overall gradient scaling of approximately 1/τ2.

Without compensation, the soft-label gradients would vanish as we increase temperature, defeating the purpose of softening. The τ_2 multiplier restores the gradient magnitudes to be comparable with the hard-label loss, ensuring that the weighting coefficient αα behaves predictably regardless of the chosen temperature.

High-temperature approximation: why distillation equals logit matching

There is an elegant mathematical insight that connects distillation to a simpler concept. At high temperatures, the softmax can be approximated using a Taylor expansion:

where

is the mean logit and K is the number of classes.

Substituting this into the KL divergence and simplifying, the soft-label loss reduces to:

This is simply a mean squared error between the logits of the student and teacher! At high temperature, knowledge distillation is approximately equivalent to logit matching, which was actually proposed in 2014 as a precursor to Hinton's method. The τ_2 factor cancels the 1/τ2 scaling, confirming that the multiplication is necessary. This result also reveals why dark knowledge transfer works: the student is effectively learning the teacher's internal ranking of all classes, not just its top prediction.

1.5 DeiT’s distillation: hard labels beat soft labels

Now that we understand the mechanics of knowledge distillation, let us see how DeiT adapts this framework. DeiT introduces two variants: soft distillation and hard distillation. Surprisingly, the simpler approach wins.

Soft distillation in DeiT

DeiT’s soft distillation uses the standard knowledge distillation framework, but with the distillation token producing separate logits:

The CLS token head learns from ground-truth labels. The distillation token head learns to match the teacher’s full soft probability distribution.

Hard distillation in DeiT

Hard distillation replaces the KL divergence with a simple cross-entropy against the teacher’s hard prediction (argmax):

This is remarkably simple: no temperature parameter, no KL divergence, no τ2 correction. Just two cross-entropy losses weighted equally. The CLS token learns from the ground truth, and the distillation token learns from the teacher’s top-1 prediction.

Why hard distillation works better

The DeiT paper found that hard distillation outperforms soft distillation by approximately +1.0--1.2% accuracy on ImageNet. This was surprising because soft labels contain strictly more information than hard labels.

The authors hypothesize that this relates to the fundamental difference between the teacher (a CNN) and the student (a transformer). CNNs and transformers process visual features in fundamentally different ways: CNNs use local convolutional filters while transformers use global self-attention. When the teacher provides a full soft distribution, it reflects the CNN’s specific way of processing the image, which may not transfer well to the transformer’s very different processing pipeline. Hard labels abstract away these architectural details, providing a cleaner supervisory signal that is easier for the transformer to learn from.

Think of it like learning to cook. Soft distillation is like watching a master chef’s exact hand movements (which depend on their specific knife and grip style). Hard distillation is like reading their recipe (the end result, abstracted from the specific process). If you have different tools and a different grip, the recipe is more useful than mimicking movements that do not suit your tools.

DeiT results

The results are compelling. Figure 1.16 shows the performance-throughput trade-off for various models.

Figure 1.16 Performance comparison between DeiT variants and competing architectures, plotting throughput (images per second, higher is better) against ImageNet top-1 accuracy (higher is better). DeiT-B with distillation (marked with the alembic symbol) achieves 85.2% accuracy, matching or exceeding EfficientNet while being significantly faster at inference. The key takeaway is that DeiT achieves competitive accuracy with ViT-level throughput, using only ImageNet-1K for training rather than the 300 million images that the original ViT required.

DeiT demonstrated that the Vision Transformer’s dependence on massive datasets was not an inherent limitation of the architecture, but a training problem with a training solution. By combining knowledge distillation with a dedicated distillation token, DeiT gave the transformer a way to absorb the inductive biases of a CNN teacher without modifying the transformer architecture itself. The CLS token learns from ground truth labels while the distillation token learns from the teacher’s predictions, and the two complementary signals together produce a model that generalizes better than either supervision source alone. This insight, that architectural inductive biases can be transferred rather than engineered, opened the door to practical vision transformers that anyone with a single multi-GPU machine could train.

1.6 Building DeiT from scratch in PyTorch

Data Efficient Image Transformer Code is available below

https://github.com/VizuaraAI/Transformers-for-vision-BOOK

[Don’t forget to star the code repo!]

With the theory firmly established, let us now build a complete DeiT implementation from scratch. We will train a small-scale version on MNIST digits to demonstrate all the key concepts: patch embedding, the distillation token, the teacher-student setup, and the knowledge distillation loss. While MNIST is far simpler than ImageNet, it lets us see the entire pipeline working end-to-end on a single machine.

Setting up imports and configuration

We begin by importing the necessary libraries and defining our hyperparameters. We will use PyTorch for the implementation, torchvision for datasets and the pre-trained teacher model, and NumPy for utility operations.

Listing 1.1 Importing libraries and setting up the device

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torch.utils.data import Subset, DataLoader
from torchvision import transforms, datasets, models
import numpy as np
import matplotlib.pyplot as plt

device = 'cuda' if torch.cuda.is_available() else 'cpu' #A

#A Automatically selects GPU if available, otherwise falls back to CPU

Next, we define the hyperparameters that control our model and training process. These values are deliberately small compared to the full DeiT paper to allow rapid experimentation on a personal machine.

Listing 1.2 Defining hyperparameters

BATCH_SIZE = 12
ATTENTION_HEADS = 4          #A
TRANSFORMER_LAYERS = 4       #B
EMBED_DIM = 16               #C
IMG_SIZE = 28                #D
PATCH_SIZE = 7               #E
CLASSES = 10                 #F
EPOCHS_STUDENT = 10
LR_STUDENT = 1e-4
TEMPERATURE = 4              #G
ALPHA = 0.1                  #H
CHANNELS = 3                 #I

#A Number of attention heads in each transformer layer

#B Number of stacked transformer encoder layers

#C Embedding dimension for patch tokens (small for demonstration)

#D MNIST images are 28x28 pixels

#E Each patch is 7x7 pixels, giving us (28/7)^2 = 16 patches

#F MNIST has 10 digit classes (0--9)

#G Temperature for softening the teacher's probability distribution

#H Weight for the KL divergence term in the distillation loss

#I Number of input channels (expanded from 1 to 3 for CNN compatibility)

Let us understand the patch arithmetic. With 28x28 images and 7x7 patches, each image is divided into

non-overlapping patches. Each patch is flattened and linearly projected into a 16-dimensional embedding. Together with the CLS token and distillation token, this gives us a sequence of 16+2=18 tokens.

Preparing the data

MNIST images are grayscale (1 channel), but our teacher model (ResNet50) expects 3-channel RGB inputs. We handle this by repeating the single channel three times. We also use only 10% of the training set to simulate a data-scarce scenario, which is the exact setting where DeiT’s distillation approach shines.

Listing 1.3 Loading and preparing the MNIST dataset

tfm = transforms.Compose([
    transforms.ToTensor(),
    transforms.Lambda(lambda t: t.repeat(3, 1, 1)),    #A
])

train_full = datasets.MNIST('./data', train=True, download=True, transform=tfm)
test = datasets.MNIST('./data', train=False, download=True, transform=tfm)

n = int(0.1 * len(train_full))                         #B
subset_idx = np.random.permutation(len(train_full))[:n]
train = Subset(train_full, subset_idx)

train_dl = DataLoader(train, batch_size=BATCH_SIZE, shuffle=True)
test_dl = DataLoader(test, batch_size=BATCH_SIZE)

#A Converts single-channel grayscale to 3-channel by repeating, making it compatible with the ResNet teacher

#B Uses only 10% of training data (6,000 images from 60,000) to simulate data scarcity

Using only 6,000 training images is deliberately challenging. This mimics the real-world scenario that motivated DeiT: how do we train a Vision Transformer effectively when data is limited? Knowledge distillation is the answer.

Setting up the teacher model

Our teacher is a pre-trained ResNet50, one of the most well-known CNN architectures. We load it with ImageNet pre-trained weights and modify the final classification layer to output 10 classes (for MNIST digits) instead of the original 1,000 ImageNet classes.

Critically, we freeze all the teacher’s parameters except the final classification layer. The teacher’s convolutional feature extractors already understand visual patterns from ImageNet pre-training. We only need to adapt the final layer to map those features to our 10 digit classes.

Listing 1.4 Setting up the pre-trained CNN teacher

teacher = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V2)
teacher.fc = nn.Linear(teacher.fc.in_features, CLASSES)    #A
teacher.to(device)

for param in teacher.parameters():                         #B
    param.requires_grad = False

for param in teacher.fc.parameters():                      #C
    param.requires_grad = True

#A Replaces the 1000-class ImageNet head with a 10-class head for MNIST

#B Freezes all layers: the convolutional backbone will not be updated

#C Unfreezes only the final classification layer so it can learn MNIST-specific mappings

This setup mirrors the DeiT paper’s approach: the teacher is a powerful CNN that has already learned rich visual representations. By freezing its backbone, we preserve the inductive biases (locality, translation equivariance, hierarchical features) that the CNN learned during ImageNet pre-training. The student will learn to mimic these capabilities through distillation.

Building the student Vision Transformer

Now we build the student model: a Vision Transformer with a distillation token. This is the heart of the DeiT architecture. Let us break it into two components: the patch embedding layer and the full ViT model.

Listing 1.5 Patch embedding module

class PatchEmbed(nn.Module):
    def __init__(self, img_size=IMG_SIZE, patch=PATCH_SIZE,
                 dim=EMBED_DIM, channels=CHANNELS):
        super().__init__()
        self.proj = nn.Conv2d(channels, dim, patch, patch)  #A
        self.n = (img_size // patch) ** 2                   #B

    def forward(self, x):
        x = self.proj(x)           #C
        x = x.flatten(2)           #D
        x = x.transpose(1, 2)     #E
        return x

#A A Conv2d with kernel_size=patch_size and stride=patch_size extracts non-overlapping patches and projects them to the embedding dimension in one operation

#B Computes the number of patches: (28/7)^2 = 16 patches

#C Applies the convolution: (B, 3, 28, 28) -> (B, 16, 4, 4)

#D Flattens the spatial dimensions: (B, 16, 4, 4) -> (B, 16, 16)

#E Transposes to get sequence format: (B, 16, 16) -> (B, 16, 16) where first 16 is sequence length and second 16 is embedding dim

The patch embedding uses a single convolution to simultaneously extract patches and project them into the embedding space. This is mathematically equivalent to extracting each patch, flattening it, and multiplying by a weight matrix, but it is computationally more efficient.

Now let us build the full ViT model with the distillation token:

Listing 1.6 The DeiT student model with distillation token

class ViT(nn.Module):
    def __init__(self, num_classes=CLASSES, dim=EMBED_DIM,
                 depth=TRANSFORMER_LAYERS, heads=ATTENTION_HEADS):
        super().__init__()
        self.patch = PatchEmbed()
        n = self.patch.n

        self.cls = nn.Parameter(torch.zeros(1, 1, dim))      #A
        self.distill = nn.Parameter(torch.zeros(1, 1, dim))   #B
        self.pos = nn.Parameter(torch.zeros(1, n + 2, dim))   #C

        self.blocks = nn.Sequential(*[                        #D
            nn.TransformerEncoderLayer(
                dim, heads, dim * 4, batch_first=True
            )
            for _ in range(depth)
        ])

        self.norm = nn.LayerNorm(dim)                         #E
        self.head_cls = nn.Linear(dim, num_classes)           #F
        self.head_dist = nn.Linear(dim, num_classes)          #G

    def forward(self, x):
        B = x.size(0)
        x = self.patch(x)                                     #H

        cls = self.cls.expand(B, -1, -1)                      #I
        dist = self.distill.expand(B, -1, -1)

        x = torch.cat([cls, x, dist], dim=1) + self.pos      #J

        x = self.blocks(x)                                    #K
        x = self.norm(x)

        cls_out = x[:, 0]                                     #L
        dist_out = x[:, -1]                                   #M

        cls_logits = self.head_cls(cls_out)                    #N
        dist_logits = self.head_dist(dist_out)

        return cls_logits, dist_logits

student = ViT().to(device)
opt_s = torch.optim.AdamW(student.parameters(), lr=LR_STUDENT)

#A Learnable CLS token: initialized to zeros, shape (1, 1, dim)

#B Learnable distillation token: the key DeiT innovation, also initialized to zeros

#C Positional embeddings for all tokens: n patches + CLS + DIST = 18 positions

#D Stack of transformer encoder layers, each with multi-head self-attention and FFN

#E Layer normalization applied after the final transformer block

#F Classification head for the CLS token (trained against ground truth)

#G Separate classification head for the distillation token (trained against teacher)

#H Convert image to patch embeddings: (B, 3, 28, 28) -> (B, 16, 16)

#I Expand special tokens to match batch size: (1, 1, 16) -> (B, 1, 16)

#J Concatenate CLS + patches + DIST and add positional embeddings: (B, 18, 16)

#K Pass through all transformer encoder layers

#L Extract the CLS token output (first position)

#M Extract the distillation token output (last position)

#N Produce class logits from each token through separate linear heads

There are several important details to notice in this implementation:

1. Two separate learnable tokens. Both self.cls and self.distill are nn.Parameter objects initialized to zeros. During training, backpropagation will update them to learn useful representations. Despite identical initialization, they will diverge because they receive different gradient signals (ground-truth loss vs. teacher loss).

2. Positional embeddings. The self.pos tensor has shape (1, n+2, dim), providing a unique positional encoding for each of the 18 positions. This is essential because, unlike CNNs, the transformer has no built-in notion of spatial arrangement.

3. Two output heads. The head_cls and head_dist are separate linear layers that map the final token representations to class logits. Each head receives its own supervision signal during training.

4. Token placement. The CLS token is placed at position 0 and the distillation token at the last position. This is a convention; both tokens interact with all patch tokens through self-attention regardless of their position.

Implementing the knowledge distillation loss

Now we implement the loss function that drives the distillation process. This is where the mathematics we developed earlier becomes code.

Listing 1.7 Knowledge distillation loss function

def kd_loss(s_logits, t_logits, y, T=TEMPERATURE, alpha=ALPHA):
    kd = F.kl_div(                                        #A
        F.softmax(s_logits / T, dim=1),                   #B
        F.softmax(t_logits / T, dim=1),                   #C
        reduction='batchmean'
    ) * (T * T)                                           #D
    ce = F.cross_entropy(s_logits, y)                     #E
    return alpha * kd + (1 - alpha) * ce                  #F

#A KL divergence between the student’s and teacher’s softened distributions #B Student’s log-probabilities at temperature T (PyTorch’s kl_div expects log-probabilities for the first argument)

#C Teacher’s probabilities at temperature T (target distribution)

#D Multiply by T^2 to compensate for the reduced gradient magnitude

#E Standard cross-entropy between student predictions and ground-truth labels

#F Weighted combination: alpha controls the distillation-vs-classification balance

Let us trace through this function to make sure we understand each step. Given student logits z_s​, teacher logits z_t​, ground-truth labels y, temperature T=4, and α=0.1:

1. F.softmax(s_logits / T, dim=1) computes σ(zs/4): the student’s softened probabilities

2. F.softmax(t_logits / T, dim=1) computes σ(zt/4): the teacher’s softened probabilities

3. F.kl_div(...) computes D_KL(teacher∥student), the divergence between these distributions

4. Multiplying by T^2=16 compensates for the gradient scaling

5. F.cross_entropy(s_logits, y) computes L_CE at τ=1

6. The final loss is 0.1×(KL term)+0.9×(CE term)

With α=0.1, we weight the cross-entropy heavily (90%) and the distillation lightly (10%). This means the student primarily learns from the ground-truth labels, with the teacher’s knowledge providing supplementary guidance.

Training the student

With all components in place, we can now train the student model. The training loop follows the standard PyTorch pattern, but with the crucial addition of generating teacher predictions on the fly.

Listing 1.8 Training the DeiT student with knowledge distillation

print("Training student...")
for e in range(EPOCHS_STUDENT):
    student.train()
    for x, y in train_dl:
        x, y = x.to(device), y.to(device)

        with torch.no_grad():                            #A
            t_logits = teacher(x)

        cls_logits, dist_logits = student(x)             #B
        loss_cls = F.cross_entropy(cls_logits, y)  

        loss_distill = kd_loss(dist_logits, t_logits, y) #C
        loss = loss_distill + loss_cls        
        opt_s.zero_grad()                                #D
        loss.backward()
        opt_s.step()

    print(f"Epoch {e+1} done")

Output

Training student...
Epoch 1 done
Epoch 2 done
Epoch 3 done
Epoch 4 done
Epoch 5 done
Epoch 6 done
Epoch 7 done
Epoch 8 done
Epoch 9 done
Epoch 10 done

#A Teacher inference with no gradient tracking: the teacher is frozen and never updated

#B Student forward pass returns two sets of logits (CLS and distillation)

#C CLS loss trains the classification head with hard labels; distillation loss aligns the DIST token's logits with the teacher's soft logits. The total loss combines both.

#D Standard gradient descent: zero gradients, compute backpropagation, update weights

There is an important detail to highlight: we pass dist_logits (from the distillation token) to the loss function, not cls_logits. The distillation token is specifically designed to learn from the teacher. In a full DeiT implementation, you would compute a separate cross-entropy loss for the CLS token against ground truth and add it to the distillation loss. In our simplified version, the kd_loss function handles both terms using the distillation token.

Notice the torch.no_grad() context manager around the teacher’s forward pass. Since the teacher is frozen, we do not need to track gradients for its computations. This saves memory and computation: we only backpropagate through the student.

Evaluating the trained model

After training, we evaluate the student by combining the predictions from both tokens. At inference time, the CLS and distillation tokens each produce their own logits, and we average the softmax outputs to get the final prediction.

Listing 1.9 Evaluating the DeiT student

student.eval()
correct = 0
total = 0
samples = []

with torch.no_grad():                                     #A
    for x, y in test_dl:
        x, y = x.to(device), y.to(device)
        cls_logits, dist_logits = student(x)

        cls_dist = (cls_logits + dist_logits) / 2          #B
        pred = cls_dist.argmax(1)                          #C

        correct += (pred == y).sum().item()
        total += y.size(0)

        if len(samples) < 15:
            samples.append((x.cpu(), pred.cpu(), y.cpu()))

acc = 100 * correct / total
print(f"Test Accuracy: {acc:.2f}%")

#A No gradient computation needed during evaluation

#B Average the logits from both tokens: this combines the ground-truth-informed CLS prediction with the teacher-informed distillation prediction

#C Take the class with highest average logit as the final prediction

The averaging of CLS and distillation logits at inference is a key DeiT design choice. Each token has learned a different perspective on the input: the CLS token is optimized for ground-truth classification, while the distillation token is optimized for mimicking the teacher. Combining them yields a prediction that benefits from both information sources.

Finally, we can visualize some predictions to qualitatively assess the model:

Listing 1.10 Displaying sample predictions

fig, axs = plt.subplots(1, len(samples), figsize=(12, 3))
for i, (img, pred, true) in enumerate(samples):
    img = img[0].permute(1, 2, 0).numpy()                 #A
    axs[i].imshow(img)
    axs[i].set_title(f"P:{pred[0].item()} T:{true[0].item()}")  #B
    axs[i].axis('off')
plt.show()

Output

#A Convert from (C, H, W) tensor format to (H, W, C) NumPy array for matplotlib

#B Display the predicted label (P) and true label (T) for each sample

With this minimal implementation, we achieve approximately 94% accuracy on MNIST using only 6,000 training images and 100 epochs in just 18 mins of training on T4 GPU. While this is modest compared to state-of-the-art results, remember that we are using an extremely small model (16-dimensional embeddings, 4 layers, 4 heads) with very little data and training time. On the full ImageNet dataset with the proper hyperparameters from section 1.5, DeiT-B achieves 83.4% top-1 accuracy, competitive with models trained on 250 times more data.

The code demonstrates all the core DeiT concepts:

• Patch embedding via convolution

• CLS and distillation tokens as learnable parameters

• Dual classification heads with separate objectives

• Knowledge distillation loss combining cross-entropy and KL divergence

• Temperature scaling for softening probability distributions

• Averaging token predictions at inference

Summary

• Inductive biases are built-in architectural assumptions that shape how a model learns. CNNs have strong inductive biases (locality and translation equivariance) that make them data-efficient but limit their ability to capture long-range dependencies

• Vision Transformers lack these biases, processing all patches globally through self-attention. This flexibility becomes an advantage with enough data (300 million images) but a severe limitation with standard datasets (1.2 million images), causing overfitting and poor generalization

• Scaling laws show that model performance follows predictable power-law relationships with model size, dataset size, and compute. These laws hold for both language models and Vision Transformers, confirming that scale can compensate for lack of inductive bias

• Knowledge distillation transfers knowledge from a large teacher model to a smaller student by training the student on the teacher’s soft probability outputs rather than just hard labels. The teacher’s soft outputs contain dark knowledge (inter-class similarity information) that helps the student learn more efficiently

• Temperature scaling softens the teacher’s probability distribution, amplifying dark knowledge. The mathematical relationship σ(zi/τ)controls the softness, and the τ2 factor in the loss compensates for reduced gradient magnitudes

• DeiT introduces a distillation token alongside the standard CLS token, creating separate pathways for ground-truth learning and teacher imitation. The two tokens develop complementary representations that are combined at inference

• Hard distillation outperforms soft distillation in DeiT, achieving approximately 1% higher accuracy on ImageNet. This surprising result is attributed to the architectural mismatch between CNN teachers and transformer students

• DeiT achieves 83.4% top-1 accuracy on ImageNet using only ImageNet-1K (1.2 million images), matching models trained on 250 times more data, and can be trained in approximately 53 hours on a single 8-GPU node

Resources

Original Paper

https://arxiv.org/pdf/2012.12877

Dr Sreedath Panat has amazing videos on the same topic.

Some More Substacks

I’m also building Audio Deep Learning projects and Exploring and Finetuning different tts,sst models, sharing and discussing them on LinkedIn and Twitter. If you’re someone curious about these topics, I’d love to connect with you all!

Mayank Pratap Singh

LinkedIn : www.linkedin.com/in/mayankpratapsingh022

Twitter/X : x.com/Mayank_022.

1. What Is Quantization, Really?

2. The Outlier Problem

3. The Key Insight — Just Rotate It

4. A Concrete Example

5. The Second Trick — Fixing the Inner Product Bias

6. Why This Matters for LLM Inference

7. Connections and Prior Art

Every few months, a paper drops that makes the ML community collectively lose its mind. This month, it’s TurboQuant (Zandieh et al. 2025): a new vector quantization method from Google Research that achieves near-optimal compression of LLM weights and KV caches at 2.5–3.5 bits per parameter.

The Twitter discourse has been… colorful. “It’s just polar coordinates!” “It’s information theory!” “It’s magic!”

None of that is quite right. The core idea is shockingly simple, and I’m going to explain it from scratch: with visuals, concrete numbers, and zero hand-waving.

What Is Quantization, Really?

Before we get to TurboQuant, let’s make sure we’re on the same page about what quantization actually does.

You have a neural network. Every weight, every activation, every KV cache entry is a number — typically stored as a 16-bit floating point value (FP16). That’s 2 bytes per number.

A model like Llama 3 70B has 70 billion parameters. At FP16, that’s:

70B × 2 bytes = 140 GB

That doesn’t fit in a single GPU. Quantization is the art of making those numbers smaller.

The simplest possible quantization? Just round to fewer decimal places:

Original:    0.2374623   0.7237428   0.5434738   0.1001233
Quantized:   0.237       0.724       0.543       0.100

You’ve lost some precision, but you’ve saved memory. Real quantization schemes are more sophisticated — they map continuous values to a discrete set of levels — but the core idea is always: reduce the precision of each number to use fewer bits.

At 4 bits per weight, our 70B model becomes 35 GB. At 2 bits, it’s 17.5 GB. The question is: how low can you go before the model stops working?

The Outlier Problem

Here’s where things get interesting — and where TurboQuant enters the picture.

In an ideal world, the values in a neural network vector would be spread roughly evenly across their range. Something like:

Nice vector:  [0.24, 0.31, 0.18, 0.27, 0.22, 0.29, 0.20, 0.25]

Every component is in a similar range. If you quantize each one to 4 bits (16 levels), you can divide the range [0.18, 0.31] into 16 uniform buckets and represent each value with minimal error.

But real neural network vectors don’t look like that. They look like this:

Real vector:  [0.0001, 0.9999, 0.0002, 0.0001, 0.0003, 0.0001, 0.0002, 0.0001]

One component is enormous. The rest are near zero.

This phenomenon goes by many names in the transformer literature:

• Massive activations (Sun et al. 2024)

• Attention sinks (Gu et al. 2024)

• Outlier features (Dettmers et al. 2022)

Whatever you call it, the effect on quantization is devastating.

Why Outliers Kill Quantization

Think about what happens when you quantize that spiky vector. Your quantization grid has to span the full range [0.0001, 0.9999]. With only 16 levels (4 bits), each bucket covers about 0.0625 of that range.

The massive component (0.9999) maps cleanly to the top bucket — no problem.

But all the tiny components (0.0001, 0.0002, 0.0003) get crushed into the very first bucket. They all become 0. The quantized vector is essentially:

Quantized:  [0, 1, 0, 0, 0, 0, 0, 0]

This is a cardinal direction — a unit vector pointing along a single axis. It contains almost no information. An 8-dimensional cardinal direction can be described with just log₂(2×8) = 4 bits total, but we spent 4 × 8 = 32 bits to represent it.

We’re wasting bits, and we’ve lost all the subtle information that was encoded in the small components.

The Key Insight — Just Rotate It

Here is the entire key idea of TurboQuant, in one sentence:

Before quantizing a vector, randomly rotate it. After dequantizing, rotate it back.

That’s it. That’s the paper. (Well, most of it — there’s a clever second trick we’ll get to.)

Let me say that again, because it sounds too simple to be a major research contribution:

1. Take your vector

2. Multiply it by a random rotation matrix

3. Quantize the rotated vector

4. To dequantize: undo the quantization, then multiply by the inverse rotation

The rotation is data-independent. It’s chosen once and applied to everything. It doesn’t need to be learned or calibrated. It’s just a random orthogonal matrix.

But Why Does This Work?

This is the beautiful part, and it requires a little geometric intuition.

Remember our problem vector?

Spiky:  [0.0001, 0.9999, 0.0002, 0.0001]

This vector is nearly aligned with a coordinate axis — it points almost exactly in the direction of the second basis vector. Geometrically, it lives very close to one of the “poles” of the unit sphere.

Now imagine randomly rotating this vector. Where does it end up?

Almost certainly nowhere near any coordinate axis.

Think about it in 3D first. If you take a vector pointing straight up (the North Pole) and apply a random rotation, it could end up pointing in literally any direction. The chance of it landing near another pole is vanishingly small — because the poles occupy a tiny fraction of the sphere’s surface area.

After rotation, the components of our vector are spread out across all dimensions:

Before rotation: [0.0001,  0.9999,  0.0002,  0.0001]
After rotation:  [0.5012, -0.4998,  0.5001, -0.4989]

The outlier has been smeared across all components. Now each component has a similar magnitude, and our quantization grid can capture all of them with roughly equal fidelity.

The Geometry: Why Random Directions Are “Spread Out”

Here’s the deep geometric reason this works, and it gets stronger in high dimensions.

In d dimensions, the unit sphere has a surface area that is overwhelmingly concentrated away from the coordinate axes. As d grows, the fraction of the sphere near any axis shrinks exponentially. A randomly oriented vector in high dimensions will have components that are all roughly the same size — around 1/√d each.

More precisely, after rotation, each coordinate of a unit vector follows a Beta distribution that concentrates tightly around zero. When you square each coordinate, the sum must equal 1 (because the vector has unit length), so no single coordinate can be much larger than the others.

This is a fundamental property of high-dimensional geometry, and it’s the mathematical engine behind TurboQuant.

A Concrete Example

Let’s trace through the full TurboQuant process with actual numbers. We’ll use a small 4D example.

Step 1: Start With a Spiky Vector

x = [0.01, 0.98, 0.02, 0.01]

Norm: ‖x‖ ≈ 0.9807

Normalized: x̂ = [0.0102, 0.9993, 0.0204, 0.0102]

Step 2: Random Rotation

We apply a random orthogonal matrix R (generated once, shared across all vectors):

x_rotated = R × x̂ = [0.5127, -0.4831, 0.5042, -0.4998]

Notice: the energy is now spread evenly across all 4 dimensions. Each component has magnitude ≈ 0.5.

Step 3: Quantize

With a 2-bit quantizer (4 levels), our grid levels might be: {-0.75, -0.25, 0.25, 0.75}

x_quantized = [0.75, -0.75, 0.75, -0.75]

Hmm, that’s crude — but every component gets a meaningful representation. Compare this to quantizing the original spiky vector:

Original quantized: [0.0, 1.0, 0.0, 0.0]  ← cardinal direction!
Rotated quantized:  [0.75, -0.75, 0.75, -0.75]  ← much more info!

Step 4: Dequantize and Un-Rotate

x_dequantized = R⁻¹ × x_quantized × ‖x‖
             ≈ [0.03, 0.95, 0.04, 0.02]

Compare to original: [0.01, 0.98, 0.02, 0.01]

The reconstruction is far better than the naive approach!

The Second Trick — Fixing the Inner Product Bias

TurboQuant doesn’t stop at rotation. There’s a second, more subtle insight that matters specifically for attention computation.

In a transformer’s attention layer, we don’t just store KV cache vectors — we compute inner products between queries and keys:

score = q · k

It turns out that even when rotation makes quantization better in terms of mean squared error (MSE), it can introduce a systematic bias in inner products. The quantized vectors tend to produce inner products that are slightly “off” — not randomly off, but consistently biased in one direction.

The Fix: Residual Quantization with QJL

TurboQuant uses a two-step approach:

1. Quantize the rotated vector using an MSE-optimal scalar quantizer (the standard part)

2. Compute the residual (the error between the true vector and the quantized version)

3. Compress the residual using a 1-bit Quantized Johnson-Lindenstrauss (QJL) transform

The QJL step takes the residual error vector and represents it with just 1 bit per dimension using a random projection. This extra bit of information is enough to correct the bias in inner product computations.

The combined system — rotation + optimal quantizer + QJL residual — achieves inner product distortion that’s within a constant factor (≈2.7×) of the information-theoretic lower bound. That’s near-optimal.

Why This Matters for LLM Inference

All of this theory is nice, but why should you care? Because TurboQuant directly addresses one of the biggest bottlenecks in LLM inference: the KV cache.

During autoregressive generation, every token you’ve processed has a key vector and a value vector stored in the KV cache. For a model like Llama 3 70B generating a sequence of 8K tokens:

KV cache = 2 (K+V) × 80 layers × 8192 tokens × 8192 dim × 2 bytes
         ≈ 21 GB

That’s 21 GB of memory just for one request’s context! For batched serving with multiple users, this quickly becomes the dominant memory cost.

TurboQuant compresses this KV cache from 16 bits to 2.5–3.5 bits per value — a 5-6× reduction — with minimal quality degradation. The paper demonstrates:

• 3.5 bits/value: Near-lossless quality on standard benchmarks

• 2.5 bits/value: Slightly degraded but still remarkably good

• Training-free: No fine-tuning or calibration needed — just rotate, quantize, and go

This is what makes TurboQuant so exciting for production inference. It’s simple to implement, requires no retraining, and achieves near-optimal compression.

Connections and Prior Art

TurboQuant isn’t the first method to use rotation for quantization. Two notable predecessors:

QuIP (Chee et al. 2023) uses random orthogonal transformations for weight quantization, with a similar intuition about spreading outlier energy. However, QuIP uses the rotation as part of a more complex optimization procedure, while TurboQuant isolates the rotation as a standalone preprocessing step with clean theoretical guarantees.

RaBitQ (Gao et al. 2024) employs random rotation for vector database compression and nearest-neighbor search. TurboQuant extends this idea with the bias-correcting QJL step, which is critical for the inner product computations in attention.

What TurboQuant contributes beyond these is: (1) a clean theoretical analysis showing the rotation + quantizer is near-optimal, (2) the bias correction via QJL for inner products, and (3) a practical demonstration on LLM KV cache compression.

The Takeaway

Neural network vectors tend to be “spiky” — nearly aligned with coordinate axes. Quantization snaps spiky vectors to cardinal directions, destroying information. A random rotation spreads the energy evenly across all dimensions, making every component equally important and equally well-served by the quantization grid.

The fix is almost embarrassingly simple: multiply by a random rotation matrix before quantizing, and multiply by its inverse after dequantizing. No training. No calibration. No data dependence. Just linear algebra.

Combined with a 1-bit residual correction for inner product bias, this yields a quantization scheme that’s within a small constant factor of the information-theoretic limit.

Sometimes the most impactful ideas in ML aren’t the most complex ones. They’re the ones that see a fundamental geometric truth that was hiding in plain sight.

• The PreNorm dilution problem: Why standard residual connections cause hidden states to grow uncontrollably with depth, drowning out individual layer contributions

• The depth-time duality: The elegant insight that information dilution across depth is structurally identical to memory loss across a sequence, and can be solved the same way

• Attention Residuals (AttnRes): How replacing fixed accumulation with softmax attention over previous layer outputs gives each layer selective, input-dependent access to earlier representations

• Block AttnRes: The practical variant that partitions layers into blocks, reducing memory from O(Ld) to O(Nd) while preserving most gains

• Quantifying the gains: Scaling law experiments showing a 1.25x compute advantage, with benchmark improvements up to +7.5 points on GPQA-Diamond

This article assumes familiarity with transformers, self-attention, and residual connections. If you have read about how standard transformer blocks work, including the alternation of attention and feed-forward sublayers connected by residual streams, you have all the background you need.

Let’s begin with a roadmap of what we will cover, as shown in figure 1.

As shown in figure 1, we will progress through six stages. We start by understanding a hidden problem inside every modern LLM, build the conceptual insight that makes the fix possible, walk through the mechanism in detail, scale it to real-world models, prove it works mathematically, and finally measure the gains.

To understand why Attention Residuals matter, we first need to confront a hidden problem lurking inside every modern LLM.

The hidden flaw in residual connections

Residual connections are one of the most important innovations in deep learning. Introduced by He et al. in 2016, they allow gradients to flow through deep networks by providing shortcut paths around each layer. Every modern LLM, from GPT to LLaMA to DeepSeek, relies on them.

Yet these connections have a fundamental design limitation that has gone largely unaddressed for a decade. Every layer output is added to the running hidden state with a fixed weight of 1.0, creating an ever-growing signal that progressively drowns out individual layer contributions. This is the PreNorm dilution problem, and it wastes model capacity while limiting the effective utilization of depth.

The residual stream

Let’s examine how residual connections work in a standard PreNorm transformer, as illustrated in figure 2.

As shown in figure 2, each transformer layer receives the current hidden state, applies RMSNorm, processes it through attention and a feed-forward network, and then adds the result back to the running sum. The formula at each layer is straightforward:

The key observation is that every layer contributes its output with the same fixed weight of 1.0. There is no mechanism for one layer to contribute more or less than any other. The running sum, often called the “residual stream,” simply accumulates all outputs equally.

For our running example, we use four tokens, “The”, “cat”, “sat”, “down”, each with an 8-dimensional embedding. The input matrix X has shape (4, 8). We trace these tokens through 6 layers.

The magnitude growth problem

Since every layer adds its output with weight 1.0, the hidden state norm grows linearly with depth. This is not a minor bookkeeping detail. It is a fundamental property that shapes how the entire model behaves.

The hidden state at the final layer is literally the sum of all previous outputs:

Let’s see this growth in action with our running example, as shown in figure 3.

As shown in figure 3, the standard residual norm (orange line) grows steadily from 1.0 at the embedding layer to over 5.0 by layer 12. Meanwhile, the layer contribution fraction (gray bars) shrinks from 1.0 at layer 0 to less than 0.1 by layer 12. The RMSNorm before each sublayer normalizes the input, but the accumulated output keeps growing unchecked.

In a 50-layer model, the hidden state is the sum of approximately 50 vectors. Each layer’s contribution represents roughly 2% of the total signal. This is the PreNorm dilution problem.

The signal dilution effect

The consequence of this magnitude growth is that each layer’s voice gets progressively quieter relative to the total signal. Let’s visualize what the final hidden state actually looks like, as shown in figure 4.

As illustrated in figure 4, the final hidden state h_6 is the sum of 7 components: the token embedding x plus 6 layer outputs f_1 through f_6. Each contributes roughly 14% of the total signal. This means that layer 1, which might encode crucial low-level syntactic features, has the same influence as layer 6, which handles high-level reasoning. There is no selectivity.

This lack of selectivity creates a deeper problem. Both the attention sublayer and the FFN sublayer within each block receive the same blended signal, even though they may benefit from very different mixtures of earlier information.

Evidence of wasted capacity

The dilution problem is not just theoretical. Research has demonstrated that entire layers can be removed from deep LLMs with minimal performance impact, as shown in figure 5.

As shown in figure 5, removing 3 middle layers from a 12-layer model results in only a 3% performance drop. This suggests that the uniform accumulation of residual connections makes many intermediate layers effectively redundant. The model cannot efficiently utilize depth because each additional layer has diminishing marginal impact on the blended hidden state.

Now that we see the problem clearly, a natural question arises: can we fix this by applying the same trick that transformers already use to solve a structurally identical problem?

The depth-time duality

The conceptual breakthrough behind Attention Residuals comes from a simple but profound observation: the problem of information dilution across network depth is structurally identical to the problem of memory loss across a sequence of tokens. And we already know how to solve the sequence version.

The sequence analogy

Before transformers, recurrent neural networks processed sequences by passing a hidden state from one time step to the next. Early tokens were progressively forgotten as the sequence grew, because each new token blended into the same fixed-size hidden state. This “forgetting” problem limited how far back an RNN could effectively look.

The transformer’s self-attention mechanism solved this elegantly, as shown in figure 6.

As illustrated in figure 6, the transformer replaced the sequential blending of an RNN with selective retrieval via attention. Instead of each position receiving only the previous position’s output, any position can directly attend to any earlier position. The attention weights are content-dependent, meaning the model learns which earlier positions are most relevant for each new token.

Rotating attention from time to depth

Now here is the key insight. In the depth dimension of a transformer, the same “forgetting” problem occurs. Early layer outputs are progressively diluted as more layers are added, just as early tokens were forgotten in RNNs. The fix is structurally identical: let each layer directly access any previous layer with learned, content-dependent weights.

This is what the authors of Attention Residuals describe as “rotating” the attention mechanism 90 degrees, as shown in figure 7.

As shown in figure 7, both mechanisms compute h = sum(alpha_i * v_i), where the alpha weights are computed via softmax over learned scores. The only difference is the dimension: standard attention operates across sequence positions (horizontal), while AttnRes operates across network depth (vertical).

Let’s visualize this rotation more concretely, as illustrated in figure 8.

As shown in figure 8, the standard self-attention matrix (left) has tokens on both axes, with the causal mask ensuring each token only attends to previous tokens. The AttnRes pattern (right) has layers on both axes, with each layer attending only to previous layers. The mathematical structure is the same, just applied to a different dimension.

This is a profound reframing. The residual stream has always been performing a kind of “attention” over depth, but the crudest possible kind: every layer gets weight 1.0 regardless of content. AttnRes upgrades this to full softmax attention with learned, input-dependent weights.

With this duality in mind, let’s see exactly how Attention Residuals work, step by step.

The mechanics of Attention Residuals: a hands-on walkthrough

We now have the intuition: replace fixed-weight residual accumulation with softmax attention over previous layer outputs. Let’s build the mechanism piece by piece, starting with a review of what standard residuals compute and ending with the full AttnRes forward pass.

Standard residuals in action

Let’s first trace through our running example with standard residual connections to see the problem concretely. The side-by-side comparison in figure 9 frames what we are about to build.

As illustrated in figure 9, the standard approach (left) produces a hidden state h_4 with growing magnitude because all four layer outputs are summed with equal weight. The AttnRes approach (right) produces a bounded h_4 because the softmax weights, computed via pseudo-query vectors w_1 through w_4, always sum to 1.

Now let’s walk through the standard computation layer by layer, as shown in figure 10.

As shown in figure 10, our tokens start with embedding norm 1.0. After layer 1, the hidden state has norm 1.4 (the sum of the embedding and layer 1’s output). By layer 3, the norm has grown to 2.3. By layer 6, it reaches 3.2. Each layer contributes roughly 0.4 to the total norm, creating steady linear growth.

We can verify this growth pattern across all four of our tokens, as shown in figure 11.

As illustrated in figure 11, all four tokens exhibit the same monotonic growth pattern. The norms rise from 1.0 at the embedding layer to between 3.4 and 3.6 by layer 6. This confirms that the dilution problem is universal across all tokens, not specific to any particular input.

The pseudo-query vector

The first new concept in AttnRes is the pseudo-query vector. Each layer l has a learned parameter vector w_l of the same dimension as the hidden state (in our example, d = 8). This vector encodes the question: “which previous layers are most relevant for my computation?”

Let’s examine this concept in detail, as illustrated in figure 12.

As shown in figure 12, the pseudo-query w_l is a fixed learned parameter, not derived from the input like a standard attention query. However, the attention weights are still input-dependent because the “keys” come from previous layer outputs (v_i), which depend on the input tokens. Different inputs produce different keys, which produce different dot products with the same pseudo-query, which produce different attention weights.

What is a pseudo-query? A pseudo-query is a learned parameter vector that functions like a query in standard attention, but instead of being derived from the input, it is a fixed vector that each layer learns during training. It asks the same “question” for every input, but gets different “answers” because the keys (previous layer outputs) vary with the input.

Computing attention weights over depth

Now let’s see the full computation at a single layer. At layer 4, we have four previous outputs to attend over: v_0 (embeddings), v_1 (layer 1 output), v_2 (layer 2 output), and v_3 (layer 3 output). The detailed computation is shown in figure 13.

As shown in figure 13, the computation proceeds in four steps:

1. Normalize: Each previous output v_i is passed through RMSNorm to produce a key k_i

2. Score: The pseudo-query w_4 is dotted with each key to produce logits [0.2, 0.8, 1.5, 0.6]

3. Softmax: The logits are converted to attention weights that sum to 1: [0.08, 0.15, 0.55, 0.22]

4. Aggregate: The output is the weighted sum: h_4 = 0.08v_0 + 0.15v_1 + 0.55v_2 + 0.22v_3

The formal equation for the attention weights is:

Let’s zoom into the attention weight computation itself, as shown in figure 14.

As illustrated in figure 14, layer 4 selectively draws 55% of its input from layer 2’s output. Notice how different this is from standard residuals, where all layers contribute equally. The model has learned that layer 2’s representation is the most useful input for layer 4’s computation, and it can focus its attention accordingly.

The full AttnRes forward pass

Now let’s trace the complete forward pass with AttnRes through all 6 layers of our running example, as shown in figure 15.

As shown in figure 15, the hidden state magnitudes stay bounded near 1.0 throughout all 6 layers. This is a dramatic contrast with the standard residual version, where magnitudes grew to 3.2. The key difference is visible in the bar chart on the right: instead of linearly increasing bars, we see roughly uniform bars that stay below the dashed line marking the standard residual growth.

Let’s confirm this with a direct comparison of the magnitude curves, as shown in figure 16.

As illustrated in figure 16, the difference is stark. The standard residual norm (orange) grows from 1.0 to 3.45 over just 6 layers. The AttnRes norm (blue) stays essentially flat near 1.0. This is the mathematical magic of Attention Residuals: by replacing fixed-weight accumulation with a softmax-weighted combination, the hidden state magnitude is automatically bounded.

Input-dependent layer selection

The final piece of the puzzle is that different tokens produce different attention weight distributions. This is what makes AttnRes strictly more expressive than fixed-weight approaches like DenseFormer. Let’s see this in action, as shown in figure 17.

As shown in figure 17, the same pseudo-query w_4 produces three very different attention patterns depending on the token:

• Token “The”: A function word that attends mostly to the embedding layer (alpha_0 = 0.60). Function words carry relatively stable syntactic information that does not change much across layers.

• Token “cat”: A content word that focuses on intermediate layers (alpha_2 = 0.55). Content words benefit from the semantic processing performed by middle layers.

• Token “down”: A context-dependent word that attends most to recent layers (alpha_3 = 0.40). Its meaning depends heavily on the context built up by previous layers.

This is the key advantage over fixed-weight approaches. DenseFormer uses the same weights for every token; AttnRes adapts its layer selection to each token individually.

We now have the full AttnRes mechanism, but there is a practical challenge: storing all L layer outputs requires O(Ld) memory. For massive models with tens of billions of parameters, we need a more efficient approach.

Block AttnRes: scaling to real-world models

Full AttnRes is elegant but expensive. At layer l, we need to store all l-1 previous outputs, each of dimension d. For a 48-layer model with d = 8192, this represents significant memory overhead. In pipeline-parallel training, where different layers live on different GPUs, the situation is even more challenging.

The solution is Block AttnRes, which partitions layers into blocks and applies attention only at block boundaries.

The memory challenge

Let’s first understand the architecture at a high level, as shown in figure 18.

As shown in figure 18, Block AttnRes uses standard residuals within blocks (the familiar h_l = h_{l-1} + f_l pattern) but applies attention-based aggregation at the boundaries between blocks. The token embeddings x are always available as one of the attention sources, providing a direct path from input to any block.

The memory savings are substantial, as illustrated in figure 19.

As illustrated in figure 19, full AttnRes requires storing 48 separate representations for a 48-layer model. Block AttnRes with N = 8 blocks stores only 9 representations (8 block outputs plus the token embeddings), a 6x reduction. The authors found that 8 blocks is the optimal trade-off between performance and overhead.

How Block AttnRes works

Let’s walk through Block AttnRes with our running example: 6 layers partitioned into 2 blocks of 3 layers each. The step-by-step computation is shown in figure 20.

As shown in figure 20, the computation proceeds in three phases:

1. Block 1 (Layers 1-3): Standard residual connections. Each layer adds its output to the running sum. The block output B_1 = h_3 has a growing magnitude, just like in a standard transformer.

2. Block boundary: The attention mechanism computes weights over the token embeddings x and block 1’s output B_1. In this example, the weights are alpha_x = 0.35 and alpha_B1 = 0.65, producing h_boundary = 0.35x + 0.65B_1. This “resets” the magnitude by forming a convex combination.

3. Block 2 (Layers 4-6): Standard residual connections resume, using h_boundary as the starting point.

Notice the critical design insight: within blocks, O(L) growth is allowed to occur. But at block boundaries, the attention-based aggregation “resets” the magnitude back to a bounded value. This means the maximum growth within any single block is limited to O(L/N), which is much smaller than O(L).

Now let’s examine the block boundary computation in more detail, as shown in figure 21.

As illustrated in figure 21, the block boundary attention follows four stages:

1. Stack: Gather the token embeddings, all previous block outputs, and the current block’s partial sum into a tensor V of shape (N_prev+1, batch, seq, d)

2. Normalize: Apply RMSNorm to stabilize the attention computation

3. Score: Use the learned projection weight (the pseudo-query) to compute logits via einsum, then apply softmax to get weights

4. Aggregate: Compute the weighted sum to produce the input for the next block

Pipeline parallelism and system design

In large-scale training, models are split across multiple GPUs using pipeline parallelism. Each GPU holds a subset of layers, typically corresponding to one or two blocks. The challenge is that block boundary attention needs representations from blocks on other GPUs.

The solution is cache-based point-to-point communication, as shown in figure 22.

As shown in figure 22, each GPU stage maintains a cache of all block representations computed so far. When a block boundary is reached, the cached representations from previous stages are available for the attention computation. The two-phase strategy ensures efficient execution:

• Phase 1: Compute the standard attention and FFN operations within the current block

• Phase 2: Apply cross-block attention using the cached block representations

During inference, online softmax is used to amortize the computation cost, keeping the overhead minimal.

Having built the complete mechanism, both the idealized full version and the practical block variant, let’s formalize the mathematics and prove why AttnRes solves the dilution problem.

The mathematics of Attention Residuals

We have built intuition for why AttnRes works. Now let’s prove it. The mathematics reveals three key properties: bounded hidden states, improved gradient flow, and a beautiful connection between standard residuals and linear attention.

Bounding the hidden state

The most important mathematical property of AttnRes is that it produces a convex combination of previous outputs. Since softmax weights sum to 1, the output magnitude is automatically bounded, as shown in figure 23.

As illustrated in figure 23, the bound is elegantly simple:

Standard residuals:

AttnRes:

This follows directly from the triangle inequality and the fact that the alpha weights are non-negative and sum to 1. The output is a weighted average of the input vectors, and a weighted average can never have a larger magnitude than the largest input. This single property eliminates the entire PreNorm dilution problem, regardless of how many layers the model has.

Gradient flow and training dynamics

The second mathematical advantage concerns gradient flow during training. In standard residuals, gradients must propagate backward through a chain of multiplicative terms. In AttnRes, each layer receives a direct gradient signal.

Let’s compare the two gradient flow patterns. The standard residual gradient flow is shown in figure 24.

As shown in figure 24, the standard gradient formula involves a product of (L-l) Jacobian terms. While the identity component of each term helps prevent complete vanishing, the gradients still tend to concentrate in the layers closest to the loss, leaving early layers with weaker gradient signals.

Now let’s examine the AttnRes gradient flow, as shown in figure 25.

As illustrated in figure 25, AttnRes provides each layer with a direct gradient path:

Every layer receives a gradient signal weighted by its attention weight alpha. Layers that contribute more to the output get proportionally stronger gradients, creating a self-reinforcing learning dynamic: useful layers train faster, become even more useful, and receive even stronger gradients.

We can verify this improved gradient distribution quantitatively, as shown in figure 26.

As shown in figure 26, the standard transformer (orange) shows a severe gradient imbalance: layer 1 receives a gradient norm of only 0.15, while layer 8 receives 1.0. The AttnRes model (blue) maintains gradient norms between 0.82 and 0.97 across all layers. This near-uniform distribution means every layer can train effectively, regardless of its position in the network.

Standard residuals as low-rank linear attention

The paper establishes a beautiful theoretical result: standard residual connections are equivalent to low-rank linear attention over depth. This reframing reveals that the residual stream has always been performing attention over depth, just a very limited form of it. The comparison is shown in figure 27.

As illustrated in figure 27, the standard residual sum is mathematically identical to linear attention with constant weights over depth. AttnRes simply upgrades this to full softmax attention, providing strictly more expressive depth-wise information routing. The residual stream has always been doing a form of attention over depth; AttnRes just makes it a much better form.

The math confirms what our intuition suggested: AttnRes fundamentally solves the dilution problem. Now, let’s see how much this matters in practice.

Quantifying the gains

Theory is important, but the question practitioners care about is: does it work, and how much does it cost? The Kimi team evaluated Block AttnRes across scaling law experiments, benchmark evaluations, and architecture search, and the results are compelling.

Scaling law experiments

The most informative experiment compares validation loss at different compute budgets. The results are shown in figure 28.

As shown in figure 28, the AttnRes curve (blue) is consistently below the standard curve (orange) at every compute budget from 0.5 to 128 x 10^18 FLOPs. The key result is annotated on the plot: Block AttnRes matches baseline performance trained with 1.25x more compute. This means that adding AttnRes to your model is equivalent to increasing your compute budget by 25%, for free.

Benchmark results

The scaling law result is confirmed by benchmark evaluations on Kimi Linear, a production model with 48B total parameters (3B activated via mixture-of-experts), trained on 1.4 trillion tokens. The results are shown in figure 29.

As illustrated in figure 29, AttnRes delivers consistent improvements across all 9 benchmarks:

• GPQA-Diamond: +7.5 points (36.9 to 44.4). This is the standout result. GPQA-Diamond tests graduate-level reasoning, and the 7.5-point improvement suggests AttnRes significantly enhances the model’s ability to chain multi-step reasoning.

• Math: +3.6 points (53.5 to 57.1). Mathematical problem-solving requires precise, sequential reasoning across many steps, exactly the kind of task where selective layer access matters most.

• HumanEval: +3.1 points (59.1 to 62.2). Code generation benefits from the improved depth-wise routing, likely because different layers specialize in different aspects of code understanding.

• C-Eval: +2.9 points (79.6 to 82.5). Even on knowledge-focused benchmarks, the improvements are meaningful.

Deeper models, narrower architectures

AttnRes does not just improve performance at a fixed architecture. It fundamentally changes the optimal architecture itself, as shown in figure 30.

As shown in figure 30, standard residuals favor shallower, wider architectures because depth is poorly utilized (due to dilution). AttnRes shifts the optimal point toward deeper, narrower architectures because depth is now effectively utilized through selective attention. This is a fundamental shift in model design philosophy: with AttnRes, you can build deeper models that actually benefit from their depth.

Overhead analysis

The final question is cost. How much overhead does Block AttnRes add? The answer, shown in figure 31, is remarkably little.

As illustrated in figure 31, the overhead is minimal:

• Training cost increase: 3.8%. The additional attention computation at block boundaries adds less than 4% to the total training cost.

• Inference latency increase: 1.9%. Block attention is fast relative to the standard attention within layers.

The 1.25x compute advantage far outweighs the less than 4% overhead. For every dollar you spend on the attention computation, you get 25 cents of equivalent compute back in model quality. The authors confirmed that 8 blocks is the optimal trade-off between performance and overhead.

These improvements did not emerge from a vacuum. Attention Residuals represent the latest step in a decade-long quest to improve how deep networks combine information across depth.

The evolution of depth-wise aggregation

To appreciate where AttnRes sits in the landscape, it helps to trace the history of how researchers have approached the problem of combining information across depth.

From ResNet to AttnRes: a ten-year journey

The progression from fixed residual connections to learned depth-wise attention has unfolded over a decade, as shown in figure 32.

As shown in figure 32, each step in this timeline adds more expressiveness to depth-wise information routing:

• ResNet (2016): Fixed weight 1.0. Every layer contributes equally. Zero overhead, but no selectivity.

• DenseNet (2016): Concatenates all previous outputs instead of summing them. Very expressive but memory-intensive. Primarily used in computer vision.

• DenseFormer (Feb 2024): Learns scalar weights per layer pair, but these weights are input-independent. Minimal overhead, but limited expressiveness.

• ResFormer (Oct 2024): Adds residual connections specifically to value vectors from the first layer. Addresses the related but different problem of attention concentration.

• DeepCrossAttention (2025): Uses full input-dependent cross-attention between layers. Claims up to 3x training speedup. Most similar to AttnRes in concept.

• AttnRes (March 2026): Full softmax attention over depth via a single pseudo-query per layer. Less than 2% inference overhead with 1.25x compute advantage.

Comparing the approaches

A detailed comparison of the key methods is shown in figure 33.

As illustrated in figure 33, AttnRes occupies a unique position in the design space: it provides full input-dependent attention over depth with less than 2% inference overhead. The key distinguishing properties are bounded magnitudes (from the convex combination) and the 1.25x compute advantage (from better depth utilization).

When does AttnRes struggle?

No technique is universally superior, and AttnRes is no exception. Ziming Liu of MIT and Caltech provides a nuanced analysis of when AttnRes excels and when it falls short.

AttnRes excels on structured tasks where skipping intermediate layers is valuable. Natural language has rich hierarchical structure, with different layers specializing in syntax, semantics, and reasoning. The attention mechanism can learn to focus on specific layers without needing to suppress intermediate representations.

However, AttnRes can struggle on pure memorization tasks where uniform blending works fine. If the task requires every layer’s contribution equally, the selective attention provides no advantage and may even hurt by constraining the representation.

There is also a risk of representation collapse: if the attention weights converge to a uniform distribution during training, AttnRes degenerates to averaging all previous hidden states. This uniform bias can limit expressive capacity in certain settings. Natural language’s structured nature likely explains why Kimi’s strong empirical results generalize well across diverse benchmarks.

Having explored the theory, results, and context, let’s look inside a trained model to see what patterns the depth-wise attention actually learns.

Inside a trained model

The most illuminating way to understand AttnRes is to examine what the model actually learns. What patterns do the attention weights develop? How do different blocks interact?

What the attention weights learn

The attention weight heatmap from a trained Kimi Linear model reveals striking patterns, as shown in figure 34.

As shown in figure 34, three distinct patterns emerge from the learned attention weights:

• Early blocks attend to embeddings: Blocks 1 and 2 place strong attention on the token embeddings (source block 0). This makes sense because early processing needs direct access to the raw input features.

• Middle blocks attend to neighbors: Blocks 3 through 5 show a strong diagonal pattern, attending primarily to their immediate predecessors. This suggests a sequential refinement strategy where each block builds on the most recent computation.

• Deep blocks develop long-range connections: Blocks 7 and 8 show renewed attention to early blocks (columns 0-1), even though they also attend to nearby blocks. This is the most intriguing pattern: deep layers “reach back” to retrieve fundamental features that may have been diluted by intermediate processing.

These patterns would be impossible with standard residual connections, where every block receives the same blended signal regardless of what it needs.

Implementation walkthrough

The implementation of Block AttnRes is remarkably concise. The annotated pseudocode is shown in figure 35.

As illustrated in figure 35, the entire Block AttnRes computation fits in five steps:

def block_attn_res(block_reprs, partial_sum, proj):
    # Step 1: Stack all sources
    V = stack(block_reprs + [partial_sum])      # (N+1, batch, seq, d)

    # Step 2: Normalize
    V = RMSNorm(V)

    # Step 3: Compute attention logits
    logits = einsum('d, n b t d -> n b t',
                    proj.weight.squeeze(), V)

    # Step 4: Softmax over block dimension
    weights = logits.softmax(dim=0)

    # Step 5: Weighted aggregation
    output = einsum('n b t, n b t d -> b t d',
                    weights, V)

    return output

The function takes the previous block representations, the current block’s partial sum, and a learned projection (the pseudo-query). The stack, normalize, score, softmax, aggregate pattern is clean and efficient. Adding Block AttnRes to an existing transformer requires modifying only the block boundary logic, making it a true drop-in replacement for standard residual connections.

Summary

We have traced the complete story of Attention Residuals, from the problem they solve to the mechanism they use to the gains they deliver. Let’s recap the key takeaways.

• The PreNorm dilution problem: Standard residual connections add all layer outputs with fixed weight 1.0, causing hidden state magnitudes to grow as O(L) with depth. This progressively dilutes each layer’s contribution, making deeper layers less effective and wasting model capacity. In a 50-layer model, each layer represents only 2% of the final hidden state.

• The depth-time duality: Information dilution across network depth is structurally identical to memory loss across a sequence. Just as self-attention solved the sequence problem by enabling selective access to any previous position, AttnRes solves the depth problem by enabling selective access to any previous layer. The mathematical structure is the same, just rotated 90 degrees.

• The AttnRes mechanism: Each layer uses a learned pseudo-query vector to compute softmax attention weights over all previous layer outputs, producing an input-dependent weighted combination. Since softmax weights sum to 1, the resulting hidden state is a bounded convex combination that can never exceed the magnitude of the largest input. This completely eliminates the O(L) growth problem while providing each layer with selective, content-aware access to earlier representations.

• Block AttnRes for scalability: The practical variant partitions layers into approximately 8 blocks, using standard residuals within blocks and attention-based aggregation at block boundaries. This reduces memory from O(Ld) to O(Nd) while recovering most of the full AttnRes gains. Pipeline-parallel training is supported through cache-based point-to-point communication of block representations, with less than 2% inference overhead.

• Quantified improvements: Block AttnRes achieves a 1.25x compute advantage, matching baseline models trained with 25% more resources. On Kimi Linear (48B total, 3B activated, 1.4T tokens), it improved GPQA-Diamond by +7.5 points, Math by +3.6 points, and HumanEval by +3.1 points. The largest gains appear in multi-step reasoning tasks. AttnRes also shifts the optimal architecture toward deeper, narrower models that more effectively utilize their depth.

Further reading

• Attention Residuals (arXiv:2603.15031) - The original paper by Chen et al. at Moonshot AI (Kimi team), March 2026. Full mathematical formulation, scaling law experiments, and Kimi Linear integration.

• MoonshotAI/Attention-Residuals on GitHub - Official implementation with PyTorch pseudocode, Block AttnRes details, and benchmark tables.

• When Does Attention Residuals Work? (Ziming Liu) - Critical analysis from MIT/Caltech showing when AttnRes excels (structured tasks) and when it struggles (pure memorization).

• DenseFormer (arXiv:2402.02622) - The predecessor using Depth-Weighted Averaging with input-independent scalar weights. Important comparison point.

• DeepCrossAttention (arXiv:2502.06785) - Google’s learnable input-dependent cross-layer weights, claiming 3x training speedup.

• Value Residual Learning / ResFormer (arXiv:2410.17897) - Residual connections on value vectors to address attention concentration, saving 10-14% parameters.

Table of Contents

• Adapting transformers to images: patch embeddings and flattening

• Positional encodings in vision

• Encoder-only structure for classification

• Benefits and drawbacks of ViT

• Real-world applications of ViT

• Hands-on: fine-tuning ViT for image classification

Finetuning Vision Transformer Code is available below

https://github.com/VizuaraAI/Transformers-for-vision-BOOK

1.1 Introduction to Vision Transformers and Comparison with CNNs

Vision Transformers adapt the transformer architecture from language modeling to images. Instead of scanning an image with small sliding filters, they treat small image patches as tokens and learn how these patches relate to one another through self attention. For now it is enough to think of a Vision Transformer as a model that can look at all parts of an image at once and decide which regions should influence each other.

The key architectural difference between Vision Transformers and convolutional networks lies in how they see an image. A convolutional layer only looks at a small neighborhood of pixels when it computes its output. Its receptive field grows only gradually as we stack more layers and pooling operations. This locality bias has been very successful for classic vision tasks, but it means that long distance relationships in an image are only captured indirectly and late in the network. Figure 1 illustrates this contrast between global self attention and local convolution on a simple bird image.

Self attention in a Vision Transformer has a global receptive field from the first layer. For any query location, the model can compare it directly with every other patch in the image and decide which ones are relevant. In the bird illustration, a single pixel or patch can immediately connect to any other region in the picture, while a convolution sees only its nearby neighborhood and must rely on many stacked layers to move information from one side of the image to the other.

The second figure we have an optical illusion painting to make this difference more concrete. A human viewer can perceive both the detailed scene of a rider by the river and the larger face formed by the entire painting. A convolutional model tends to focus on the local textures of rocks, water, and fur, while a Vision Transformer can link distant regions that together form the face. Figure 2 shows how self attention attends across the whole painting, while a convolution still operates on a restricted local view.

The classic story of several blind people trying to describe an elephant gives another intuitive picture of this difference. Each person touches a single part and concludes that the elephant is a rope, a wall, or a snake, depending on whether they hold the tail, the body, or the trunk. A convolutional network behaves in a similar way, since each unit only has access to a small patch and builds its understanding from many separate local views. A Vision Transformer behaves more like a group that can share information freely. Even if each observer starts with a limited view, self attention lets them combine their observations and agree on the full shape of the elephant.

In the rest of this chapter we shift from this high level comparison to the inner workings of Vision Transformers. We will unpack how images are converted into patch tokens, how positional information is added, and how self attention layers process these tokens. Step by step, we will build up the full Vision Transformer encoder and its classification head so that the complete architecture becomes clear and concrete.

Since we've already covered the fundamentals of Transformers, I won't be going into too much detail here. If you're new to the architecture, I highly recommend reading my previous post, to get a solid foundation before we dive into Vision Transformers.

From Text Transformers to Vision Transformers

Transformers for text and transformers for images share the same core ideas. In language models such as GPT, we start from a sequence of tokens, embed them, and apply masked self attention so that each token can only attend to tokens at its current or earlier positions. This matches the next token prediction objective, where the final context vector of the sequence is used to predict the following token. In Vision Transformers, an image is first tokenized into a sequence of patches, self attention is applied without masking so that every patch can attend to every other patch, and a special class token provides a single representation that is passed to a small MLP head for classification.

BERT provides a second reference point for Vision Transformers. Instead of predicting the next word, BERT is trained to recover masked tokens inside a sentence, so it uses unmasked self attention over the entire sequence to capture bidirectional context. Vision Transformers adopt this encoder style design from BERT, but apply it to image patches together with a class token, giving a close analogue of BERT style sequence understanding in the image domain.

Now that we have compared transformers for text and images at a high level, we can focus on how a Vision Transformer actually sees an image. In the next section we will follow an image as it is cut into small patches, converted into vectors, and arranged into a sequence that looks very much like a sentence of tokens. This patch embedding stage is the first step that lets a standard transformer encoder operate directly on images.

1.2 Adapting transformers to images: patch embeddings and flattening

Adapting a transformer to images starts with a simple question: how can we turn a 2D image into the kind of 1D token sequence that a text transformer expects? Vision Transformers answer this by cutting the image into a grid of fixed size patches and treating each patch as a token. In this section we work through that patch embedding step using the 640 by 640 cat image and then show two common implementations: a flatten plus linear projection and an equivalent convolutional layer.

For the cat example, suppose the image has height and width 640 pixels and three color channels. We choose square patches of size 160 pixels. The image is divided into a 4 by 4 grid of non overlapping patches, each 160 by 160 by 3. In general, for an image of height H, width W, and patch size P, the number of patches N is

assuming H and W are multiples of P. For a square image of side S this reduces to

With H = W = 640 and P = 160 we obtain N = (640 / 160)² = 4² = 16 patches.

Once the image has been split into patches, each patch must be converted into a vector that lives in a common embedding space of dimension D, just as word tokens are embedded in text models. One straightforward approach, used in the original Vision Transformer paper, is to flatten each patch and apply a linear projection. A single patch has shape P by P by C, where C is the number of channels. Flattening gives a vector of length P²C.

In the transformer it is convenient to line these patches up in a fixed order and think of them as a one-dimensional sequence rather than a two-dimensional grid. For instance, we might start from the top-left patch, then move left to right, row by row, until we reach the bottom-right patch.

Patch embeddings without convolution

In the first subsection we will build patch embeddings without any convolutional layers. We will treat each 160×160×3 patch as a tiny image, flatten all of its pixels into a long vector, and pass that vector through a learnable linear layer to obtain a D-dimensional patch embedding. This view mirrors word embeddings in language models: each patch is simply another token whose embedding is learned directly from data.

To keep things concrete, return to the 640×640 cat image from the previous section. We divide it into a 4×4 grid of non-overlapping patches, each of size 160×160 pixels.

Each of these sixteen patches will become one token for the transformer.

Consider patch 10, the square around the cat’s eye. As an RGB patch it is a small 3-D tensor with shape

where C=3 is the number of colour channels. The image in Figure 7 shows this patch split into its red, green, and blue planes, and then shows a few individual pixels as (R,G,B) triplets. The first step of the non-convolutional patch embedding is to flatten this tensor into a single vector by concatenating all pixel values from all channels in a fixed order. If we denote the flattened version of the i-th patch by flat_patch_i then its length is

For a 160×160 RGB patch this means

numbers per patch. Flattening does not learn anything; it is just a reshaping operation that turns a 160×160×3 block of pixels into a vector of length 76,800.

A transformer, however, does not want raw pixel vectors of length 76,800. It expects a much shorter embedding vector of some dimension D, such as D=32 in our toy diagrams or D=768 in a ViT-Base model. The simplest way to obtain that is to apply a shared linear layer to every flattened patch. We introduce a weight matrix W_patch​ and a bias vector b_patch​ and define the patch embedding for the i-th patch as

with

The dimensions line up in the natural way: W_patch takes a length-P^2C vector and maps it down to a length-D vector, while b_patch shifts the result. In the cat example, if we choose D=32, then W_patch​ has shape 32×76,800 and b_patch​ has length 32. The output

is a D-dimensional patch embedding. The same parameters W_patch and b_patch​ are reused for every patch in every image, so they are heavily shared and trained end-to-end with the rest of the model. You can think of each row of W_patch​ as a learned template that looks at the entire patch and responds with a single number; stacking D such responses gives the embedding vector.

At this point we have turned one patch into one token. Repeating the same flatten-and-project operation for all N patches yields a collection of patch embeddings

We arrange them in a fixed, deterministic order, for example row by row across the image, and stack them into a matrix

This matrix is the direct analogue of the token-embedding matrix in a text transformer. The only difference is that each row now encodes an entire 160×160 coloured region of the image instead of a word or subword.

It is helpful to interpret this as describing the whole pipeline. Before embedding, the model sees N raw patches, each of shape 3×P×P. After flattening we conceptually have N vectors of length 3P^2. After the linear projection we have N patch embeddings of length D. Once we add the special class token and positional embeddings in the following sections, this becomes an (N+1)×D matrix that is finally presented to the transformer encoder.

Patch embeddings with convolution

In this subsection we will see that a single Conv2D layer can perform almost the same job in one shot. By choosing the kernel size and stride to match the patch size, a convolution turns the 640×640 image into a 4×4 grid of feature maps whose channel dimension is exactly our embedding dimension (for example, 32 channels). Each spatial location in this feature map is then interpreted as one patch token.

The idea is that a convolution with kernel size P and stride P can visit each patch exactly once, compress it, and write the output into a grid of size (H/P,W/P)

In our running example we use kernels of spatial size 160×160 and stride 160 The input is a tensor of shape (3,640,640). We apply a convolution layer whose RGB image we set C=3. The number of output channels is chosen to be our desired embedding dimension, for example D=32. Both the kernel size and the stride are set to 160, which means the kernel slides over the 640×640 image in non-overlapping 160×160 steps, and we use zero padding so that the image is neatly tiled into patches without any extra border pixels being added.

The convolution layer contains D separate kernels. Each kernel is a learnable weight tensor of shape (C,P,P)=(3,160,160). When the layer processes the image, each kernel slides over the image in steps of 160 pixels, producing one response per patch. Because the stride equals the kernel size, there is no overlap between neighbouring receptive fields. After the convolution, the output tensor has shape

Each spatial location (u,v) in this 4×4 grid corresponds to one patch of the original image. The vector of 32 numbers at that location comes from the 32 kernels, which together play the role of the rows of the linear matrix W_patch​ in the previous subsection. If we flatten the 4×4 grid of spatial locations into a length-16 sequence, and read out the 32-dimensional vector at each location, we obtain the same set of patch embeddings as before.

Taken together, these two constructions show that Vision Transformers do not depend on any particular patch-extraction trick: what matters is ending up with a sequence of N patch embeddings of dimension D. Whether those embeddings come from flattening plus a linear layer or from a carefully configured convolution is largely an implementation choice; once we have the N×D matrix of patch tokens, the rest of the Vision Transformer proceeds in exactly the same way.

Adding the class token and forming the sequence

So far we have obtained N patch embeddings x1,…,xN, each of dimension D. For classification tasks the Vision Transformer introduces one extra token, called the class token. This token does not come from any particular patch; it is a learned vector that is added to the front of the sequence and is meant to gather information from all other tokens through self-attention.

We denote the class-token embedding by

This vector is a trainable parameter, initialized randomly when we create the model and optimized along with all other weights. Once we have x_0 and the patch embeddings x1,…,x_N​, we can form the full sequence matrix.

The number of tokens entering the transformer encoder is therefore

In the 640 by 640 example with patch size 160, this gives N = 16 plus one class token, so there are 17 tokens in total, each of dimension D=32.

that will become the input to the Vision Transformer, once we add positional information.

1.3 Positional encodings in Vision Transformers

Self-attention in a Vision Transformer has no built-in notion of order. In the left panel of Figure 10, we could shuffle the cat patches so that a tile from the ear region swaps places with a tile from the plain background, and the encoder would happily process this jumbled sequence as if nothing were wrong. For images this is clearly problematic: a patch showing the cat’s eye carries very different meaning from a patch showing only empty purple background. To give the model a sense of where each token comes from in the original cat image grid, we add a positional embedding to every token before it enters the transformer encoder.

So after patch embedding and adding the class token we have a sequence of N+1 tokens, each of dimension D. We collect them in a matrix X, where x_0​ is the class token and x1,…,x_N are the patch embeddings. The Vision Transformer introduces a learnable positional embedding matrix

where each row p_i​ is a trainable vector that represents the position of token i. During training these vectors are updated like any other parameter in the model. For a mini batch of size B we broadcast the same positional matrix over the batch and form the final input to the encoder by simple elementwise addition

The tensor sent into the transformer encoder therefore has shape B×(N+1)×D. The sequence length and embedding size are unchanged, but each token now carries two kinds of information at once: the visual content of its patch and the location of that patch in the original grid.

1.4 Encoder-only structure for classification

In a Vision Transformer we only keep the encoder side of the original transformer architecture. The image is converted into a sequence of tokens and these tokens pass through a stack of L identical encoder blocks. Each block contains multi-head self attention, a feed-forward MLP, and residual connections with layer normalization, but there is no decoder that predicts future tokens. Instead, we prepend a single learnable class token to the sequence and treat the encoder as a feature extractor. After the last encoder block we read out only the final hidden state of this class token and feed it into a small MLP head that produces the class logits for the image. In this sense a ViT is an encoder-only model trained for classification.

The entire path from patch embeddings to context vectors

From the previous sections we already have a matrix of embedded tokens that combines patch information and positional information. We then add one extra token for classification, so the total sequence length is N+1. Each token has embedding dimension D. If we stack all token vectors row wise we obtain a matrix

With the class token we therefore have N+1= 17 tokens. If we choose an embedding dimension D=32, the matrix E has shape 17×32. This matrix is the input to the transformer encoder stack and from the encoder’s perspective it looks exactly like the token embeddings of a language model: a batch of sequences, each of length 17, each token represented by a 32-dimensional vector.

The encoder does not change the sequence length. After each encoder block we still have a matrix of shape (N+1)×D, but every row vector has been updated to incorporate information from all other tokens through self attention and the MLP. These updated vectors are what we call context vectors, because they encode both the content of a token and the context supplied by the other tokens in the sequence.

Transformer encoder and attention

To understand what happens inside one encoder block it is helpful to zoom into the self-attention sublayer. At the input of a block we have a matrix

where j denotes the depth of the block in the stack. The rows of this matrix are the current representations of the tokens. Self attention transforms this matrix into a new matrix of the same shape by letting every token look at every other token and decide how much to pay attention to it.

The first step is to project the token representations into three new spaces called queries, keys and values. Concretely, we multiply Z^(j) by three learnable weight matrices

where d_h​ is the head dimension for a single attention head. These weight matrices are shared across all positions in the sequence and are learned during training. Applying them gives three new matrices

each of shape (N+1)×dh​. Intuitively, the query vector q_i for token i encodes what that token is looking for in its context, the key vector k_i encodes what that token offers to others, and the value vector v_i​ encodes the actual information that will be blended into other tokens when they attend to it.

The second step is to turn queries and keys into attention weights. For a given query vector q_i​ we compute its similarity with every key k_j​ using a scaled dot product, which produces a scalar score for each pair of positions

The softmax function along the index j converts these scores into a probability distribution

so that

The coefficient α_ij​ can be read as “how much token i attends to token j”

The third step is to use these attention weights to blend the value vectors. For token i we take a weighted sum of all values

The vector

is the new representation of token i produced by this attention head. It contains a mixture of the value vectors of all tokens, with larger weights coming from positions that were judged more relevant by the scaled dot product. If a head learns to focus on the cat’s eye, for example, the value vectors from patches around the eye will receive larger coefficients when computing the context vector for the class token.

In practice a Vision Transformer uses multi-head attention rather than a single head. This means we repeat the procedure above several times in parallel with different sets of projection matrices

​Each head has its own head dimension d_h​, so after computing

for every head we concatenate the results and use another learned projection to return to the original embedding dimension D. This gives the attention sublayer output matrix of shape (N+1)×D, which is then passed through the MLP sublayer and residual connections to form the updated matrix

Context vectors and output dimensions

To keep track of how representations evolve through the encoder stack, it is useful to introduce a simple notation. Let

be the initial matrix of token embeddings after patch embedding and positional embedding. After the j-th encoder block we write

where

is the context vector for token i at depth j. The index i runs from 0 to N. When i=0 the vector corresponds to the class token. When i ≥ 1 it corresponds to one of the image patches. Because the encoder stack never changes the sequence length, every Z^(j) has exactly the same shape: (N+1)×D. For our 640×640 cat image with P=160 and D=32 that means each encoder block takes a 17×32 matrix as input and produces another 17×32 matrix as output.

After the final encoder block we obtain Z^(L). The most important vector in this matrix is z_0^(L)​, the last context vector of the class token. During training this vector has learned to aggregate information from all patch tokens through the repeated layers of self attention and MLPs. As a result it acts as a compact summary of the entire image. We feed z_0^(L)​ into a small MLP head that maps the D-dimensional vector to a vector of class logits, for example of dimension 1000 for ImageNet-1k. A softmax over these logits then gives a probability distribution over classes.

MLP head and classification

By the time the sequence has passed through the transformer encoder, all of the heavy lifting has already happened. Starting from our cat image, we created N = 16 patch tokens, added a single learnable class token at the beginning, and mapped everything into a D-dimensional embedding space. After L encoder blocks, we obtain the final sequence matrix

Each row

is a context vector for token i, where i=0 corresponds to the class token and i = 1,…,N correspond to the image patches. In our running toy example N+1=17 and D=32 , so Z^(L) has shape 17×32

For image classification we do not feed all seventeen context vectors into a separate network. Instead, we follow the original ViT design and use only the final context vector of the class token,

This vector has attended to every patch token at every encoder layer, so it acts as a learned summary of the entire image. Using a single summary vector keeps the architecture simple and keeps the number of parameters in the final classifier small. In principle one could pool or concatenate all patch context vectors, but this would increase the dimensionality of the classifier input and did not bring clear benefits in the ViT experiments.

The MLP head is an ordinary feed-forward classifier that takes

as input and outputs one logit for each class. In the simplest case it consists of a single linear layer with weight matrix

and bias vector

where C is the number of labels (for example, cat, dog, bird, and so on). The logits vector is then

Many practical ViT implementations insert a small two-layer MLP here instead of a single linear layer. In that case we first project

to a hidden dimension D_mlp​, apply a nonlinearity such as GELU, optionally apply dropout for regularisation, and then project down to C dimensions. The overall effect is to give the classifier a bit more capacity to reshape the representation coming from the transformer encoder before turning it into class scores.

The output y is a vector of unnormalised scores, or logits, one per class. At inference time we usually take the index of the largest logit as the predicted label. During training we pass y through a softmax to obtain a probability distribution over classes and compute a cross-entropy loss against the true label. Gradients from this loss flow back through the MLP head into the transformer encoder and further into the patch and positional embeddings, allowing the entire Vision Transformer to be trained end to end.

1.5 Benefits and drawbacks of ViT

Vision Transformers offer a conceptually clean and flexible alternative to convolutional networks by modeling images as sequences of tokens and relying entirely on self-attention to capture relationships between image regions. One of their key strengths is global context modeling: from the very first encoder layer, every image patch can attend to every other patch. This makes ViTs particularly effective at capturing long-range dependencies, such as relationships between distant parts of an object or interactions between foreground and background regions. In addition, the ViT architecture scales extremely well with data and model size. When trained on large-scale datasets, Vision Transformers often surpass convolutional networks in accuracy, showing that explicit convolutional inductive biases are not strictly necessary when sufficient data is available. Their architectural simplicity is another advantage: apart from the patch embedding stage, the model closely mirrors standard transformer encoders used in language models, making it easy to reuse ideas, optimizations, and tooling across vision and language domains.

However, these benefits come with important trade-offs. Vision Transformers are generally less data-efficient than convolutional networks, especially on small or medium-sized datasets. Without the strong locality and translation-equivariance biases of convolutions, ViTs must learn many visual regularities directly from data, which can lead to poorer performance when training data is limited. Self-attention also introduces higher computational and memory costs, as attention scales quadratically with the number of patches. For high-resolution images, this can quickly become a bottleneck. As a result, many practical ViT variants introduce hierarchical structures, windowed attention, or hybrid CNN–Transformer designs to mitigate these issues. In short, Vision Transformers excel when data and compute are abundant, but require careful design choices to remain competitive in more constrained settings.

1.6 Real-World Applications of Vision Transformers

Vision Transformers are now widely used across a broad range of real-world vision tasks, particularly in settings where large datasets and pretraining are available. In image classification, ViTs and their variants have become strong alternatives to deep convolutional networks, achieving state-of-the-art performance on large benchmarks when pretrained on massive image collections and fine-tuned on downstream tasks. Beyond classification, Vision Transformers have proven highly effective in object detection and image segmentation, where global context is especially valuable. Tasks such as detecting small objects in cluttered scenes or segmenting large, spatially distributed structures benefit from the ability of self-attention to relate distant patches directly.

In industrial and applied domains, Vision Transformers are increasingly used in medical imaging, remote sensing, and autonomous systems. In medical imaging, ViTs help model complex spatial relationships in high-resolution scans, such as MRI or histopathology images, where long-range dependencies can be diagnostically important. In satellite and aerial imagery, they are used for land-use classification, change detection, and large-scale scene understanding. Vision Transformers are also central to modern multimodal systems, where images must be aligned with text, audio, or other modalities. Models such as image–text encoders rely on ViT backbones to produce visual representations that integrate naturally with language transformers. As a result, Vision Transformers have become a foundational component in systems for image captioning, visual question answering, and large multimodal models, reinforcing their role as a unifying architecture across perception tasks.

1.7 Hands-on: fine-tuning ViT for image classification

Finetuning Vision Transformer Code Repo Link available below

https://github.com/VizuaraAI/Transformers-for-vision-BOOK

In this section we fine-tune a pretrained Vision Transformer on a real-world, high-resolution image classification task. The Oxford-IIIT Pet dataset provides sufficiently detailed visual structure to match the inductive biases of Vision Transformers, making it an ideal dataset for demonstrating practical ViT fine-tuning. We adapt a ViT-Base model pretrained on ImageNet and fine-tune it to classify pet images into breed categories, following a standard transfer-learning workflow used in modern vision systems.

Dataset and problem setup

The Oxford-IIIT Pet Dataset consists of 7,349 images of cats and dogs across 37 fine-grained breed classes. Images vary in resolution but are typically larger than 200×200 pixels and contain rich texture, shape, and spatial cues. Each image is labeled with a single breed, making this a multiclass classification problem. Although the dataset images are already reasonably high-resolution, pretrained Vision Transformers expect inputs of size 224×224, so we standardize all images to this resolution during preprocessing.

This dataset is a good fit for demonstrating ViT fine-tuning for several reasons:

• Fine-grained categories. Distinguishing between 37 pet breeds requires the model to attend to subtle visual differences in fur pattern, ear shape, and body proportion, exactly the kind of long-range spatial reasoning that self-attention handles well.

• Sufficient visual complexity. The images contain natural backgrounds, varying poses, and different lighting conditions, giving the model a realistic transfer learning challenge.

• Manageable size. With roughly 3,680 training images and 3,669 test images, the dataset is small enough to fine-tune on a single GPU in reasonable time, yet large enough to produce meaningful results.

Installing dependencies and setting constants

Before we write any model code, we install the required libraries and define the hyperparameters that will stay fixed throughout the experiment:

Listing 1.1 Installing dependencies and defining constants

!pip install torchmetrics -q
import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
import numpy as np

from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from transformers import ViTForImageClassification, ViTImageProcessor
from transformers import get_cosine_schedule_with_warmup

from torchmetrics.classification import MulticlassAccuracy
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay

from tqdm.auto import tqdm

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

torch.manual_seed(42)
torch.cuda.manual_seed_all(42)

NUM_CLASSES = 37
IMAGE_SIZE = 224
BATCH_SIZE = 32
EPOCHS = 50

We fix the random seed so that results are reproducible, and we set NUM_CLASSES = 37 to match the 37 pet breeds in the Oxford-IIIT dataset. The image size of 224 matches the resolution that the pretrained ViT-Base model expects.

Loading and exploring the dataset

We use torchvision.datasets.OxfordIIITPet to download and load the dataset. The dataset provides both a train-val split (used for training) and a test split (used for validation):

Listing 1.2 Loading the Oxford-IIIT Pet dataset

raw_train = datasets.OxfordIIITPet(
    root="./data",
    split="trainval",
    target_types="category",
    download=True
)

class_names = raw_train.classes
print(len(class_names))
print(class_names[:10])

this prints

37
['Abyssinian', 'Bengal', 'Birman', 'Bombay', 'British_Shorthair',
'Egyptian_Mau', 'Maine_Coon', 'Persian', 'Ragdoll', 'Russian_Blue']

Visualizing sample images

This listing displays a small grid of pet images to highlight the dataset’s visual diversity and resolution.

Listing 1.3 Visualizing sample images from the dataset

plt.figure(figsize=(12, 6))

for i in range(8):
    img, label = raw_train[i]
    plt.subplot(2, 4, i + 1)
    plt.imshow(img)
    plt.title(class_names[label])
    plt.axis("off")

plt.show()

output

From this visualization we observe that images contain rich textures, variable poses, and complex backgrounds. These properties make long-range attention over image patches particularly valuable.

Preprocessing and data loaders

Pretrained Vision Transformers are sensitive to the normalization statistics used during pretraining. We load the ViTImageProcessor to extract the correct mean and standard deviation, and then build separate transforms for training and validation:

Listing 1.4 Building preprocessing pipelines with ViTImageProcessor

processor = ViTImageProcessor.from_pretrained(
"google/vit-base-patch16-224"
)
print(processor)

Output

ViTImageProcessor {
  "do_convert_rgb": null,
  "do_normalize": true,
  "do_rescale": true,
  "do_resize": true,
  "image_mean": [
    0.5,
    0.5,
    0.5
  ],
  "image_processor_type": "ViTImageProcessor",
  "image_std": [
    0.5,
    0.5,
    0.5
  ],
  "resample": 2,
  "rescale_factor": 0.00392156862745098,
  "size": {
    "height": 224,
    "width": 224
  }
}

image_mean = processor.image_mean
image_std = processor.image_std

train_transforms = transforms.Compose([
    transforms.RandomResizedCrop(IMAGE_SIZE),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=image_mean, std=image_std)
])

val_transforms = transforms.Compose([
    transforms.Resize(IMAGE_SIZE),
    transforms.CenterCrop(IMAGE_SIZE),
    transforms.ToTensor(),
    transforms.Normalize(mean=image_mean, std=image_std)
])

The training transforms apply random resized cropping and horizontal flipping for data augmentation, while the validation transforms use a deterministic resize and center crop so that evaluation is reproducible. Both pipelines normalize using the ImageNet statistics that the pretrained model was trained with.

After preprocessing, every image has shape

which matches the ViT input specification.

We now construct the training and validation datasets and wrap them in PyTorch DataLoaders.

Listing 1.5 Constructing training and validation data loaders

train_dataset = datasets.OxfordIIITPet(
    root="./data",
    split="trainval",
    target_types="category",
    transform=train_transforms
)

val_dataset = datasets.OxfordIIITPet(
    root="./data",
    split="test",
    target_types="category",
    transform=val_transforms
)

train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)

Loading the pretrained model

We load a ViT-Base model pretrained on ImageNet using the Hugging Face transformers library. The key argument num_labels=NUM_CLASSES tells the library to replace the original 1000-class ImageNet head with a new linear head that outputs 37 logits,one per pet breed.

The ignore_mismatched_sizes=True flag suppresses the warning about the size mismatch in the classification head:

Listing 1.6 Loading a pretrained ViT-Base model with a new classification head

model = ViTForImageClassification.from_pretrained(
    "google/vit-base-patch16-224",
    num_labels=NUM_CLASSES,
    ignore_mismatched_sizes=True
).to(device)

print(model)

The output model structure would be

ViTForImageClassification(
  (vit): ViTModel(
    (embeddings): ViTEmbeddings(
      (patch_embeddings): ViTPatchEmbeddings(
        (projection): Conv2d(3, 768, kernel_size=(16, 16), stride=(16, 16))
      )
      (dropout): Dropout(p=0.0, inplace=False)
    )
    (encoder): ViTEncoder(
      (layer): ModuleList(
        (0-11): 12 x ViTLayer(
          (attention): ViTAttention(
            (attention): ViTSelfAttention(
              (query): Linear(in_features=768, out_features=768, bias=True)
              (key): Linear(in_features=768, out_features=768, bias=True)
              (value): Linear(in_features=768, out_features=768, bias=True)
            )
            (output): ViTSelfOutput(
              (dense): Linear(in_features=768, out_features=768, bias=True)
              (dropout): Dropout(p=0.0, inplace=False)
            )
          )
          (intermediate): ViTIntermediate(
            (dense): Linear(in_features=768, out_features=3072, bias=True)
            (intermediate_act_fn): GELUActivation()
          )
          (output): ViTOutput(
            (dense): Linear(in_features=3072, out_features=768, bias=True)
            (dropout): Dropout(p=0.0, inplace=False)
          )
          (layernorm_before): LayerNorm((768,), eps=1e-12, elementwise_affine=True)
          (layernorm_after): LayerNorm((768,), eps=1e-12, elementwise_affine=True)
        )
      )
    )
    (layernorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)
  )
  (classifier): Linear(in_features=768, out_features=37, bias=True)
)

Freezing the backbone and training only the head

A common transfer-learning strategy is to freeze all pretrained parameters and train only the newly initialized classification head. This is fast, requires little memory, and often produces strong results when the pretrained features already capture the visual concepts needed for the target task:

Listing 1.7 Freezing the backbone and counting trainable parameters

for param in model.parameters():
    param.requires_grad = False
for name, param in model.named_parameters():
    if "classifier" in name:
       param.requires_grad = True

We can verify the freeze with a quick parameter count:

def print_model_parameters(model):
    trainable_params = 0
    frozen_params = 0
    all_param = 0

    for _, param in model.named_parameters():
        num_params = param.numel()
        all_param += num_params

        if param.requires_grad:
            trainable_params += num_params
        else:
            frozen_params += num_params

    print(f"trainable params: {trainable_params:,}")
    print(f"frozen params:    {frozen_params:,}")
    print(f"all params:       {all_param:,}")
    print(f"trainable%:       {100 * trainable_params / all_param:.2f}%")

# Run the function
print_model_parameters(model)

The output shows that only the classification head is trainable, roughly 28,000 parameters out of the model’s 86 million total:

trainable params: 28,453
frozen params: 85,798,656
all params: 85,827,109
trainable%: 0.03%

Sanity check: pre-training inference

Before any fine-tuning, we run the model on a single validation image to establish a baseline. Since the classification head is randomly initialized, we expect the prediction to be essentially random:

Listing 1.8 Sanity check: pre-training inference on one image

model.eval()

image, label = val_dataset[0]
image = image.unsqueeze(0).to(device)

with torch.no_grad():
    outputs = model(pixel_values=image)
    logits = outputs.logits
    probs = torch.softmax(logits, dim=1)
    pred = probs.argmax(dim=1).item()
    confidence = probs.max(dim=1).values.item()

print(
    f"[Pre-training inference]\n"
    f"  Ground truth class : {class_names[label]}\n"
    f"  Predicted class    : {class_names[pred]}\n"
    f"  Prediction confidence : {confidence:.2f}"
)

Output

[Pre-training inference]
  Ground truth class : Abyssinian
  Predicted class    : Shiba Inu
  Prediction confidence : 0.05

The model predicts an incorrect class with low confidence, confirming that the head needs training.

Setting up the optimizer, scheduler, and loss

We use AdamW with a learning rate of 3 × 10−4 and a cosine schedule with linear warmup. The warmup phase helps stabilize early training when the head weights are still random:

Listing 1.9 Setting up AdamW optimizer, cosine scheduler, and loss function

optimizer = optim.AdamW(
    filter(lambda p: p.requires_grad, model.parameters()),
    lr=3e-4,
    weight_decay=1e-4
)


total_steps = len(train_loader) * EPOCHS
warmup_steps = int(0.1 * total_steps)

scheduler = get_cosine_schedule_with_warmup(
    optimizer,
    num_warmup_steps=warmup_steps,
    num_training_steps=total_steps
)

criterion = nn.CrossEntropyLoss()
accuracy = MulticlassAccuracy(num_classes=NUM_CLASSES).to(device)

AdamW for fine-tuning

AdamW is a variant of Adam that decouples weight decay from the gradient update. This prevents the regularization from interfering with the adaptive learning rate, leading to better generalization. It is the standard optimizer for both pretraining and fine-tuning transformers.

The training loop

The training loop follows a standard PyTorch pattern: for each epoch, iterate over mini-batches, compute the cross-entropy loss, back-propagate, and update the classification-head weights. After each epoch, we evaluate on the validation set:

Listing 1.10 The main training and validation loop

train_losses, val_accuracies = [], []

for epoch in range(EPOCHS):
    model.train()
    running_loss = 0.0

    for imgs, labels in tqdm(train_loader):
        imgs, labels = imgs.to(device), labels.to(device)

        outputs = model(pixel_values=imgs)
        loss = criterion(outputs.logits, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        scheduler.step()

        running_loss += loss.item()

    train_losses.append(running_loss / len(train_loader))

    model.eval()
    accuracy.reset()
    with torch.no_grad():
        for imgs, labels in val_loader:
            imgs, labels = imgs.to(device), labels.to(device)
            preds = model(pixel_values=imgs).logits.argmax(dim=1)
            accuracy.update(preds, labels)

    val_accuracies.append(accuracy.compute().item())

    print(f"Epoch {epoch+1}: "
          f"Loss={train_losses[-1]:.4f}, "
          f"Val Acc={val_accuracies[-1]:.4f}")

Note how we call model.train() at the start of each epoch to enable dropout and batch-norm updates, and model.eval() before validation to disable them. The scheduler.step() call happens after each optimizer step (not each epoch), which is the correct behavior for the cosine-with-warmup schedule.

Plotting training progress

After training, we plot the training loss and validation accuracy curves to assess convergence:

Listing 1.11 Plotting training loss and validation accuracy curves

plt.figure(figsize=(12, 4))

plt.subplot(1, 2, 1)
plt.plot(train_losses)
plt.title("Training Loss")
plt.xlabel("Epoch")
plt.ylabel("Loss")

plt.subplot(1, 2, 2)
plt.plot(val_accuracies)
plt.title("Validation Accuracy")

plt.xlabel("Epoch")
plt.ylabel("Accuracy")
plt.tight_layout()
plt.show()

Output plot

A healthy run will show the training loss decreasing steadily and the validation accuracy climbing over the first several epochs before leveling off. Since we are only training the classification head, convergence is typically fast.

Post-training inference

We can now verify that the fine-tuned model makes correct predictions on validation images:
Listing 1.12 Post-training inference on a validation image

model.eval()

image, label = val_dataset[0]
image = image.unsqueeze(0).to(device)

with torch.no_grad():
    logits = model(pixel_values=image).logits
    pred = logits.argmax(dim=1).item()

print("After training → Pred:", class_names[pred],
      "| GT:", class_names[label])

Output

After training → Pred: Abyssinian | GT: Abyssinian

Confusion matrix evaluation

To understand where the model excels and where it struggles, we compute a confusion matrix over the entire validation set. This reveals which breed pairs are most easily confused:

Listing 1.13 Computing and displaying the confusion matrix

all_preds, all_labels = [], []

model.eval()
with torch.no_grad():
    for imgs, labels in val_loader:
        imgs = imgs.to(device)
        # Get predictions
        preds = model(pixel_values=imgs).logits.argmax(dim=1)

        all_preds.extend(preds.cpu().numpy())
        all_labels.extend(labels.numpy())

cm = confusion_matrix(all_labels, all_preds)
disp = ConfusionMatrixDisplay(cm, display_labels=class_names)


fig, ax = plt.subplots(figsize=(30, 30))

disp.plot(ax=ax, xticks_rotation=45, colorbar=True)


plt.tight_layout()

plt.show()

Output

The diagonal entries show correct predictions; off-diagonal entries indicate confusions. Breeds with visually similar features (for example, different shorthaired cats) will typically show higher off-diagonal values. This kind of finegrained analysis is valuable for deciding whether to invest in more data, stronger augmentation, or a larger model.

Saving the fine-tuned model

Finally, we save the fine-tuned weights so that the model can be reloaded later for inference without retraining:

Listing 1.14 Saving the fine-tuned model weights

torch.save(model.state_dict(), "vit_finetuned_final.pth")
print("Model saved successfully.")

The .pth file contains only the model’s state_dict, a dictionary mapping each layer name to its parameter tensor. To reload the model, we would create a new ViTForImageClassification instance with the same configuration and call model.load_state_dict(torch.load("vit_finetuned_final.pth")).

Resources

Original Paper

An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

Dr Sreedath Panat

Vision Transformer Paper Dissection

Build Vision Transformer from Scratch

1.8 Summary

• Vision Transformers adapt the transformer architecture to images by treating fixed-size image patches as tokens, enabling global self-attention from the first layer. Unlike convolutional networks, which build receptive fields gradually through stacked layers, ViTs can relate any two image regions directly.

• Patch embedding converts a 2D image into a 1D sequence of token vectors. This can be done by flattening each patch and applying a linear projection, or equivalently by using a single convolution with kernel size and stride equal to the patch size.

• A learnable class token is prepended to the sequence and accumulates information from all patches through self-attention. After the final encoder block, the class token’s context vector serves as a compact summary of the entire image and is passed to an MLP head for classification.

• Learnable positional embeddings are added to each token so that the model retains spatial information about where each patch originated in the original image grid.

• The encoder-only architecture processes the full patch sequence through L identical blocks of multi-head self-attention and feed-forward layers. Each block preserves the sequence length and embedding dimension, progressively refining token representations.

• Vision Transformers scale well with large datasets and model sizes but are less data-efficient than CNNs on small datasets. Practical variants address the quadratic attention cost through hierarchical designs and windowed attention.

• Fine-tuning a pretrained ViT on a downstream classification task follows a standard transfer-learning workflow: freeze the pretrained backbone, replace the classification head, and train only the head on the target dataset using a cosine learning rate schedule with warmup.

Some of More Substacks

I’m also building Audio Deep Learning projects and Exploring and Finetuning different tts,sst models, sharing and discussing them on LinkedIn and Twitter. If you’re someone curious about these topics, I’d love to connect with you all!

Mayank Pratap Singh

LinkedIn : www.linkedin.com/in/mayankpratapsingh022

Twitter/X : x.com/Mayank_022.

The Transformer Architecture

• 1.1 Introduction to Large Language Models

• 1.2 Anatomy of the transformer block

• 1.3 Tokenization

• 1.4 Byte Pair Encoding

• 1.5 Word Embedding

• 1.6 Transformer Block

• 1.7 The Need for Attention Mechanism

• 1.8 Self Attention Mechanism

• 1.9 Understanding the Input Embedding Matrix

• 1.10 From Embeddings to Queries, Keys & Values

• 1.11 A Quick Note on Matrix Multiplication

• 1.12 Why Scale Attention Scores?

• 1.13 Causal & Masked Attention

• 1.14 Causal Attention with Dropouts

• 1.15 Summary of Self-Attention

• 1.16 Intuition of Multi-Head Attention

• 1.17 Layer Normalization

• 1.18 FeedForward Network

• 1.19 Shortcut connections

• 1.20 Why Transformers Scale Better Than RNNs and CNNs

• 1.21 Pretraining, Fine Tuning, and Transfer Learning in Transformers

• 1.22 Limitations and Challenges of Transformers

• 1.23 Hands On Coding a Miniature Transformer for Sequence Classification

• 1.24 Summary

You can fine the code notebook here

https://github.com/VizuaraAI/Transformers-for-vision-BOOK

1.1 Introduction to Large Language Models

Figure 1.1 A Large Language Model takes a sequence of words as input and predicts the most likely next word, generating text one token at a time.

Large Language Models are neural networks trained on vast text datasets to perform a fundamental task: predicting the next word in a sequence. This simple objective drives the sophisticated capabilities we see in systems like GPT and ChatGPT.

Figure 1.2 Autoregressive text generation. The model predicts the next word, appends it to the input, and repeats the process to produce entire paragraphs.

When you interact with an LLM, it generates responses one word at a time. Given a prompt like “The cat sat on the,” the model predicts the next word, perhaps “mat.” This word is added to the sequence, becoming “The cat sat on the mat,” which then serves as input for predicting the following word. Through this iterative process, LLMs produce entire paragraphs and complex responses.

The transformer architecture enables these models to consider both immediate context and long range dependencies throughout the input sequence, maintaining coherence across extended text generation.

Despite the apparent simplicity of next word prediction, this mechanism gives rise to remarkable language understanding and generation capabilities. Understanding how transformers accomplish this task is essential to grasping how modern language models work.

Predicting the Next Word with OpenAI’s LLM

Let’s see a simple example to see how an LLM predicts the next word given a partial sentence:

You can refer to the full source code notebook for this exercise on Colab.

Predicting_the_next_word_notebook

Using the given code, we can predict the next word in a sentence based on probabilities assigned by a Large Language Model (LLM). Let’s say our Input sentence is

“After years of hard work, your effort will take you”

Figure 1.3 Input sentence fed to the LLM for next-word prediction.

if you will observe the top 10 predicted next words along with their probabilities (refer the notebook)

Figure 2.4 Top 10 predicted next words and their probabilities. The token “to” dominates at 90.7%, reflecting the most natural continuation.

The probabilistic nature of Large Language Models becomes clear when examining how they rank potential next words. The first token, such as “to,” might have the highest probability at 90.7 percent because it represents the most natural continuation based on the given context. As we look at alternative word choices, the probabilities gradually decrease, with each subsequent option representing a less common but still valid completion.

This distribution reveals the fundamental mechanism of Large Language Models: they function as probabilistic engines, predicting the most likely next token based on learned patterns. Rather than selecting a single correct answer, LLMs evaluate every possible next word and assign likelihood scores based on the vast patterns learned during training. This probabilistic approach enables models to generate diverse, contextually appropriate text while maintaining flexibility in their outputs.

Why is There “Large” in LLMs?

The term “Large” in Large Language Models reflects a fundamental principle: size directly impacts performance. Scaling laws show that model capabilities improve predictably with more parameters, enabling complex tasks like reasoning and code generation that smaller models cannot perform. Most critically, emergent properties such as arithmetic reasoning and multilingual understanding appear only when models cross certain size thresholds. This relationship between scale and capability explains why billions of parameters are essential for achieving sophisticated language understanding.

Figure 1.5 Scaling laws demonstrate a predictable relationship between model size and performance across a range of benchmarks

LLMs have billions to trillions of parameters. The first major paper to explore scaling laws was the GPT-3 paper (Language Models are Few-Shot Learners). The research demonstrated that as we increase the model size, from 1.3B parameters to 13B to 175B, the model’s performance dramatically improves.

Figure 1.6 Exponential growth in the size of language models from the 1950s to today. Orange dots represent language models, some of which have already crossed one trillion parameters.

Over the years, we have seen an exponential increase in the size of LLMs, from the 1950s to today. In the above graph, the orange dots represent language models, showing how their size has increased drastically over time. Some models have already crossed 1 trillion parameters!

Why do we care about the size of LLMs?

The size of Large Language Models matters primarily because of emergent properties: abilities that are absent in smaller models but spontaneously appear when models reach certain scales. These emergent capabilities fundamentally distinguish large models from their smaller counterparts. As LLMs grow beyond specific parameter thresholds, they suddenly acquire skills like solving complex arithmetic equations, translating between languages with nuanced understanding, and unscrambling letters into meaningful words. These abilities do not gradually improve with size but rather emerge abruptly at particular scales, making model size not just a technical detail but a critical factor in determining what tasks an LLM can perform.

Figure 1.7 Emergent abilities of large language models. Performance on certain tasks remains near zero until the model reaches a critical size, after which accuracy jumps sharply.

In the figure above , the X-axis represents model size (or computational power), and we can observe a pickup point, a stage where models suddenly start performing significantly better at these tasks. Emergent Abilities of Large Language Models

Figure 1.8 At larger scales, LLMs move beyond simple word prediction to excel at specialized tasks such as multilingual translation, text summarization, and grammar correction.

At larger scales, LLMs transcend simple word prediction to excel at specialized tasks like multilingual translation, text summarization, and grammar correction. This evolution from basic prediction to complex language understanding drives the race to build increasingly larger models. The direct correlation between parameter count and performance across diverse NLP tasks makes scale a critical competitive advantage.

1.2 Anatomy of the transformer block

Figure 1.9 The original transformer architecture from the “Attention Is All You Need” paper, consisting of an encoder stack on the left and a decoder stack on the right, connected through cross-attention.

The transformer architecture, introduced in the groundbreaking 2017 paper “Attention Is All You Need,” revolutionized artificial intelligence and natural language processing. This paper, now with over 200,000 citations, proposed the concept of self attention, fundamentally changing how we implement NLP systems. The transformer architecture consists of two main components: encoders and decoders. Encoder architectures power models like BERT, while decoder architectures form the basis of GPT and ChatGPT.

At the heart of modern LLMs lies this transformer architecture, which replaced traditional models like LSTMs and GRUs with self attention mechanisms. This innovation brought crucial advantages: the ability to capture long range dependencies in text, parallel processing that enables faster training, and unprecedented scalability that allows building increasingly powerful models. Understanding how the decoder portion works essentially reveals how GPT models function, as they are decoder only architectures.

The transformer block itself contains several key components working in sequence. Input text is first tokenized and converted to embeddings, which are then combined with positional encodings. These flow through layers of multi head attention, normalization, and feed forward networks, with dropout applied for regularization. The output layer finally produces logits for next token prediction. While the complete architecture diagram may appear complex with its numerous modules and connections, each component serves a specific purpose in transforming input text into meaningful predictions.

Figure 1.10 A simplified decoder-only transformer showing the main components: token and positional embeddings, transformer blocks with multi-head attention, feed-forward networks, layer normalization, and dropout, followed by the output layer.

The decoder only architecture, which powers models like GPT, can be understood by examining a simplified version of the transformer’s decoder component. While the complete architecture may appear complex with numerous interconnected modules, we can break it down into three manageable parts for clarity. This modular approach allows us to examine each component systematically rather than attempting to grasp the entire system at once. By focusing on these three core sections sequentially, we can build a comprehensive understanding of how the decoder transforms input text into predictions.

The three parts of an LLM’s architecture are Input, Processing, and Output.

Figure 1.11 The three stages of an LLM: the Input stage (tokenization and embeddings), the Processing stage (transformer blocks), and the Output stage (linear layer and softmax for next-token prediction).

So everything begins with the input stage, where several key transformations take place before it enters the processing unit, commonly known as the Transformer block.

First, the raw text undergoes tokenization, a process where the sentence is broken down into smaller units called tokens, these could be words, subwords, or characters, depending on the tokenization method used. This step ensures that the model can handle language efficiently,

Figure 1.12 The input pipeline: raw text is tokenized into subword units, each token receives a numerical embedding, and positional embeddings are added to encode sequence order.

Next, each token is converted into a numerical representation through token embeddings. These embeddings assign a unique vector to each token, capturing semantic meaning and relationships between words. However, since token embeddings alone do not preserve the sequence order, we introduce positional embeddings. These embeddings encode the position of each token within the sentence, allowing the model to understand the order and structure of the input.

With tokenization, token embeddings, and positional embeddings in place, the input is now fully prepared for the Transformer block, where deep learning mechanisms, such as multi-head attention and feed-forward neural networks, process the text to generate meaningful predictions.

1.3 Tokenization

Before text enters a transformer model, it undergoes tokenization, a process that converts raw text into tokens which are then assigned unique IDs. There are three main tokenization approaches, each with distinct characteristics.

Figure 1.13 Three tokenization strategies applied to the word “tokenization”: word-based (one token per word), character-based (one token per character), and subword-based (meaningful subword units).

Word based tokenization treats each complete word as a separate token, creating a dictionary of all words in the vocabulary. While intuitive, this approach struggles with vocabulary size and cannot handle new or misspelled words effectively. Character based tokenization breaks text down to individual characters, making each character a token. This creates a very small vocabulary but produces extremely long sequences that are computationally expensive to process.

Subword based tokenization, the preferred method for modern LLMs, breaks words into meaningful subword units.

Figure 1.14 Subword tokenization example: the word “playground” splits into “play” and “ground,” each a reusable meaningful unit.

A subword is a smaller meaningful unit that can be reused across different words. For example, “playground” might split into “play” and “ground,” while “unhappiness” could become “un” and “happiness.” This approach allows models to understand new words by recognizing familiar components. The word “neural” might tokenize as “ne” and “ural,” enabling the model to handle variations and new combinations it has never seen before.

The advantage of subword tokenization becomes clear when dealing with related words that share common roots. Instead of treating each variation as a completely new token, the model can leverage shared subword patterns. This reduces vocabulary size while maintaining the ability to represent any text, making it the optimal choice for Large Language Models. Tools like the TikTokenizer demonstrate how original text gets broken down into these subword tokens, revealing the building blocks that LLMs use to understand and generate language.

1.3.1 Problems with Tokenization Methods

Word Based Tokenization Limitations

Figure 1.15 Limitations of word-based tokenization: related words like “learn,” “learn-ing,” “learned,” and “learnt” are treated as entirely separate tokens, and out-of-vocabulary words cannot be processed.

Word based tokenization treats each word as an independent unit, creating fundamental challenges for language models. The most significant issue is the failure to recognize relationships between related words. Words sharing common roots like “learn,” “learning,” “learned,” and “learnt” are treated as entirely separate tokens, forcing the model to learn each variation independently without understanding their connection.

The vocabulary explosion presents another critical problem. English alone requires over 200,000 word tokens, with filler words like “this,” “is,” and “a” consuming valuable vocabulary space despite contributing minimal semantic value. Most critically, the out of vocabulary problem renders models helpless when encountering unseen words. Simple spelling mistakes transform “running” into the unrecognizable “runing,” while new terms or proper nouns become impossible to process, leaving the model unable to make educated guesses about meaning

Character Based Tokenization Drawbacks

Figure 1.16 Character-based tokenization reduces the vocabulary to 256 ASCII characters but dramatically increases sequence length.

Character tokenization solves vocabulary size by using only 256 ASCII characters, but creates severe new problems.

Figure 1.17 Sequence length explosion with character tokenization: “Hello, world!” grows from a few word tokens to thirteen character tokens, and individual letters carry no semantic meaning.

Sequence length explodes dramatically: “Hello, world!” grows from two (or six) word tokens to thirteen character tokens. This expansion makes processing computationally expensive and quickly exhausts context windows in longer texts.

More fundamentally, character tokenization destroys semantic understanding. Individual letters carry no meaning, forcing models to reconstruct word boundaries and meanings from scratch. The model cannot recognize that “lowest” and “highest” share the meaningful suffix “est” indicating superlatives. When presented with “Hello,world!” as individual characters, the model sees meaningless symbols rather than a greeting, losing the essence of language structure entirely.

The Subword Tokenization Solution

Figure 1.18 Subword tokenization splits “modernization” into “modern” and “ization,” both reusable components found across many English words.

Subword tokenization provides the optimal balance, breaking words into meaningful components. “Modernization” becomes “modern” and “ization,” both reusable parts appearing across many words. This approach maintains reasonable vocabulary size while preserving meaning, handles new words through familiar components, and keeps token counts manageable. The model can now understand misspellings and new terms by recognizing known subword patterns.

Figure 1.19 The tokenization challenge: how should “learning” be split? Byte PairEncoding provides a systematic, data-driven answer.

The challenge remains: how should “learning” tokenize? As one token, as “learn” plus “ing,” or broken further? Byte Pair Encoding provides the systematic answer, using frequency analysis to determine optimal splits that balance vocabulary efficiency with semantic preservation.

1.4 Byte Pair Encoding

Byte Pair Encoding transforms the challenge of tokenization into a systematic process. Originally developed as a text compression algorithm in the 1990s, BPE now serves as the foundation for tokenization in models like GPT. The algorithm iteratively merges the most frequent character pairs, building a vocabulary from the bottom up.

For LLMs, BPE builds vocabularies systematically. Consider this corpus with word frequencies:

Figure 1.20 A small corpus with word frequencies used to illustrate the BPE algorithm: “old” appears 7 times, “older” 3 times, “finest” 9 times, and “lowest” 4 times.

Step 1: Add End-of-Word Markers

Figure 1.21 End-of-word markers (</w>) appended to each word to distinguish word boundaries and preserve morphological information.

The first step in BPE adds end-of-word markers (</w>) to distinguish word boundaries. Words transform as: old becomes old</w>, older becomes older</w>, finest becomes finest</w>, and lowest becomes lowest</w>. This boundary marker is crucial because the same character sequence carries different meanings based on position. The sequence “est” functions as a suffix in “lowest</w>” (indicating superlative) but as a prefix in “esteem” (with completely different meaning). Without these markers, the tokenizer cannot distinguish between identical character sequences that serve different linguistic roles, losing critical information about word structure and morphology.

Step 2: Split into Characters

Figure 1.22 Each word is decomposed into individual characters, providing atomic units for iterative merging.

After adding end-of-word markers, each word is decomposed into individual characters, treating each as a separate token. The word old</w> becomes the sequence [o, l, d, </w>], while older</w> splits into [o, l, d, e, r, </w>]. Similarly, finest</w> breaks down to [f, i, n, e, s, t, </w>] and lowest</w> to [l, o, w, e, s, t, </w>]. This character-level decomposition serves as the starting point for BPE, providing the atomic units from which larger, more meaningful tokens will be built through iterative merging based on frequency patterns in the data.

Step 3: Count Character Pairs & Merge

Figure 1.23 Counting adjacent character pairs across the corpus, weighted by word frequency. The pair “es” appears 13 times and is selected for the first merge.

The algorithm now counts all adjacent character pairs across the corpus, weighted by word frequency. The pair “es” appears 13 times (9 from “finest” plus 4 from “lowest”), as does “st” with the same distribution. The pairs “ol” and “ld” each appear 10 times (7 from “old” plus 3 from “older”), while “ne” and “in” from “finest” contribute 9 occurrences each.

With “es” as the most frequent pair, the algorithm performs its first merge, creating a new token “es” and updating the representations: finest</w> becomes [f, i, n, es, t, </w>] and lowest</w> becomes [l, o, w, es, t, </w>].

Figure 1.24 After the first merge (“es”), pairs are recounted and “est” emerges as the next most frequent pair, triggering a second merge.

After recounting pairs with this new token, “est” emerges as highly frequent, triggering a second merge. The token “est” replaces the “es” and “t” sequences, transforming finest</w> into [f, i, n, est, </w>] and lowest</w> into [l, o, w, est, </w>]. Through these iterative merges, BPE progressively builds larger, more meaningful tokens from the most frequent patterns in the data, creating an efficient vocabulary that captures common linguistic structures.

Step 4: Building the Complete Vocabulary

The merging process continues iteratively, identifying increasingly complex patterns. Common prefixes like “old” become single tokens when they appear frequently across multiple words. Suffixes with end markers like “est</w>” are preserved as units to maintain their grammatical function. Frequent character sequences like “low” merge into single tokens regardless of their position.

After multiple iterations, the final vocabulary becomes a hierarchical collection of tokens at different granularities. It contains individual characters [o, l, d, e, r, f, i, n, w, s, t] for handling rare sequences, common subwords [es, est, old, low, fin] that appear across multiple words, and complete frequent words [old</w>, finest</w>] that occur often enough to warrant their own tokens. This multi-level vocabulary enables efficient encoding of common patterns while maintaining the flexibility to tokenize any possible input.

Figure 1.25 GPT-2’s final vocabulary of 50,257 tokens, built from 50,000 BPE merges. Each token is mapped to a unique numerical ID used by the model internally.

GPT-2 performs 50,000 merges to build its vocabulary, creating a rich token set that balances compression with expressiveness. Each token in this vocabulary is assigned a unique token ID, a numerical identifier that the model uses internally. For example, in GPT-2’s vocabulary, common words like “Building” map to ID 25954, while special tokens like “</endoftext>” receive IDs like 50256, creating a complete dictionary of 50,257 token-ID pairs that serves as the bridge between text and numerical processing.

When the model encounters an unfamiliar word, it gracefully degrades to smaller subwords or individual characters, ensuring robust handling of misspellings, neologisms, or foreign terms. This fallback mechanism makes BPE remarkably resilient, capable of processing any text while maintaining efficiency for common patterns.

With our text now converted into meaningful tokens through BPE and mapped to numerical IDs, the next challenge is transforming these discrete symbols into continuous numerical representations that neural networks can process, leading us to the crucial concept of embeddings.

1.5 Word Embedding

After tokenization transforms text into discrete symbols and assigns them numerical IDs, we face a fundamental challenge: these IDs are merely labels that convey no semantic information. The token ID 25954 for “Building” tells the model nothing about buildings, construction, or architecture. To enable neural networks to process language meaningfully, we need to convert these discrete tokens into continuous numerical representations that capture semantic relationships. This is where word embeddings become essential.

The Limitations of Simple Encoding

Early approaches to numerical representation revealed critical limitations. One-hot encoding represents each token as a vector of zeros with a single one at the token’s position. For a vocabulary of 50,000 tokens, “cat” might be encoded as 50,000 zeros except for a single one at position 3. While this eliminates arbitrary ordering, it creates sparse, high-dimensional vectors where every word is equally distant from every other word. The vectors for “cat” and “dog” are as orthogonal as those for “cat” and “quantum”, providing no semantic signal. Similarly, bag-of-words models count word occurrences but lose all sequential information, treating “dog bites man” and “man bites dog” identically despite their opposite meanings.

Learning Meaning Through Context

Figure 1.26 high-dimensional embeddings transform semantic meaning into geometric coordinates. In this 768-dimensional space, linguistic relationships are defined by proximity, grouping concepts like animals and fruits into distinct neighborhoods.

The breakthrough came from the distributional hypothesis: words appearing in similar contexts tend to have similar meanings. If “coffee” frequently appears near “morning,” “cup,” and “brew,” while “tea” appears near similar words, a model can learn that coffee and tea are related concepts. Word2Vec revolutionized this approach by training neural networks to predict words from context (CBOW) or context from words (Skip-gram). Through millions of training examples, the network’s hidden layer learns to position similar words near each other in vector space. After training, “king” naturally clusters near “queen” and “prince,” while “banana” groups with “apple” and “fruit.” Most remarkably, these embeddings capture analogical relationships geometrically: the vector arithmetic “king - man + woman” yields a vector nearly identical to “queen,” demonstrating that the model has learned abstract concepts like gender and royalty as directions in space.

Embeddings in Large Language Models

Figure 1.27: Dense embedding vectors transform token IDs into high-dimensional representations where specific dimensions encode learned semantic features. This layer functions as a learned lookup table that maps tokens to unique vectors, enabling the model to capture nuanced semantic attributes like parts of speech or categorical relationships.

Modern LLMs transform token IDs into dense embedding vectors, typically ranging from 768 to 4096 dimensions, where each dimension encodes aspects of meaning learned during training. Unlike Word2Vec’s static embeddings where each word has one fixed representation, transformer models employ contextual embeddings that dynamically adjust based on surrounding tokens. The word “bank” receives different vector representations when appearing in “river bank” versus “investment bank,” enabling the model to disambiguate meaning through context. These embeddings are learned end-to-end during training, with the model discovering optimal representations that maximize its ability to predict the next token. The embedding layer becomes a learned lookup table that maps each of the 50,000+ token IDs to a unique vector in high-dimensional space, where semantic similarity translates to geometric proximity.

The power of LLM embeddings lies in their ability to encode multiple layers of linguistic information simultaneously. Each vector captures semantic meaning (cat near dog), syntactic roles (verbs clustering separately from nouns), conceptual relationships (similar terms grouping together), and even abstract patterns like sentiment or formality. Through billions of training examples, the model learns to position tokens in this space such that vector operations correspond to meaningful transformations. This geometric structure enables transformers to perform complex reasoning by manipulating these vectors through attention mechanisms and feed-forward networks, turning language understanding into mathematical computation.

However, embeddings alone cannot capture the sequential nature of language, where word order fundamentally changes meaning. This limitation leads us to positional embeddings, which encode each token’s location in the sequence, enabling transformers to understand that “dog bites man” differs crucially from “man bites dog.”

Positional Embedding

The Need for Positional Information

figure 1.28: Positional embeddings introduce sequential context into the Transformer architecture by adding unique position-aware vectors to token embeddings. This summation allows the model to distinguish identical tokens appearing at different sequence locations, enabling the architecture to capture syntactic and referential relationships despite its inherently parallel, set-based processing nature.

In natural language, word order fundamentally shapes meaning. Consider the sentences “The dog chased the cat” versus “The cat chased the dog.” While both sentences contain identical words, their meanings differ entirely based on word positioning. Traditional sequential models like RNNs inherently capture this ordering through their recurrent nature. However, the Transformer architecture processes all tokens simultaneously through self-attention, treating input as an unordered set. Without explicit positional information, a Transformer would produce identical representations for “dog” regardless of its position in the sentence, making it impossible to distinguish between different occurrences or understand sequential relationships.

This limitation becomes particularly problematic when dealing with pronouns and references. In “The dog chased the ball but it could not catch it,” the two instances of “it” refer to different entities based solely on their positions relative to other words. To address this fundamental limitation, Transformers incorporate positional embeddings that encode sequence order information directly into the model’s representations.

Integer Positional Encoding: The Simplest Approach

figure 1.29: The additive method of injecting positional data by combining token embeddings with integer-based position vectors, while noting the drawback that large integer values can interfere with and potentially confuse the model regarding the original word's semantic meaning.

The most straightforward solution involves assigning each position a unique integer value. In this scheme, if a token appears at position 300 in the sequence, we create a positional embedding vector where every dimension contains the value 300. This vector, matching the token embedding dimensions, gets added element-wise to the token embedding.

For a concrete example with an 8-dimensional embedding space, the token “dog” at position 300 would receive a positional embedding of [300, 300, 300, 300, 300, 300, 300, 300]. The final input representation becomes the sum of the token embedding and this positional embedding.

However, this approach suffers from a critical flaw: scale mismatch. Token embeddings typically contain small values clustered around zero, carefully learned to capture semantic nuances. Position values, especially for longer sequences, can grow arbitrarily large. When position 500 adds [500, 500, ...] to delicate token embeddings with values like [0.23, -0.15, 0.08, ...], the positional signal completely overwhelms the semantic information. The model loses the ability to distinguish between different words, focusing instead on their positions.

Binary Positional Encoding: Constraining the Range

figure 1.30: This technique represents positions using binary bit strings to keep values between 0 and 1, but it creates sudden jumps in the embedding space that complicate the training process for the model.

To address the magnitude problem inherent in integer encoding, binary positional encoding represents positions using their binary representation, naturally constraining all values between 0 and 1. This approach transforms each position number into its binary form and uses those bits directly as the positional embedding vector.

Figure 1.31: This demonstrates how binary representations change across consecutive integer positions. It highlights the rapid bit flipping in the least significant bit position which creates frequency oscillations that make optimization more difficult for the model.

Consider the visualization showing positions 64 through 75 with their 8-bit binary representations. Position 64, which equals 01000000 in binary, becomes the embedding vector [0, 1, 0, 0, 0, 0, 0, 0]. Here, each bit position corresponds to a dimension in the embedding space, with i=8 representing the most significant bit (MSB) and i=1 representing the least significant bit (LSB).

Looking at the pattern across consecutive positions reveals a fascinating structure. Position 64 starts with [0, 1, 0, 0, 0, 0, 0, 0]. Position 65 becomes [0, 1, 0, 0, 0, 0, 0, 1], position 66 transforms to [0, 1, 0, 0, 0, 0, 1, 0], and position 67 yields [0, 1, 0, 0, 0, 0, 1, 1]. The rightmost bit (i=1) flips with every single position increment, creating a rapid alternation between 0 and 1.

The second bit from the right (i=2) follows a different rhythm, maintaining its value for two positions before flipping. It stays 0 for positions 64-65, switches to 1 for positions 66-67, returns to 0 for positions 68-69, and so forth. The third bit (i=3) changes every four positions, remaining stable from 64-67, then flipping for 68-71.

This creates a hierarchical encoding scheme where each bit position operates at a different frequency. The LSB oscillates most rapidly, capturing fine-grained positional differences between adjacent tokens. Moving leftward through the bits, oscillation frequencies decrease exponentially. The fourth bit changes every 8 positions, the fifth every 16 positions, the sixth every 32 positions, and the seventh every 64 positions. The MSB (i=8) remains constant for 128 consecutive positions before flipping.

In the visualization, this pattern becomes immediately apparent. The rightmost column shows constant flickering between (0) and (1) for every position. The i=2 column displays pairs of same cells. The i=3 column shows groups of four, and this doubling pattern continues across all dimensions. The leftmost column (i=8) remains uniformly across the entire visible range, as positions 64-75 all share the same MSB value of 0.

This encoding elegantly solves the scale problem that plagued integer encoding. Instead of values potentially reaching into the thousands, every dimension now contains either 0 or 1. When added to token embeddings clustered around zero, these binary values preserve the semantic information while injecting positional signals at a comparable scale.

The hierarchical structure provides the model with positional information at multiple granularities simultaneously. Lower-indexed dimensions encode local sequential relationships, helping the model understand which tokens appear near each other. Higher-indexed dimensions capture global positional context, indicating whether tokens appear in the first half versus second half of the sequence, or in early versus late quarters.

However, binary encoding introduces a critical limitation: discontinuity. The hard transitions between 0 and 1 create step functions rather than smooth gradients. When the model needs to learn relationships between positions 67 ([0, 1, 0, 0, 0, 0, 1, 1]) and 68 ([0, 1, 0, 0, 0, 1, 0, 0]), multiple dimensions flip simultaneously. These abrupt changes complicate gradient-based optimization, as the loss landscape contains sharp edges and discontinuous regions.

During backpropagation, these discrete jumps prevent smooth gradient flow. Small parameter updates cannot gradually transition the model’s understanding between binary states. The optimizer must navigate around these discontinuities, potentially getting stuck in suboptimal configurations or requiring careful learning rate scheduling to handle the non-smooth optimization landscape.

Despite these challenges, binary encoding demonstrates the key insight that positional information can be encoded through patterns of oscillation at different frequencies. This conceptual breakthrough, showing that different dimensions can operate at different temporal scales, directly inspired the development of sinusoidal positional encoding, which maintains these beneficial oscillatory patterns while ensuring continuous, differentiable representations throughout the embedding space.

Sinusoidal Positional Encoding: Continuous Representations

Figure 1.32: Sinusoidal PE applies trigonometric sine and cosine functions to generate continuous and bounded positional vectors, which enables the model to learn sequential relationships while avoiding the optimization challenges and discontinuities inherent in integer and binary positional encodings.

The breakthrough in positional encoding came with the sinusoidal approach, introduced in the seminal “Attention Is All You Need” paper. This method preserves the oscillatory patterns discovered in binary encoding while ensuring smooth, continuous values bounded between -1 and 1, eliminating the discontinuity problems that hindered optimization.

The Mathematical Foundation

The sinusoidal formulation employs alternating sine and cosine functions across dimensions:

Where pos denotes the token’s position in the sequence, i represents the dimension index, and d_model indicates the total embedding dimensionality. The constant 10000 serves as the base for creating a geometric progression of wavelengths across different dimensions.

Frequency Spectrum Analysis

Taking GPT-2’s architecture as an example, with d_model = 768 and maximum context length = 1024, we can observe how different dimensions encode positional information at varying frequencies. For any given position, we compute 768 values using the alternating sine-cosine formulas.

At the lowest dimension (i=0), the formula simplifies to sin(pos/1) = sin(pos), creating rapid oscillations. The adjacent dimension uses cos(pos/1) = cos(pos). As the dimension index increases, the denominator 10000^(2i/768) grows exponentially, progressively slowing the oscillation frequency.

Sinusoidal Patterns Across Different Dimensions

Figure 1.33: The visualization reveals how positional encodings behave across four different dimension indices:

At i=1, both sine and cosine components oscillate extremely rapidly, appearing as dense vertical lines that alternate between approximately -1 and 1. This high-frequency pattern changes with nearly every position, capturing fine-grained local relationships between adjacent tokens.

At i=50, the oscillation frequency decreases noticeably. The sine and cosine waves create regular patterns with periods spanning roughly 20-30 positions. These medium-frequency components encode relationships at the phrase or sentence level.

At i=150, the waves become smooth and gradual, with clear sinusoidal curves visible. The sine (green) and cosine (blue) components maintain their 90-degree phase offset, completing only 2-3 full cycles across the entire 1024-position range. These dimensions capture broader structural information about whether tokens appear in early, middle, or late portions of the sequence.

At i=250, the oscillation becomes extremely slow, with the functions barely completing a single cycle across the full context. The cosine component remains nearly constant around 1, while the sine component stays close to 0, providing stable anchoring for global position context.

Sinusoidal encoding creates a hierarchical representation where each position receives a unique 768-dimensional fingerprint. Lower dimensions oscillate rapidly between positions, capturing local token relationships and word order, while higher dimensions change gradually, encoding broader context like paragraph boundaries and document structure. This combination of multiple sine-cosine pairs at different frequencies generates a unique signature for every position. Unlike binary encoding’s abrupt 0-to-1 transitions, sinusoidal encoding provides smooth, continuous functions that enable stable gradient flow during backpropagation, dramatically improving training efficiency. The bounded range between -1 and 1 keeps positional signals at a scale comparable to token embeddings, preventing positional information from overwhelming semantic content while allowing the optimizer to make incremental refinements.

The sinusoidal approach offers significant practical advantages: it requires no learned parameters, reducing model complexity and training overhead, and its mathematical formulation naturally extends to arbitrary sequence lengths, potentially enabling generalization beyond training context sizes. In practice, positional encodings are precomputed for the maximum sequence length and stored as a lookup table. During processing, these encodings are retrieved and added element-wise to token embeddings, preserving semantic information while injecting positional signals. This simple yet elegant solution simultaneously addresses multiple challenges: maintaining bounded values, ensuring smooth optimization, providing unique position identification, and encoding multiscale temporal information. These properties have established sinusoidal positional encoding as a cornerstone of the Transformer architecture, inspiring numerous variations while remaining widely used in its original form across modern language models.

1.6 Transformer Block

Figure 1.34 Inside the transformer block: multi-head attention, feed-forward network, layer normalization, and dropout layers work in sequence to process token representations

Having converted our raw text into meaningful numerical representations through tokenization and embeddings, we now enter the heart of the language model: the Transformer Block. This is where the real magic happens. The block contains several components working in sequence, including layer normalization, dropout layers, and feed forward networks. However, before we dive into these supporting elements, we need to understand the star of the show: the attention mechanism. The multi head attention layer is what gives transformers their remarkable ability to understand context and relationships between words, no matter how far apart they appear in a sentence. Once we grasp how attention works and explore the feed forward network that follows, we can then circle back to understand how the other components like dropout and layer normalization help stabilize and improve the overall system. For now, let’s focus on what makes transformers truly powerful: their attention mechanism.

1.7 The Need for Attention Mechanism

Figure 1.35 Timeline of sequence modeling methods from early RNNs to LSTMs, attention with RNNs, transformers, and GPT models.

Feedforward neural networks see every input as independent. For a sentence such as “The cat sat on the mat” the model processes each word separately and has no built in notion of order or context. This is not enough for language, where meaning depends on how words are arranged.

Recurrent neural networks introduce a hidden state that is passed along the sequence. The encoder reads tokens one by one, updates its hidden state at each step, and hands the final state to a decoder. The decoder must use this single vector as a summary of the entire input sentence. As sequences get longer, early information is squeezed into this fixed size state and gradually fades. This is the context bottleneck.

LSTMs improve the situation with a cell state and gates that control what to store and what to forget. They maintain information over longer spans than basic RNNs, but they still process tokens step by step and still rely on compressed hidden states. Long sentences can still overwhelm this bottleneck.

Figure 1.36: Encoder-decoder RNN for the sentence “I will eat” showing encoder hidden states h1, h2, h3 and decoder states that must rely on a single summary vector.

To see the bottleneck more concretely, consider an encoder decoder model that translates the English sentence “I will eat into French”. The encoder produces hidden states h1, h2, h3 for the three input tokens and a final state that is passed to the decoder. Without attention, the decoder can only use this final state when generating the first French word. It has no direct way to reach back to h1 or h2.

Attention

Attention removes the hard bottleneck by giving the decoder direct access to all encoder states. At each decoding step the model compares the current decoder state with every encoder state and produces attention scores. After a softmax these scores become attention weights that sum to one.

The decoder then forms a context vector as a weighted sum of the encoder states. If the first input word is most relevant for the current output, its weight may be close to one while the others are close to zero. At the next step the weights are recomputed and the model can shift its focus to a different part of the sentence.

In the translation example, when the decoder produces the first French word it might focus almost entirely on h1. When it moves on to the second French word it can focus more on h2, and so on. Instead of depending on a single final state, the decoder now has a flexible view over the entire input sequence at every step.

Bahdanau attention

Figure 1.37: Bahdanau attention architecture: the encoder produces a sequence of states and the decoder combines its own state with a context vector formed as a weighted sum of all encoder states

Bahdanau attention was the first widely adopted implementation of this idea. The encoder is still a recurrent network that produces a sequence of hidden states. The decoder is also recurrent, but before predicting each target token it computes alignment scores between its current state and every encoder state. These scores become attention weights, and their weighted sum is the context vector used for prediction.

The attention weights can be visualized as a matrix whose rows correspond to target words and columns correspond to source words. Each cell shows how strongly the model attends to a particular source word when generating a particular target word. This view reveals attention as a soft alignment between the two sentences.

Figure 1.38: Attention heatmaps for a French-to-English sentence pair. The left grid shows overall alignments; the right grid highlights how the model focuses on “European Economic Area” when generating “zone economique europeenne,” capturing word reordering.

These heatmaps show that many words align along a near diagonal, indicating similar order in both languages. Off diagonal patterns reveal reordered phrases. For example, the French adjective corresponding to European appears last in the phrase, but its attention weights point back to the first English word. This ability to align by meaning rather than position is what allows attention based models to handle flexible word order and long range dependencies.

Finally, it is helpful to remember where this attention block lives inside the full transformer model from earlier chapters. The transformer encoder and decoder both contain stacked attention and feedforward sublayers that operate on token and positional embeddings.

We are now ready to see why attention became the central idea in modern language models. Starting from simple recurrent networks and LSTMs, we saw how the context bottleneck makes it hard to remember all the details of a long sentence. Bahdanau attention solved this by letting the decoder look back at every encoder state and learn soft alignments between source and target words, which we visualized through attention weights and heatmaps. So far, attention has connected two different sequences, such as English and French sentences. In the next section we will study self attention in detail and see how letting every token attend to every other token becomes the core operation of the transformer.

1.8 Self Attention Mechanism

What does Self Attention actually means ?

Now that we understand the mechanics of attention, let’s clarify what makes self-attention special, the key concept behind modern language models like Transformers.

Two Types of Attention

Figure 1.39: Two types of attention: cross-attention connects words across different sequences (e.g., translation), while self-attention connects words within the same sequence.

To understand self-attention, we first need to see where attention was used before. There are two fundamental ways attention can work:

Between Sequences: Attention connects words across different sequences, think of translating from one language to another.

Within a Sequence: Attention connects words within the same sequence to capture relationships and context.

Attention in Translation

Figure 1.40: Cross-attention in translation: the English phrase “The next day is bright” is aligned to its French counterpart, with attention determining which source words correspond to which target words.

In traditional translation tasks, attention operates between two sequences. Imagine translating the English phrase “The next day is bright” into French. The word order might change. “Day” might align with “jour,” but its position in the French sentence could be different. Attention helps the model figure out these cross-language alignments, which English word corresponds to which French word. This works beautifully for translation. But what happens when we’re not translating at all?

Enter Self-Attention

Figure 1.41: Self-attention: every word in a single sentence attends to all other words in that same sentence to build contextual understanding.

Consider a different task: predicting the next word in a sentence. Or understanding what a pronoun refers to. Or simply trying to grasp the meaning of a sentence. Here, we don’t have two separate sequences. We have just one, the sentence itself. This is where self-attention comes in.

Self-attention means that every word in a sentence attends to all other words in that same sentence. Instead of looking across two different sequences (like English and French), the model examines how words relate to each other within a single sequence. The word “day” attends to “next,” to “bright,” to “the”, to everything in its own sentence. It’s attention turned inward. The sequence attending to itself. That’s why we call it self-attention. We cannot encode these complex relationships directly in the attention mechanism using just the raw input embeddings. The connections between words depend on context, grammar, meaning, and a dozen other subtle factors that shift from sentence to sentence.

So what do we do when faced with complexity we can’t hard-code? We let the model learn it. We leave it to weight matrices that can be trained. Before we dive into the mechanics, let’s be clear about our goal. We start with input embeddings, numerical representations of words. But here’s what we want to end up with: context vectors.

What’s the difference?

An input embedding represents a word in isolation. The embedding for “bank” is always the same, whether you’re talking about a financial institution or the side of a river. But a context vector represents a word as it appears in a specific sentence, infused with information from the words around it.

Think about

The input embedding for the second “it” doesn’t know what “it” refers to, it’s just a generic representation. But the context vector we’re building will carry information from “ball,” from “catch,” from the entire sentence.

Figure 1.42: From input embedding to context vector: the static representation of a word is enriched with information from all surrounding words through self-attention.

It will understand that this particular “it” refers to the ball. So our entire journey with self-attention, the queries, the keys, the attention scores we’re about to explore, all of it serves one purpose: transforming static input embeddings into dynamic context vectors that understand meaning in context.

1.9 Understanding the Input Embedding Matrix

Figure 1.43: For each word, the word embedding and positional embedding are summed to produce the input embedding vector.

As we’ve already seen, for each word in our sentence, we have an embedding vector combined with positional information, that is, the word embedding plus the positional embedding that tells us where the word sits in the sequence. The sum of these two gives us our input embedding vector for each word.

When we stack all these input embedding vectors together for an entire sentence, we get what’s called the input embedding matrix.

Let’s say we’re working with the sentence

that’s five words. Our input embedding matrix would have dimensions (5, 8).

Figure 1.44: The input embedding matrix for the sentence “The next day is bright” has shape (5, 8): five rows (one per word) and eight columns (the embedding dimension).

What do these numbers mean?

The 5 rows come from having 5 words. Simple enough. Each word gets its own row in the matrix. If we had ten words, we’d have ten rows. The number of rows always matches the number of words in our sequence.

The 8 columns represent the dimensionality we’ve chosen for our embeddings. Each word is represented as an 8-dimensional vector, eight numbers that capture its meaning. This dimension is something we decide when building our model. It’s a design choice.

In GPT-2, for instance, the embedding dimension varies: 768 for GPT-2 Small, all the way up to 1,600 for GPT-2 XL. Larger dimensions can capture more nuanced information, but they also require more computation.

The Problem We’re Solving

So here we are with our input embedding matrix. Each word has its 8-dimensional vector. But here’s what’s missing: these vectors exist in isolation. They don’t know about each other.

Look at the word “day” in our sentence “The next day is bright.” Its input embedding vector is just a generic representation of the word “day.” It doesn’t know it should pay attention to “bright.” It doesn’t know that “next” right before it gives it temporal context. It has no idea how much importance it should give to “the” or “is” or any other word in the sentence.

This is exactly why we need to transform input embeddings into context vectors. We need to integrate information from all the other words. We need each word’s representation to reflect not just what it is, but what it means in this particular sentence, surrounded by these particular neighbors. That’s the journey we’re about to take.

Before we can perform any attention calculations, we must first define our input sequence and its corresponding embedding matrix. We will use the PyTorch library to create a tensor that holds this information for our example sentence: “The next day is bright”. Each word is represented by an 8-dimensional vector

Listing 1.1 Defining the input embedding matrix

import torch

words = [’The’, ‘next’, ‘day’, ‘is’, ‘bright’]

inputs = torch.tensor([
    [0.32, 0.21, 0.43, 0.21, 0.86, 0.67, 0.98, 0.23], # The
    [0.43, 0.56, 0.43, 0.56, 0.69, 0.21, 0.56, 0.21], # next
    [0.56, 0.21, 0.43, 0.21, 0.54, 0.12, 0.89, 0.98], # day
    [0.87, 0.34, 0.18, 0.32, 0.75, 0.12, 0.54, 0.92], # is
    [0.76, 0.21, 0.85, 0.34, 0.98, 0.23, 0.68, 0.34]  # bright
], dtype=torch.float32)



print(”Input Embedding Matrix:”)
print(inputs)
print(”\nMatrix Shape:”)
print(inputs.shape)

Running the previous code prints the following output

Input Embedding Matrix:
tensor([
[0.3200, 0.2100, 0.4300, 0.2100, 0.8600, 0.6700, 0.9800, 0.2300],  [0.4300, 0.5600, 0.4300, 0.5600, 0.6900, 0.2100, 0.5600, 0.2100],
[0.5600, 0.2100, 0.4300, 0.2100, 0.5400, 0.1200, 0.8900, 0.9800],
[0.8700, 0.3400, 0.1800, 0.3200, 0.7500, 0.1200, 0.5400, 0.9200],
[0.7600, 0.2100, 0.8500, 0.3400, 0.9800, 0.2300, 0.6800, 0.3400]
])

Matrix Shape:
torch.Size([5, 8])

The output shows our input object is a tensor with a shape of torch.Size([5,8]). This confirms we have a matrix with 5 rows, one for each of our tokens, and 8 columns, representing the 8-dimensional embedding vector for each token. This matrix is the starting point for the self-attention mechanism, but as noted, these vectors exist in isolation and lack any contextual information from their neighbors.

1.10 From Embeddings to Queries, Keys & Values

Here’s where we meet the heart of self-attention: three trainable weight matrices called Queries, Keys, and Values.

Figure 1.45: The input embedding matrix is multiplied by three separate weight matrices Wq, Wk, and Wv to produce Query, Key, and Value matrices

You might wonder, why three? Why not just use the input embeddings directly?

The answer lies in a fundamental principle of neural networks: they’re universal function approximators. They can learn complex patterns if we give them the right structure. So instead of trying to hand-code how words should relate to each other, we do something smarter.

We initialize three weight matrices with random values at the start. Then we let the training process figure it out. During training, these matrices learn how to transform embeddings in ways that capture meaningful relationships. The Query matrix learns to create vectors that “ask questions.” The Key matrix learns to create vectors that “answer” whether they’re relevant. And the Value matrix? It learns what information should actually be passed along once we know which words matter. We’re not telling the model how attention should work, we’re giving it the tools to learn it on its own.

Let’s understand this by taking one example

Figure 1.46: The sentence “The next day is bright” with the word “next” highlighted as the current focus of the attention mechanism.

When we focus on a specific word say, “next”.

Figure 147: The word “next” acts as the Query, asking how much attention it should pay to each other word in the sentence.

we need to decide how much attention it should pay to all the other words in the sentence. This is where our terminology becomes important. The word we are focusing on (“next”) is called the Query (Q).

Figure 148: The other words, “the,” “day,” “is,” “bright”, serve as Keys that the query evaluates for relevance.

The other words in the sentence, “the,” “day,” “is,” “bright,” are called Keys (K). These are the words that the query will evaluate. They’re potential sources of information.

Figure 149: Attention scores α between the query word “next” and all keys. Each score quantifies how strongly “next” should attend to each other word.

Now comes the crucial part: the attention score (α). This score determines how much importance “next” should give to each of these other words. Should “next” pay more attention to “day” (the word right after it) or to “bright” (further away)? The attention scores tell us exactly this.

So “next” uses these attention scores to focus on other words in the sentence, weighing some as more important, others as less so. This is how a word builds its understanding of context.

For example, the attention score α₂₁ means:

• “Next” (X₂) is attending to “The” (X₁).

• The first 2 represents “next” (position 2 in the sentence).

• The second 1 represents “the” (position 1 in the sentence).

The goal of self-attention is to take these attention scores (α values) and use them to modify the original input embeddings, creating context vectors that contain richer information.

• Input Embedding (X₂ - “next”): Just represents the word itself.

• Context Vector (C₂ - “next”): Now contains information from all relevant words around it, based on attention scores.

Instead of just knowing “next” as an isolated word, the context vector of “next” now understands:

• How much “next” relates to “day” (α₂₃)

• How much “next” relates to “the” (α₂₁)

• How much “next” relates to “is” (α₂₄)

This transformation from input embeddings to context vectors is what makes self-attention so powerful, it helps the model understand relationships between words, not just individual tokens.

Context Vector is an enriched embedding vector. It combines information from all other input elements

The Dimensions of Query, Key, and Value Matrices

Now let’s talk about the actual shape and size of these weight matrices. Understanding their dimensions is crucial to grasping how self-attention works mathematically.

Figure 1.50: Dimensions of the weight matrices: Wq, Wk, and Wv each have shape (d_in, d_out), where din matches the embedding dimension and d_out is a design choice.

If we look at the dimensions of the Query, Key, and Value matrices (Wq, Wk, and Wv), we’ll notice something interesting.

The number of rows in each of these matrices equals the number of columns in our input embedding matrix. Remember, our input embedding matrix has dimensions (5, 8), where 8 is our embedding dimension. So our weight matrices will have 8 rows.

The number of columns in these weight matrices, however, can be anything we choose. This is a design decision.

Figure 1.51: The terminology d_in and d_out: din = 8 is the input embedding dimension, and d_out is the chosen output dimension for queries, keys, and values.

When coding language models like GPT-2 or GPT-3, we use specific terminology for these dimensions:

d_in (Input Dimension): The dimension of our input embeddings. In our example, this is 8.

d_out (Output Dimension): The dimension we want for our query, key, and value vectors. This is the number of columns in our weight matrices.

Here’s an important point: you can choose any value for d_out. In practice, it’s often set equal to d_in for simplicity. So if our input dimension is 8, we might set the output dimension to 8 as well. But we don’t have to. In our example, we’re using d_out = 4. Why? To demonstrate that the output dimension is flexible. You have the freedom to choose what works best for your model.

Listing 1.2: Extracting a Token Embedding and Setting Dimensions

x_2 = inputs[1]          # embedding for “next”
d_in = inputs.shape[1]   # input dimension
d_out = 4                # dimension for Q, K, V in this toy example

print(x_2)
print(d_in)
print(d_out)

Here you select the second row of the input matrix, which is the 8 dimensional embedding for the word “next”. The variable d_in confirms that the embedding dimension is 8, matching the theory. The variable d_out is set to 2, which means each query key and value vector will live in a 2 dimensional space in the following examples. In real models d_out is much larger, but using 2 keeps the printed tensors readable.

Output

tensor([0.4300, 0.5600, 0.4300, 0.5600, 0.6900, 0.2100, 0.5600, 0.2100])
8
4

How These Matrices Learn

Figure 1.52: Weight matrices are initialized with random values and updated during training through backpropagation, learning to produce meaningful query, key, and value representations.

At the beginning, all the values in these weight matrices are initialized randomly. They start with no knowledge of language or attention patterns. But here’s where the magic of training comes in.

As we train the model using backpropagation, these random values gradually update themselves. The matrices learn which transformations help the model understand language better.

1.11 A Quick Note on Matrix Multiplication

Before we dive into multiplying our embedding matrices, let’s make sure we’re all on the same page about how matrix multiplication actually works. If you already know this, feel free to skip ahead. But if matrices feel a bit fuzzy, stick with me for a moment.

Figure 1.53: Matrix multiplication: Matrix A of shape (3, 2) multiplied by Matrix B of shape (2, 3) produces a result of shape (3, 3). The inner dimensions must match.

We have Matrix A with dimensions (3, 2) and Matrix B with dimensions (2, 3). Notice something important: the number of columns in Matrix A (which is 2) matches the number of rows in Matrix B (also 2). This isn’t a coincidence. For matrix multiplication to work, these inner dimensions must match.

When we multiply them, we get a result with dimensions (3, 3). The outer dimensions survive: 3 rows from Matrix A and 3 columns from Matrix B.

Figure 1.54: Element-wise computation in matrix multiplication: each output entry is the dot product of a row from the first matrix and a column from the second matrix.

To compute each element in the result, you take a row from the first matrix and pair it with a column from the second matrix. Multiply corresponding elements together, then sum them up.

Example: To find the element at position (1, 1):

• Take row 1 from Matrix A: [1, 2]

• Take column 1 from Matrix B: [7, 10]

• Calculate: (1 × 7) + (2 × 10) = 7 + 20 = 27

Another example: For position (2, 1):

• Take row 2 from Matrix A: [3, 4]

• Take column 1 from Matrix B: [7, 10]

• Calculate: (3 × 7) + (4 × 10) = 21 + 40 = 61

You repeat this pattern for every position. Row meets column, multiply and sum. That’s the entire process.

[Step 1] Creating Query, Key, and Value Vectors

The first step in converting our input embedding matrix into context embeddings is straightforward: matrix multiplication. Let’s walk through this process carefully, starting with how we create query vectors.

Figure 1.55 The input embedding matrix (5,8) is multiplied by the query weight matrix W_q (8,4) to produce the query matrix (5,4).

We take our input embedding matrix and multiply it by the query weight matrix (W_q). This transformation gives us our query vectors.

Figure 1.56: Detailed view of the matrix multiplication: each 8-dimensional word embedding is projected through the weight matrix to produce a 4-dimensional query vector

Each row of the input matrix represents one word with its 8-dimensional embedding.

Figure 1.57: Step-by-step computation showing how one row of the input matrix multiplied by W_q produces one row of the query matrix.

When we multiply this row by the weight matrix, we get a new row in the output, a 4-dimensional query vector for that word. This happens for all five words simultaneously.

Figure 1.58: The resulting query matrix: five words now have 4-dimensional query vectors, transformed from the original 8-dimensional embeddings.

The result? A query matrix where each of our five words now has its own query vector, transformed from 8 dimensions down to 4.

The Complete Picture

Figure 1.59: All three projections in parallel: the input embedding matrix is multiplied by W_q, W_k, and W_v simultaneously to produce Query, Key, and Value matrices, each of shape (5, 4)

Creating the query vectors is just the beginning. The same transformation process happens two more times, each with its own weight matrix, all operating in parallel.

For the key vectors, we multiply our input embedding matrix by the key weight matrix (W_K). Same dimensions, same process. Each word gets its own key vector.

For the value vectors, we multiply the input embedding matrix by the value weight matrix (W_V). Each word now has a value vector too.

Here’s something important to recognize: we’ve moved from an 8-dimensional space to a 4-dimensional space. More significantly, we’ve moved into a different kind of space altogether. We’re no longer dealing with input embeddings, those static representations of words. We’re now working with query, key, and value vectors. Each lives in its own transformed space, optimized for a specific purpose in the attention mechanism.

This might seem like an odd detour. Why transform our embeddings at all? Why not just work with them directly?

This trick, transforming data into different spaces is fundamental to deep learning, and it’s powerful for a simple reason: sometimes the patterns we need aren’t visible in the original data. Think about it this way. In computer vision, early systems used hand-crafted features like edges and corners. Then convolutional neural networks came along and learned to discover their own features automatically, finding patterns humans never thought to look for. That’s what’s happening here. We’re not stuck with the fixed relationships in our input embeddings. Instead, we let the model learn, through training, what transformations actually help it understand language.

Think of it like passing our input through three different lenses simultaneously. Each lens, each weight matrix transforms the same input embeddings in a different way, extracting different aspects of meaning. When all three transformations are complete, we have three new matrices sitting side by side, all sharing the same dimensions of (5, 4). These three matrices are now ready for the next step in the attention mechanism. The queries and keys will interact to figure out who should pay attention to whom. But that’s a story for the next section.

Listing 1.3 Initializing Query, Key, and Value Weight Matrices

torch.manual_seed(123)

W_query = torch.nn.Parameter(torch.randn(d_in, d_out), requires_grad=False)
W_key   = torch.nn.Parameter(torch.randn(d_in, d_out), requires_grad=False)
W_value = torch.nn.Parameter(torch.randn(d_in, d_out), requires_grad=False)

print(”W_query:”)
print(W_query)
print(”\nW_key:”)
print(W_key)
print(”\nW_value:”)
print(W_value)

Output

W_query:
Parameter containing:
tensor([[ 0.2961,  0.5166, -0.0973,  0.2340],
        [ 0.2517,  0.6886,  0.0451, -0.4128],
        [ 0.0740,  0.8665,  0.3210,  0.0185],
        [ 0.1366,  0.1025, -0.2314,  0.5642],
        [ 0.1841,  0.7264, -0.1035,  0.3399],
        [ 0.3153,  0.6871,  0.2478, -0.1520],
        [ 0.0756,  0.1966,  0.5142,  0.0813],
        [ 0.3164,  0.4017, -0.0879,  0.2904]])

W_key:
Parameter containing:
tensor([[ 0.1186,  0.8274,  0.1040, -0.3055],
        [ 0.3821,  0.6605, -0.2103,  0.1428],
        [ 0.8536,  0.5932, -0.1449,  0.3170],
        [ 0.6367,  0.9826,  0.2553, -0.0872],
        [ 0.2745,  0.6584,  0.0342,  0.5051],
        [ 0.2775,  0.8573, -0.2984,  0.1907],
        [ 0.8993,  0.0390,  0.1206,  0.2843],
        [ 0.9268,  0.7388, -0.0721,  0.3419]])

W_value:
Parameter containing:
tensor([[ 0.7179,  0.7058, -0.1630,  0.3310],
        [ 0.9156,  0.4340,  0.0982, -0.2753],
        [ 0.0772,  0.3565,  0.2056,  0.1468],
        [ 0.1479,  0.5331, -0.0925,  0.2391],
        [ 0.4066,  0.2318,  0.0194,  0.1844],
        [ 0.4545,  0.9737, -0.3086, -0.0417],
        [ 0.4606,  0.5159,  0.1274,  0.0219],
        [ 0.4220,  0.5786, -0.0853,  0.3640]])

This code creates the three trainable weight matrices that turn input embeddings into query, key and value vectors.
The variable d_in is the size of each input embedding, here 8. The variable d_out is the size we want for the query, key and value vectors, here 4.

The tensors W_query, W_key and W_value are wrapped in torch.nn.Parameter, which tells PyTorch that these tensors are learnable weights. During training, gradient descent will update these matrices so that they learn useful transformations.
The shapes printed at the end confirm that each weight matrix has shape 8, 4, matching the description in the theory where the number of rows equals d_in and the number of columns equals d_out.

Listing 1.4: Computing Query, Key, and Value Vectors

queries = inputs @ W_query   # shape: (5, 4)
keys    = inputs @ W_key     # shape: (5, 4)
values  = inputs @ W_value   # shape: (5, 4)

print(”queries.shape:”, queries.shape)
print(”keys.shape   :”, keys.shape)
print(”values.shape :”, values.shape)

print(”\nqueries:”)
print(queries)
print(”\nkeys:”)
print(keys)
print(”\nvalues:”)
print(values)

Output

queries.shape: torch.Size([5, 4])
keys.shape   : torch.Size([5, 4])
values.shape : torch.Size([5, 4])

queries:
tensor([[ 0.8840,  2.0469,  0.3419,  0.5755],
        [ 0.8723,  2.0443,  0.3241,  0.3917],
        [ 0.9738,  1.9925,  0.3154,  0.7673],
        [ 1.1051,  2.1175,  0.2772,  0.9102],
        [ 0.9692,  2.5180,  0.2741,  0.8069]])

keys:
tensor([[ 2.2381,  2.7132,  0.0507,  0.8830],
        [ 2.1564,  2.7927, -0.0064,  0.8684],
        [ 2.6705,  2.7739,  0.0891,  1.0862],
        [ 2.4087,  3.1074,  0.0683,  1.0651],
        [ 2.6735,  3.1745,  0.0538,  1.2434]])

values:
tensor([[ 2.1213,  2.5079, -0.2830,  0.7693],
        [ 2.0476,  2.1981, -0.2301,  0.4927],
        [ 2.1971,  2.4733, -0.2806,  0.9121],
        [ 2.4207,  2.5415, -0.3020,  0.9874],
        [ 2.2625,  2.5902, -0.2714,  0.9625]])

Here we apply the three weight matrices to the full input matrix. Each row of inputs is an embedding for one word.
The matrix multiplication inputs @ W_query takes every word embedding and projects it into query space. The result is a query matrix with shape 5, 4, one query vector of length four for each of the five words. The same happens for keys and values.

This mirrors the explanation in the text that we now have three new matrices, each of size number of tokens, d_out. We are no longer working with raw input embeddings but with transformed representations tailored for searching, being searched and being blended.

[Step 2] Computing Attention Scores

Now that we have our query, key, and value vectors, we’re ready for the heart of the attention mechanism: figuring out which words should pay attention to which other words. Remember, each word has a query vector and a key vector . When we compute the dot product between a query and a key, we get a number that represents how well they align. A high dot product means strong alignment, which translates to high attention. A low dot product means weak alignment, which means less attention.

Figure 1.60: The dot product between a query vector and a key vector produces a scalar attention score indicating how well they align.

Here’s where we hit a small technical hurdle. We want to compute all these dot products at once using matrix multiplication. Our Query matrix has dimensions (5, 4) and our Keys matrix also has dimensions (5, 4). If we try to multiply them directly, Query × Keys, we run into a problem. For matrix multiplication to work, the number of columns in the first matrix must equal the number of rows in the second matrix. But Query has 4 columns and Keys has 5 rows. They don’t match. The multiplication simply won’t work. The solution of this problem is to transpose the Keys matrix.

A Quick Note on Matrix Transpose

Figure 1.61: Matrix transpose: rows become columns and columns become rows, converting a (3, 2) matrix into a (2, 3) matrix.

If you’re already comfortable with matrix transpose, feel free to skip ahead to the next section. But if transpose feels unfamiliar or you want a quick refresher, stay with me for just a moment.

Figure 1.62: Transposing the Keys matrix from (5, 4) to (4, 5) so that it can be multiplied with the (5, 4) Query matrix.

When we transpose a matrix, we flip it along its diagonal. Rows become columns, and columns become rows. If you have a matrix with dimensions (3, 2), its transpose will have dimensions (2, 3). The first row of the original matrix becomes the first column of the transposed matrix. The second row becomes the second column, and so on. It’s like rotating the entire matrix 90 degrees and reflecting it. This simple operation is incredibly useful because it lets us align dimensions for matrix multiplication when they wouldn’t otherwise match.

Figure 1.63: The Keys matrix transposed: each word’s key vector becomes a column, converting the (5, 4) matrix into (4, 5).

When we transpose the Keys matrix, each row becomes a column. Notice how the first row for “The” [1.4, 1.0, 1.8, 2.2] in the original Keys matrix becomes the first column [1.4, 1.0, 1.8, 2.2] reading downward in the transposed version. The same happens for every word, ”next” becomes the second column, “day” becomes the third column, and so on, transforming our (5, 4) matrix into a (4, 5) matrix ready for multiplication.

Figure 1.64: Multiplying Query (5, 4) by K_T (4, 5) produces the (5, 5) attention scores matrix, capturing every possible word-to-word relationship.

The result of the dot product of Query and Keys vectors is a (5, 5) attention scores matrix. This matrix captures every possible relationship between words.

Interpreting the Attention Scores Matrix

Now that we have our attention scores matrix, let’s understand what it actually tells us.

Figure 1.65: The (5, 5) attention scores matrix: rows represent queries and columns represent keys. Entry (i, j) shows how much word i attends to word j.

Each number in this (5, 5) matrix represents how much one word should attend to another word. The key to reading this matrix is simple:

rows represent queries, and columns represent keys.

Let’s look at some concrete examples.

Figure 1.66: Reading the attention matrix: the entry at row 2, column 1 gives the attention score from “next” (query) to “The” (key).

Finding attention between “next” and “The”:

The word “next” is in position 2, so we look at row 2. The word “The” is in position 1, so we look at column 1. The value at position (2, 1) tells us how much “next” attends to “The.”

Figure 1.67: The entry at row 2, column 2 shows the self-attention score: how much “next” attends to itself.

Finding attention between “next” and itself:

Same word, but the pattern holds. Row 2 for “next” as the query, column 2 for “next” as the key. The value at position (2, 2) shows how much “next” attends to itself.

Here’s where it gets interesting. Each row tells a complete story about one word’s attention pattern.

Figure 1.68: Each row tells a complete story about one word’s attention pattern across all other words in the sentence.

Take the second row, for example. This entire row represents the attention scores between “next” (the query) and all other words (the keys). As you move across the columns, you see how much “next” should attend to “The,” then to “next” itself, then to “day,” then to “is,” and finally to “bright.”

The first row does the same for “The.” The third row does it for “day.” Every row follows this pattern.

Problem with Attention Scores Matrix

We have our attention scores matrix, and it captures the relationships between words. But there’s a fundamental issue we need to address. Look at the second row, which represents how much “next” should attend to all other words. The values might be something like 1.3, 0.9, 1.9, 1.9, and 1.2. These numbers tell us relative importance, but they’re hard to interpret.

Figure 1.69: Raw attention scores do not sum to one and cannot be interpreted as probabilities. We need normalization to create a proper attention distribution.

What we really want is to make clear, intuitive statements like “next should give 30% of its attention to ‘day,’ 25% to ‘is,’ 20% to itself, 15% to ‘The,’ and 10% to ‘bright.’” We want percentages that sum to 100%, or in mathematical terms, probabilities that sum to 1.

This would let us understand attention as a distribution, like slices of a pie chart, where we can immediately see which words matter most. Right now, our raw scores don’t have this property (11.0 + 10.2 + 10.1 + 10.4 + 11.2), you get 52.9, not 1. The values in each row don’t sum to one, and some might even be negative. We can’t interpret them as percentages or probabilities, and this creates two problems.

First, there’s the interpretability issue. We can’t make clear statements about attention distribution. We can’t say “next pays 22% attention to ‘bright’” when the numbers don’t represent percentages. Second, there’s a training stability issue. When training large language models, it’s better if the numbers stay in a controlled range, ideally between 0 and 1. This makes the training process much more stable. The gradients behave better, and the model learns more reliably.

That’s the problem we need to solve, and the solution is to convert attention scores into attention weights. Attention weights have two key properties: they sum to 1 for each row, and each individual weight lies between 0 and 1. This transformation is called normalization.

A Quick Note on Simple Normalization and Softmax

Simple Normalization

The simplest approach to normalization is straightforward. Take each value in a row and divide it by the sum of all values in that row.

Formula

Example

Figure 1.70: Simple normalization divides each value by the row sum, preserving proportions exactly. Softmax exponentiates first, dramatically amplifying differences so the largest value dominates.

These values sum to 1, which is good. The differences are proportional. The value 6 is twice as large as 3, and after normalization, 0.5 is twice as large as 0.25. The proportions are preserved exactly.

Softmax Normalization

Softmax takes a different approach. Instead of dividing by the sum directly, it first exponentiates each value, then normalizes by the sum of exponentials.

Formula

Example

Figure 1.71: Softmax normalization: the largest value (6) receives 93% of the weight while smaller values are heavily suppressed, creating a sharp and decisive attention distribution.

Notice something dramatic. The largest value (6) now dominates completely with 93% of the attention, while the smaller values are heavily suppressed(see the bar graph in the figure). This is the key difference between simple normalization and softmax.

Why Softmax Works Better

Amplification of Differences

Softmax has a crucial property: it amplifies differences. Larger values get disproportionately larger weights, and smaller values get disproportionately smaller weights. This makes the resulting distribution sharper and more decisive.

In our example with simple normalization, the value 6 was six times larger than 1, and after normalization, its weight (50%) was also six times larger than the weight for 1 (8.3%). The proportions stayed exactly the same.

But with softmax, the value 6 gets 93% of the weight, while 1 gets only 0.6%. That’s a ratio of over 150 times! The difference got amplified dramatically. This amplification is exactly what we want in attention mechanisms. When one word should clearly attend to another, softmax makes that relationship strong and clear. The model can make bold, decisive choices about where to focus attention.

Handling Negative Values

Softmax has another important advantage: it handles negative numbers gracefully.

Simple Normalization

We have a problem. A negative probability (-60%) doesn’t make sense. Probabilities must be between 0 and 1.

Softmax

All values are positive! The negative score (-3) simply gets suppressed to nearly zero (0.03%), while the largest value (5) dominates. Softmax automatically ensures all outputs are valid probabilities regardless of input values.

Listing 1.5: Computing Raw Attention Scores

# attention scores for all five tokens
attn_scores = queries @ keys.T     # shape (5, 5)

print(”Attention scores matrix:”)
print(attn_scores)

# attention scores only for the word “next”
idx = 1                            # index 1 is “next”
query_next = queries[idx]          # shape (4,)

keys_T = keys.T                    # shape (4, 5)
attn_scores_next = query_next @ keys_T

print(”\nAttention scores for ‘next’:”)
print(attn_scores_next)

Output

Attention scores matrix:
tensor([[10.3564, 10.4507, 10.9581, 11.3706, 11.7368],
        [10.0395, 10.1182, 10.6117, 11.0026, 11.3353],
        [10.2672, 10.3379, 10.8698, 11.2815, 11.6372],
        [10.5471, 10.6268, 11.1764, 11.5984, 11.9785],
        [11.0007, 11.1326, 11.7058, 12.2208, 12.6405]])

Attention scores for ‘next’:
tensor([10.0395, 10.1182, 10.6117, 11.0026, 11.3353])

The matrix attn_scores contains all raw attention scores before scaling or softmax. Each row corresponds to one query token. Each column corresponds to one key token. Entry (i, j) is the dot product between the query vector for token i and the key vector for token j.

Computing the full matrix in one step is just the matrix form of what we did for a single token earlier. In the second part of the code we select the query for “next” and multiply it with the transposed key matrix. The resulting vector attn_scores_next is simply row one of the full score matrix and shows how strongly “next” matches the key for each word in the sentence, including itself.

[Step 3] Converting Attention Scores to Attention Weight

Figure 1.72: Applying softmax row by row converts raw attention scores into normalized attention weights that sum to one, creating an interpretable probability distribution.

Now let’s apply softmax to convert our attention scores into attention weights. We’ll work through this for one row to see exactly how it works.

Looking at our attention scores matrix, let’s take the second row for “next.” The values are:

These represent how much “next” should attend to “The,” “next,” “day,” “is,” and “bright” respectively.

Step 1: Exponentiate each score

Step 2: Calculate the sum

Step 3: Divide each exponential by the sum

Now we can make clear, interpretable statements: “next” pays 33.6% of its attention to “bright,” 27.6% to “The,” 15.2% to “is,” 12.4% to itself, and 11.2% to “day.”

Figure 1.73: The complete attention weights matrix after softmax: every value is between 0 and 1, and every row sums to 1.

We apply this same softmax operation to every row in our attention scores matrix. Each row gets its own independent softmax transformation, converting raw scores into normalized attention weights that sum to 1. The result is our attention weights matrix, where every value is between 0 and 1, every row sums to 1, and we can finally interpret the numbers as meaningful probabilities.

Figure 1.74: The scaled dot-product attention formula from the original transformer paper.

Before we move forward, there’s something important to clarify about the attention formula we’ve been building. What we just covered was the softmax operation, which converts attention scores into attention weights. But in practice, there are two additional operations that happen before we apply softmax: scaling by the square root of the key dimension, and adding a mask.

Figure 1.75: The complete attention pipeline: compute QK^T , scale by √ d_k, optionally apply a mask, then apply softmax to obtain attention weights

You might wonder why we’re mentioning this now, after already explaining softmax. The reason is pedagogical. Understanding softmax first makes it much easier to grasp why these additional steps matter. If we had introduced all three operations at once, the picture would have been muddier. By learning them in this order, you’ll see not just what these operations do, but why they’re necessary.

If this sounds a bit abstract right now, don’t worry. The next section will clarify everything. We’ll walk through both scaling and masking step by step, and by the end, you’ll understand exactly how they fit into the complete attention mechanism.

1.12 Why Scale Attention Scores?

Figure 1.76: The scaling factor normalizes the variance of the dot product, preventing softmax from producing extremely sharp distributions

In the Transformer model, the attention mechanism calculates scores using the formula:

A critical component of this formula is the scaling factor,

This scaling is not arbitrary; it is essential for stabilizing the training process.

The Problem with Unscaled Scores

Figure 1.77 : As the key dimension d_k increases, the variance of the dot product grows, causing large scores that push softmax into saturation where gradients vanish.

The attention scores are calculated from the dot product of a query vector (Q) and a key vector (K^T). A dot product is the sum of element wise products:

The problem is that as the dimension (d_k) increases, the variance of this dot product also increases. A larger dimension means more terms are being added together, which can lead to the final scores being very large in magnitude.

These large scores are then passed into the softmax function. The softmax function is sensitive to large inputs. If one score is significantly larger than the others, softmax will assign it a probability very close to 1.0, while all other scores will be assigned a probability very close to 0.0. This is known as saturation.

When this happens, the attention becomes “hard” and “spiky,” focusing on only one position. This makes it difficult for the model to learn, as the gradients during backpropagation become extremely small, effectively vanishing and halting the training process for that attention head.

The Statistical Solution

If we assume the components of (Q) and (K) are independent random variables with a mean of 0 and a variance of 1, then their dot product (Q K^T) will have a mean of 0 but a variance of d_k. To normalize this, we want to scale the dot product so that its variance remains 1, regardless of the dimension d_k.

The standard deviation of the dot product is the square root of its variance, which is

By dividing the dot product (Q K^T) by its standard deviation (, we ensure the input to the softmax function has a stable variance of 1.

Let’s compute the scaled attention matrix

Figure 1.78: Computing the scaled attention scores, key dimension is 4, so √ d_k = 2, and every raw score is divided by 2 before softmax.

“Keys Vectors” a matrix with the shape (5, 4).

This means there are 5 key vectors, and each vector has a dimension of 4.

To get the scaled scores, we must divide every number in “Attention Scores” matrix by our new scaling factor, 2.

This is the “Attention Scores”

Calculation (Dividing by 2):

Final “Scaled Attention Scores” (based on d_k=4):

Listing 1.6: Scaling Scores and Computing Attention Weights

d_k = keys.shape[-1]   # key dimension, 4

# scale scores and convert to attention weights for all tokens

scaled_scores = attn_scores / d_k**0.5
attn_weights = torch.softmax(scaled_scores, dim=-1)

print(”Attention weights matrix:”)
print(attn_weights)
print(”\nRow sums:”, attn_weights.sum(dim=-1))

# same thing, but shown explicitly for the word “next”

scaled_scores_next = attn_scores_next / d_k**0.5
attn_weights_next = torch.softmax(scaled_scores_next, dim=-1)

print(”\nScaled scores for ‘next’:”)
print(scaled_scores_next)
print(”Attention weights for ‘next’:”)
print(attn_weights_next)
print(”Sum of weights for ‘next’:”, attn_weights_next.sum())

Output

Attention weights matrix:

tensor([[0.1331, 0.1401, 0.1781, 0.2413, 0.3074],
        [0.1375, 0.1430, 0.1779, 0.2361, 0.3056],
        [0.1342, 0.1394, 0.1797, 0.2406, 0.3060],
        [0.1307, 0.1352, 0.1776, 0.2426, 0.3139],
        [0.1180, 0.1271, 0.1660, 0.2553, 0.3336]])

Row sums: tensor([1., 1., 1., 1., 1.])

Scaled scores for ‘next’:

tensor([5.0198, 5.0591, 5.3059, 5.5013, 5.6677])
Attention weights for ‘next’:

tensor([0.1375, 0.1430, 0.1779, 0.2361, 0.3056])

Sum of weights for ‘next’: tensor(1.)

Raw dot products can grow large when the key dimension increases and they are difficult to interpret. The first step therefore scales the scores by dividing by the square root of the key dimension d_k. This keeps the variance of the scores roughly constant and prevents softmax from producing extremely sharp distributions.

The call to torch.softmax then turns each row of scaled scores into a proper probability distribution. All entries are between zero and one and each row sums to one, as confirmed by the printed row sums. The attention weight at position (i, j) now expresses the fraction of token i’s attention that is assigned to token j.

For example, the vector attn_weights_next shows how the word “next” distributes its attention across the five tokens. In the example above it puts about thirty percent of its weight on the last word, with the remaining seventy percent spread over the earlier words.

[Step 4] From Attention Weights to Context Vectors

Figure 1.79: The final step of self-attention: multiply the attention weights matrix by the Value matrix to produce context vectors.

Brief Note: The calculations described in this section illustrate the core mechanism of scaled dot-product attention. For simplicity, this example does not apply causal or look-ahead masking, which would be essential in a decoder-based model (like a GPT) to prevent a token from “seeing” future tokens.

In the self-attention mechanism, the final step is to compute the context vector for each token. A common misconception is that attention is applied to the original input embeddings. Instead, attention is used to create a weighted sum of a new, transformed representation of the input. This new representation is called the Value (V) matrix.

The Role of the Value (V) Matrix

Figure 1.80: The Value matrix is created by multiplying the input embeddings by a separate weight matrix W_V . It provides the representations that are blended according to attention weights.

Just as we created the Query (Q) and Key (K) matrices by multiplying our input embeddings (X) with trainable weight matrices (W_q and W_k), we create the Value (V) matrix by multiplying the input embeddings with its own trainable weight matrix, W_v.

This transformation is crucial. It allows the model to learn a representation of the input tokens that is specifically optimized for constructing the final contextualized output. While the Key matrix is designed for “being searched” and the Query matrix is for “searching,” the Value matrix is designed to “be blended.”

Instead of directly blending the input vectors, we blend these new Value vectors. This gives the model more flexibility and expressive power.

Calculating the Context Matrix

Figure 1.81: The context matrix is computed as a single matrix multiplication: Attention Weights (5, 5) times Values (5, 4) yields the (5, 4) context matrix containing one context vector per token.

The calculation of the final context matrix is a single matrix multiplication:

Let’s break this down using the dimensions from our examples.

1. Attention Weights (A): This is the (5, 5) matrix of normalized scores we calculated previously. Each row (i) of this matrix represents the “attention” that token (i) pays to every other token (including itself).

2. Value (V) Matrix: This is the (5, 4) matrix of transformed input vectors. Each row corresponds to a token, but it now exists in the “value space” with a dimension of 4. The multiplication is therefore:

The resulting (5, 4) Context matrix contains our five new context vectors, one for each input token. Each of these new vectors has a dimension of 4, matching the dimensionality of our value space.

What is a Context Vector?

Each row in the final Context matrix is the new, “context-aware” vector for its corresponding token. This new vector is a weighted sum of all the Value vectors in the sequence. Let’s illustrate by calculating the context vector for the third token, “day” (row 3).

Figure 1.82: Computing the context vector for “day”: its attention weights are multiplied with the corresponding Value vectors and summed to produce a new representation that blends information from all tokens.

1. Get the Weights: We take row 3 from the Attention Weights matrix. These are the weights from “day” to all other tokens:

2. Get the Values: We use the entire (5, 4) Value matrix.

3. Perform the Weighted Sum: The new context vector for “day” is calculated by multiplying its attention weights by the corresponding Value vectors and summing the results:

Figure 1.83: Detailed calculation showing how the first element of the context vector for “day” is computed as a weighted sum of the first elements of all Value vectors.

Since matrix calculations can sometimes feel overwhelming due to the number of values involved, You can see how how the value in the first column of the context vector for the token ‘day’ is calculated.

This new vector is a blend, or a weighted average, of all the tokens’ “value” representations. The blend is dictated by the attention scores. In this case, the new meaning of “day” is most heavily influenced by the value of “bright” (35%), followed by its own original value (9%), and the values of the other tokens.

Figure 1.84 : The complete transformation: input embeddings (static, context-free) are converted into context vectors (dynamic, context-aware) through the self-attention mechanism.

We began with an input embedding matrix, where each token’s vector represented its meaning in isolation, unaware of its surroundings. The self-attention mechanism transforms these static inputs by projecting them into three new spaces: Query, Key, and Value. By comparing the Query and Key matrices, the model generates a matrix of attention weights. This weight matrix acts as a precise “blending recipe,” quantifying the exact relevance and relationship of every token to every other token in the sequence. The result is the context vector matrix, where each token’s original vector is replaced by a new, context-aware representation. This fundamental transformation from isolated, static meaning to a rich, contextualized representation is the central power of the self-attention mechanism.

Listing 1.7: Computing Context Vectors from Attention Weights

# compute context vectors for all tokens
context = attn_weights @ values      # shape (5, 4)

print(”Context vectors:”)
print(context)
print(”context.shape:”, context.shape)

# context vector for the word “next” only
context_next = attn_weights_next @ values
print(”\nContext vector for ‘next’:”)
print(context_next)

Output

Context vectors:
tensor([[ 2.2118,  2.4971, -0.2800,  0.8843],
        [ 2.2038,  2.4754, -0.2794,  0.8741],
        [ 2.2062,  2.4855, -0.2796,  0.8797],
        [ 2.2100,  2.4933, -0.2802,  0.8884],
        [ 2.2229,  2.5209, -0.2808,  0.9064]])
context.shape: torch.Size([5, 4])

Context vector for ‘next’:
tensor([ 2.2038,  2.4754, -0.2794,  0.8741])

The final step in self attention is to combine the value vectors using the attention weights as coefficients. Every context vector is a weighted sum of all value vectors, where the weights come from the corresponding row in the attention matrix.

The matrix product

context = attn_weights @ values

implements this operation for all five tokens at once. Since attn_weights has shape 5, 5 and values has shape 5, 4, the result has shape 5, 4. Each row in context is the new context aware representation of one token in the sentence.

The row context_next shows the updated representation for “next”. It lives in the same 4 dimensional space as the value vectors, but it now encodes information aggregated from all tokens in the sentence according to the learned attention pattern. This is exactly the transformation the theory section describes when it talks about going from static input embeddings to dynamic context vectors.

Listing 1.8: Packaging Self-Attention into a PyTorch Module

import torch.nn as nn

class SelfAttention(nn.Module):

    def __init__(self, d_in, d_out):
        super().__init__()
        self.W_query = nn.Parameter(torch.randn(d_in, d_out))
        self.W_key   = nn.Parameter(torch.randn(d_in, d_out))
        self.W_value = nn.Parameter(torch.randn(d_in, d_out))

    def forward(self, x):
        queries = x @ self.W_query    # (seq_len, d_out)
        keys    = x @ self.W_key      # (seq_len, d_out)
        values  = x @ self.W_value    # (seq_len, d_out)

        attn_scores  = queries @ keys.T
        d_k = keys.shape[-1]
        attn_weights = torch.softmax(
            attn_scores / d_k**0.5, dim=-1
        )

        context = attn_weights @ values
        return context

torch.manual_seed(123)
sa = SelfAttention(d_in=8, d_out=4)
out = sa(inputs)
print(out.shape)

output

torch.Size([5, 4])

This class collects the individual steps of scaled dot product self attention into a reusable component. The constructor creates three trainable parameter matrices, each of shape d_in, d_out. When the module is part of a model, these parameters will be updated by the optimiser during training.

The forward method implements the same pipeline we derived by hand. It projects the input embeddings into queries, keys and values, computes all attention scores with a single matrix product, scales and normalises them with softmax, and finally uses the resulting weights to blend the value vectors into context vectors.

The last two lines create an instance of the layer with d_in equal to 8 and d_out equal to 4, apply it to the input sentence and print the shape of the output. The result 5, 4 confirms that for a sequence of five tokens the layer returns five context vectors, each living in the 4 dimensional space of the attention head. This is exactly the representation that will be passed on to the feed forward network in the transformer block.

1.13 Causal & Masked Attention

In the preceding section, we examined how self-attention transforms input embeddings into context-aware vectors. However, that explanation omitted a vital component for generative models: causal attention. This mechanism is fundamental, as it ensures that the model respects the sequential order of text and does not “cheat” by looking at future tokens.

Now that you have a solid understanding of the full self-attention pipeline, it is the perfect time to introduce this concept. This masking step is applied directly to the attention scores, just before the softmax function, to block any information from subsequent positions in the sequence.

Figure 1.85: Causal masking ensures that when processing token i, the model can only attend to tokens at positions 0 through i, preventing information leakage from future tokens.

Large language models like ChatGPT generate text by predicting one token at a time. Each predicted token is appended to the input, creating a growing context window used to predict the next token.

Figure 1.86: Sequential text generation: the model predicts one token at a time, appending each prediction to the input before generating the next.

This sequential process imposes a fundamental constraint: when computing the context vector for any token, only that token and preceding tokens should have influence. Future tokens must not contribute, as they have not yet been generated.

Figure 1.87 The causal constraint: when processing “science,” the model can only access “Computer” and itself. “is,” “the,” and subsequent tokens are masked out.

Consider the sequence “Computer” , “science”, “is” , “the”, “study...” When processing “science” , the model should only access itself and “Computer” . It must not see “is”, “the”, or any subsequent tokens. Similarly, “Computer” should only attend to itself, while “is” can attend to “Computer”, “science” , and itself, but not to the tokens that follow it.

Figure 1.88: A lower-triangular attention matrix enforces the causal constraint. All entries above the diagonal are zero, ensuring that no token attends to future positions.

To enforce this constraint, we use masked attention. For each token acting as a query, we mask all keys corresponding to future positions by setting their attention scores to zero. The token “Computer” has a non-zero attention score only with itself. The token “science” has attention scores with “Computer” and itself, but zero scores with all future tokens like “is” and “the” . This masking creates a lower triangular attention matrix where all entries above the diagonal are zero.

After masking, we normalize the remaining attention weights in each row to sum to one. For “Computer” , the single remaining weight is set to one. For “science”, the two remaining weights are normalized so their sum equals one. This normalization is achieved by summing the non-masked weights in each row and dividing each weight by this sum, creating a proper probability distribution over the tokens each query can attend to. This mechanism, called causal attention or masked self attention, enables language models to generate coherent text while respecting the sequential nature of prediction.

Implementing Causal Attention through Zero Masking

To implement causal attention, we can begin with the raw attention scores. Consider the 5x5 attention score matrix for the sequence “computer science is the study”

Step 1: Initial Attention Scores

This is the raw, unnormalized matrix.

Figure 1.89: The raw 5×5 attention score matrix before any masking is applied.

In PyTorch, we can construct a lower triangular mask using the torch.tril()

torch.tril(MatrixA)

Figure 1.90: The torch.tril() function creates a lower-triangular mask: ones on and below the diagonal, zeros above.

This function creates a matrix where elements on and below the diagonal are ones, while elements above the diagonal are zeros. The size of this mask matrix must match the dimensions of our attention score matrix, which is determined by the context length.

The context length is simply the number of tokens currently in the sequence. For the example sequence “computer,” “science,” “is,” “the,” “study,” the context length is five.

Listing 1.9: Understanding Triangular Masks with torch.tril

import torch

# A simple 3x3 example matrix

A = torch.tensor([
    [1., 2., 3.],
    [4., 5., 6.],
    [7., 8., 9.],
])

print(”A:”)
print(A)

# Lower-triangular version of A

A_tril = torch.tril(A)

print(”\ntorch.tril(A):”)
print(A_tril)

# A pure mask built from ones

mask_ones = torch.tril(torch.ones_like(A))

print(”\nLower-triangular mask from ones:”)
print(mask_ones)

# Using the mask to zero out the upper triangle

A_masked = A * mask_ones

print(”\nA * mask_ones:”)
print(A_masked)

Output

A:
tensor([[1., 2., 3.],
        [4., 5., 6.],
        [7., 8., 9.]])

torch.tril(A):
tensor([[1., 0., 0.],
        [4., 5., 0.],
        [7., 8., 9.]])

Lower-triangular mask from ones:
tensor([[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.]]).tril()
tensor([[1., 0., 0.],
        [1., 1., 0.],
        [1., 1., 1.]])

A * mask_ones:
tensor([[1., 0., 0.],
        [4., 5., 0.],
        [7., 8., 9.]])

The function torch.tril returns the lower triangular part of a matrix: everything on and below the main diagonal is kept, everything above it is set to zero.

torch.tril(A) takes the original data in A and zeroes the entries above the diagonal.

torch.tril(torch.ones_like(A)) builds a mask: ones on and below the diagonal, zeros above.

Multiplying A by this mask with A * mask_ones keeps the lower triangle and zeroes out the upper triangle.

Causal attention uses exactly this idea. Instead of masking a 3×3 matrix, we mask a seq_len × seq_len attention score matrix so that token i can only see tokens 0..i and not future tokens.

Building a causal mask for a 5-token sequence

Assume we have a sequence of five tokens, such as:

[”computer”, “science”, “is”, “the”, “study”]

The attention scores attn_scores are a 5 x 5 matrix.

Listing 1.10: Building a Causal Mask for a 5-Token Sequence

seq_len = 5

# Build a 5x5 causal mask

causal_mask = torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool))

print(”Causal mask (True = allowed, False = masked):”)
print(causal_mask)

Output

Causal mask (True = allowed, False = masked):
tensor([[ True, False, False, False, False],
        [ True,  True, False, False, False],
        [ True,  True,  True, False, False],
        [ True,  True,  True,  True, False],
        [ True,  True,  True,  True,  True]])

This mask allows each token to only attend to itself and all previous tokens, blocking attention to future tokens to prevent information leakage.

Step 2: Create and Apply the Mask

Figure 1.91: Applying the lower-triangular mask to the attention scores through element-wise multiplication zeros out all future-token entries.

The mask is created as a 5x5 lower triangular matrix of ones.

We then apply this mask to our attention scores through element-wise multiplication.

This operation sets all elements in the upper triangle to zero while preserving the lower triangular elements.

After multiplication, the attention scores that corresponded to future tokens are now zero, while scores for current and past tokens remain unchanged.

Figure 1.92: After masking, the rows are not yet normalized. Each row must be divided by its sum to form a valid probability distribution.

However, this masking alone is insufficient. The rows are not normalized and do not sum to one, which violates the requirement that attention weights form a probability distribution. We must normalize each row by dividing every element by the row sum, ensuring that the non-zero weights in each row sum to one.

Step 3: Row-wise Normalization

Figure 1.93: Row-wise normalization of the masked attention scores: dividing each entry by its row sum produces attention weights that sum to one.

First, we find the sum of each row in the masked matrix:

Now, we divide each element by its row sum to get the final weights.

This normalization step completes the implementation of this (flawed) version of causal attention.

The Problem: Data Leakage in Attention Computation

At first glance, the masked self-attention approach appears to solve the problem of preventing queries from attending to future tokens. We mask the upper triangular portion of the attention matrix and normalize the remaining weights. However, a closer examination reveals a critical flaw in this approach.

To understand the problem, we must revisit how attention weights are computed. The process begins with the construction of query, key, and value matrices. We compute the dot product between the query matrix and the transpose of the key matrix, producing attention scores that indicate how much each token attends to every other token. These scores are then scaled by dividing each element by the square root of the key dimensionality, yielding the scaled dot product. The crucial step occurs when we convert the scaled dot product into attention weights by applying the softmax function row-wise.

Here lies the problem. When we apply softmax to a row in the scaled dot product matrix, the denominator considers all elements in that row, including those corresponding to future tokens. Consider the first row, which corresponds to the token “computer.” When computing the softmax, the summation in the denominator includes the scaled dot product values for all tokens, including “science,” “is,” “the,” and “study.” Similarly, for the second row corresponding to “science,” the softmax denominator includes values from future tokens “is,” “the,” and “study.”

Figure 1.94: Data leakage: when softmax is computed before masking, the denominator already includes contributions from future tokens, subtly influencing the attention weights.

By the time we obtain the attention weights, each element has already been influenced by future tokens through the softmax normalization. Masking the attention weights after this computation does not eliminate this influence. The information from future tokens has already leaked into the computation through the softmax denominator.

This phenomenon is termed data leakage. We intended masked self-attention to prevent queries from accessing information about future tokens, but this prevention fails because the leakage occurs during the softmax computation itself.

Figure 1.95: To prevent data leakage, masking must be applied before softmax so that future tokens are excluded from the softmax denominator entirely.

To properly implement causal attention, we must intervene before applying softmax. The softmax denominator for each row should only consider elements up to and including the current token position. Future keys must be excluded from this summation entirely. The masking operation must occur at the scaled dot product stage, before the softmax function is applied, to truly prevent data leakage.

The Solution: Masking with Negative Infinity

The solution to the data leakage problem involves a clever technique that applies masking before the softmax operation. Instead of zeroing out attention weights after computing softmax, we assign negative infinity values to the positions we want to mask in the scaled dot product matrix.

Figure 1.96: Negative infinity masking: upper-triangular entries are set to −∞ before softmax, which maps them to exactly zero probability while correctly normalizing over visible tokens.

The process works as follows. After computing the attention scores through the dot product between the query and key matrices, and before applying softmax, we set all upper triangular elements to negative infinity. These negative infinity values persist even after scaling by the square root of the key dimensionality, since dividing negative infinity by any finite number still yields negative infinity.

To understand why this approach works, consider how the softmax function behaves with negative infinity values.

Formula

Suppose we have a row containing the values 2, 3, and 5.

Then the values would be like below

Now consider what happens when we want to mask the last two elements. We replace them with negative infinity, giving us the sequence 2, negative infinity, negative infinity.

When we apply softmax, the first element is like this

The key insight is that

Therefore, the first element simplifies to 1

What happens to the masked elements?

Consider the second position containing negative infinity, it becomes zero. The third position similarly becomes zero as well.

After softmax, our sequence transforms into 1, 0, 0.

Let’s take one more example,

Now we are going to mask only the third element

The non-masked softmax values, 0.2689 and 0.7311, add up to one, which confirms they form a proper probability distribution over the unmasked elements

The elegance of this method lies in its automatic normalization property. By setting masked positions to negative infinity before softmax, the resulting attention weights naturally satisfy two requirements.

Figure 1.97: The complete causal masking pipeline: (1) compute scaled dot products, (2) set upper-triangular entries to −∞, (3) apply softmax. Masked positions become exactly zero with no data leakage.

First, all masked positions become exactly zero after softmax. Second, the remaining non-masked weights in each row automatically sum to one, as the softmax function guarantees normalization over all finite input values. This eliminates the need for any additional normalization step after masking, solving the data leakage problem while maintaining the mathematical properties required for attention weights.

Now we apply the causal mask to the scaled attention scores.

Listing 1.11: Applying the Causal Mask to Scaled Attention Scores

import math

scaled_scores = attn_scores / math.sqrt(d_k)
print(”Scaled scores (unmasked):”)
print(scaled_scores)

# Build the boolean causal mask again

seq_len = attn_scores.size(0)
causal_mask = torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool))

# Apply the mask: set disallowed positions to -inf

masked_scaled_scores = scaled_scores.masked_fill(~causal_mask, float(”-inf”))

print(”\nScaled scores with causal mask:”)
print(masked_scaled_scores)

Output

Scaled scores (unmasked):
tensor([[0.5000, 1.0000, 1.5000, 2.0000, 2.5000],
        [0.7500, 1.2500, 1.7500, 2.2500, 2.7500],
        [1.0000, 1.5000, 2.0000, 2.5000, 3.0000],
        [1.2500, 1.7500, 2.2500, 2.7500, 3.2500],
        [1.5000, 2.0000, 2.5000, 3.0000, 3.5000]])

Scaled scores with causal mask:
tensor([[0.5000,   -inf,   -inf,   -inf,   -inf],
        [0.7500, 1.2500,   -inf,   -inf,   -inf],
        [1.0000, 1.5000, 2.0000,   -inf,   -inf],
        [1.2500, 1.7500, 2.2500, 2.7500,   -inf],
        [1.5000, 2.0000, 2.5000, 3.0000, 3.5000]])

We first apply the usual scaling factor 1/sqrt(d_k) to the scores.

Then

masked_scaled_scores = scaled_scores.masked_fill(~causal_mask, float(”-inf”))

does two things:

• ~causal_mask inverts the boolean mask. Positions that were False (future tokens) become True.

• masked_fill writes -inf into those positions.

All allowed positions (on and below the diagonal) keep their original scaled scores. Disallowed positions become negative infinity. This guarantees that, when we apply softmax next, future tokens will contribute zero probability.

Now let’s apply softmax to make causal attention weights from masked scores

Listing 1.12: Computing Causal Attention Weights with Softmax

attn_weights_causal = torch.softmax(masked_scaled_scores, dim=-1)

print(”Causal attention weights:”)
print(attn_weights_causal)
print(”\nRow sums:”, attn_weights_causal.sum(dim=-1))

Output

Causal attention weights:
tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.3777, 0.6223, 0.0000, 0.0000, 0.0000],
        [0.1863, 0.3072, 0.5065, 0.0000, 0.0000],
        [0.1015, 0.1674, 0.2760, 0.4551, 0.0000],
        [0.0580, 0.0956, 0.1577, 0.2599, 0.4288]])

Row sums: tensor([1.0000, 1.0000, 1.0000, 1.0000, 1.0000])

Softmax is now applied to rows that contain finite scores only on and below the diagonal, while masked positions are set to −inf so their exponentials become zero. As a result, the first token can only attend to itself, giving a distribution like [1, 0, 0, 0, 0]; the second token attends only to the first two positions, and those two weights sum to one; the last token can attend to all five positions, so its row is a full probability distribution over the sequence. In every case, all entries above the diagonal are exactly zero, so no token ever attends to future tokens, and each row still sums to one, so every row is a valid attention distribution. Because this masking is applied before softmax, there is no data leakage from future positions, which gives us true causal attention.

1.14 Causal Attention with Dropouts

Concept of Dropout

Before exploring how dropout is applied in causal attention, we first review the concept of dropout and its purpose in neural networks. Dropout is a regularization technique designed to prevent overfitting and ensure that all neurons in a network contribute meaningfully to the learning process.

Figure 1.98: Dropout randomly deactivates neurons during training, forcing lazy neurons to participate and preventing the network from relying on a few dominant connections.

Consider a neural network layer where certain neurons dominate the computation while others contribute minimally. For example, in a layer with five neurons, one neuron might have very large weights while the other two have small weights. This dominant neuron effectively controls the output of the layer, while the other neurons become what we call lazy neurons. These lazy neurons do not significantly influence the forward pass and consequently do not learn useful representations during training. The network essentially overfits by relying too heavily on a subset of neurons.

Dropout addresses this problem by randomly deactivating neurons during training. In each forward pass through the network, neurons are switched off with a certain probability, typically 0.5. This means that statistically, half of the neurons will be deactivated in any given training iteration. The selection is probabilistic and automatic, not manual.

When a previously dominant neuron is switched off, the lazy neurons must participate in the forward propagation. During backpropagation, the weights of these previously inactive neurons must now be adjusted to minimize the loss. Without dropout, if the forward propagation relies entirely on one or two dominant neurons that already produce low loss, the weights of the lazy neurons would never be updated. By forcing different subsets of neurons to be active across training iterations, dropout ensures that all neurons learn to extract useful features from the data.

This technique prevents the network from becoming overly dependent on specific neurons and encourages the development of more robust, distributed representations across the entire network.

Why Dropout Matters in Attention Mechanisms

When we build attention mechanisms for language models, we sometimes encounter a problem where certain words become overly dependent on each other.

Figure 1.99: Without dropout, the word “study” may pay excessive attention to specific earlier words, memorizing patterns rather than learning general language rules.

Consider the sentence “computer science is the study.” If the word “study” pays excessive attention to specific earlier words, the model might memorize these particular patterns rather than learning general language rules. This excessive dependency between tokens can hurt the model’s ability to generalize to new sentences.

Dropout offers an elegant solution to this problem. By randomly removing some attention connections during training, we force the model to learn more robust patterns that don’t rely on any single strong connection.

Figure 1.100: Causal attention with dropout applied: some valid attention connections (shown in red) are randomly zeroed out during training while the causal mask is preserved.

Let’s walk through where dropout fits in. Consider our example with 5 input tokens: “computer”, “science”, “is”, “the”, “study”. Each token gets represented as a vector, and as you have seen in the self attention section, we have Attention weights.

For unidirectional attention (also called causal attention), we mask the upper triangle of this matrix. This ensures that each token can only attend to previous tokens and itself, not future ones. Looking at the visualization, “computer” can only see itself , “science” can see “computer” and itself, “is” can see the first three words, and so on. The gray areas represent these masked positions where attention is blocked.

Here’s where dropout comes in. After obtaining the attention weights, we randomly set some of them to zero with probability p. In the right image, we see unidirectional attention with dropout applied. The red boxes highlight positions where dropout has been applied, effectively zeroing out those attention connections.

Notice how dropout respects the causal structure. It only affects the valid attention weights (the lower triangular part), never touching the already masked upper triangle. Some attention weights that were previously active are now dropped out (shown within red boxes), forcing the model to rely on different connection patterns.

The Scaling Factor Explained

Figure 1.101: Dropout scaling: when connections are dropped with probability p, the remaining weights are scaled by 1/( 1−p) to maintain consistent expected output magnitude between training and inference.

When we apply dropout, there’s a crucial detail:

we need to scale up the remaining weights. This scaling maintains consistent behavior between training and inference. Suppose in row 4, we originally have attention weights distributed across the first four positions. If dropout with probability 0.5 removes half these connections, the remaining weights need to be doubled to maintain the same expected output magnitude. For example, referring to the matrix, if the fourth row originally had attention weights [0.25, 0.26, 0.25, 0.24] for its four non-zero positions, and dropout with probability 0.5 zeros out two of them, such as the second and fourth values, the remaining two might become [0.50, 0, 0.50, 0] after scaling by 2. This ensures the total signal strength remains consistent.

After applying dropout and scaling, each row of the attention matrix still represents a valid probability distribution over the tokens that can be attended to, just with fewer active connections.

This technique has become a standard component in transformer architectures, contributing to their remarkable success in natural language processing tasks. The simple act of randomly removing connections, combined with proper scaling, creates a powerful regularization effect that helps these models achieve better performance on unseen data.

Here let’s add dropout to the causal attention weights and then compute the context vectors.
We keep the same attn_weights_causal as above and assume we already have a values matrix of shape (5, 4) from the self-attention section.

Listing 1.13: Applying Dropout to Causal Attention Weights

dropout = torch.nn.Dropout(p=0.5)

torch.manual_seed(0)  # to get a stable example mask
attn_weights_causal_drop = dropout(attn_weights_causal)

print(”Causal attention weights before dropout:”)
print(attn_weights_causal)

print(”\nCausal attention weights after dropout (training mode):”)
print(attn_weights_causal_drop)
print(”Row sums after dropout:”, attn_weights_causal_drop.sum(dim=-1))

# Use the dropped weights to compute context vectors
context_causal = attn_weights_causal_drop @ values

print(”\nCausal context vectors with dropout:”)
print(context_causal)
print(”context_causal.shape:”, context_causal.shape)

Output

Causal attention weights before dropout:
tensor([[1.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.3777, 0.6223, 0.0000, 0.0000, 0.0000],
        [0.1863, 0.3072, 0.5065, 0.0000, 0.0000],
        [0.1015, 0.1674, 0.2760, 0.4551, 0.0000],
        [0.0580, 0.0956, 0.1577, 0.2599, 0.4288]])

Causal attention weights after dropout (training mode):
tensor([[2.0000, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.7554, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.3726, 0.6144, 0.0000, 0.0000, 0.0000],
        [0.2030, 0.3348, 0.0000, 0.9102, 0.0000],
        [0.0000, 0.1912, 0.3154, 0.5198, 0.0000]])

Row sums after dropout:
tensor([2.0000, 0.7554, 0.9870, 1.4480, 1.0264])

Causal context vectors with dropout:
tensor([[ 4.2400,  5.0200, -0.5600,  1.5400],
        [ 1.6020,  1.7000, -0.1600,  0.3700],
        [ 2.2640,  2.4060, -0.2700,  0.7000],
        [ 2.4240,  2.6460, -0.2800,  0.9700],
        [ 2.2670,  2.5410, -0.2800,  0.9100]])
context_causal.shape: torch.Size([5, 4])

Dropout randomly sets some attention weights to zero during training. Each weight is kept with probability 1 - p and dropped with probability p, and the remaining weights are scaled by 1 / 1 - p so that the expected total attention remains unchanged. Importantly, dropout never unblocks future positions, so all entries above the diagonal stay at zero and the causal structure is preserved. It only thins out the valid lower triangular connections. After this regularisation step, the dropped attention weights are still used in the usual way to compute the context vectors by multiplying them with the value matrix.

Finally, the context vectors are computed as

context_causal = attn_weights_causal_drop @ values

Each row of context_causal is the context vector for a token under causal attention with dropout. These vectors have shape (5,4), matching the number of tokens and the attention head dimension, and are what you feed into the feed-forward network inside a transformer block during training of generative models.

1.15 Summary of Self-Attention

Figure 1.102: Complete self-attention pipeline: Input embeddings are projected into Q, K, V; attention scores are computed via QK^T , scaled, optionally masked, passed through softmax (with optional dropout), and multiplied by V to produce context vectors.

Before jumping into what multi-head attention is, let’s summarize the self-attention mechanism that we covered in the previous section. So the Process begins when the input tokens are converted into the Input Embedding matrix. This matrix is linearly projected into three distinct matrices: Query (Q), Key (K), and Value (V). This is achieved by multiplying the Input Embedding matrix by three separate, trainable weight matrices. Once Q, K, and V are generated, the input embedding matrix is no longer needed.

The core attention calculation starts by computing the Attention Scores. This is done, as shown in the first ‘MatMul’ step, by taking the dot product of the Q matrix with the transpose of the K matrix. These raw scores are then normalized in the ‘Scale’ step, where they are divided by the square root of the Key’s dimension. This scaling is crucial for stabilizing the gradients during training.

Following scaling, an ‘Optional Mask’ can be applied. This step is essential for implementing causal attention, where it masks out all scores corresponding to future tokens, ensuring a token can only attend to itself and previous tokens. Next, the SoftMax function is applied across each row of the scaled (and possibly masked) scores. This converts the scores into positive values that sum to one, effectively turning them into the final attention weights. An ‘Optional Dropout’ layer can be applied here to prevent overfitting.

In the final step, these attention weights are multiplied by the Value (V) matrix, as shown in the second ‘MatMul’ operation. This produces the Context Vector (Z). Each row of this Z matrix is a new, contextually enriched vector for the corresponding input token, as it now contains a weighted combination of information from all other tokens it was allowed to attend to.

Limitations of Self-Attention Mechanisms

Figure 1.103: Linguistic ambiguity: “The artist painted the portrait of a woman with a brush” has two valid interpretations, the brush is the tool for painting, or the woman in the portrait holds a brush.

A significant problem with a single self-attention mechanism is its limited ability to effectively handle linguistic ambiguity. This challenge can be illustrated with the sentence:

This statement has two distinct and valid interpretations.

The first interpretation is that the artist used a brush as a tool to perform the action of painting. In this context, the phrase “with a brush” modifies the verb “painted”.

The second interpretation is that the subject of the painting is a woman who is holding a brush. Here, “with a brush” modifies the “woman” or “portrait”.

A single self-attention layer may struggle to capture both of these potential relationships simultaneously. It might incorrectly average these dependencies or fixate on only one, resulting in a contextual vector that fails to represent the full nuance of the sentence.

Figure 1.104: A single attention head can only produce one attention matrix, either capturing the tool interpretation or the subject interpretation, but not both simultaneously.

The first interpretation, where the artist uses the brush, would generate an attention score matrix where the word “artist” has a high attention score with “brush”.

The second interpretation, where the woman in the in the portrait is holding a brush, would generate a completely different matrix where “woman” and “portrait” have high attention scores with “brush”.

A single self-attention layer can only produce one of these attention matrices. It will either settle on one perspective or create an unhelpful average of the two. This results in a context vector that fails to represent the full richness and potential ambiguity of the input, limiting the model’s ability to understand the multiple angles or meanings present in complex language.

This limitation demonstrates the need for a more robust mechanism, which multi-head attention addresses.

1.16 Intuition of Multi-Head Attention

Figure 1.105: Multi-head attention: the same input is fed into multiple independent self-attention heads in parallel, each learning a different set of relationships. Their outputs are concatenated into a single enriched context matrix.

Since a single self-attention mechanism is limited to capturing only one perspective from an input sequence, the solution is to use multiple self-attention mechanisms in parallel. This architecture is known as multi-head attention. The core idea is that the same input embedding matrix is fed into several independent self-attention “heads”. Each head produces its own distinct context vector matrix, effectively learning a different set of relationships or focusing on a different aspect of the input, such as one head capturing verb-centric relationships while another captures a different semantic nuance. These individual context vector matrices, each representing a unique perspective, are then combined or merged. This process results in a single, final context vector matrix that is much richer, as it amalgamates the multiple perspectives captured by all the individual heads, leading to a more comprehensive representation of the input.

To implement multi-head attention, we must adapt the query, key, and value matrix operations to support multiple attention mechanisms operating in parallel. Having established the limitations of a single self-attention layer, the goal is to see how we can practically implement a system with multiple heads, for instance, a two-head attention mechanism. The procedure will demonstrate how to generate two independent sets of attention scores and two corresponding context vector matrices. This parallel processing is the core of multi-head attention, as it allows the model to produce multiple, distinct representations, with each head capturing a different perspective or set of relationships from the input sequence.

Figure 1.106: From single-head to two-head attention: the input embedding matrix is processed by separate weight matrices for each head, producing independent sets of Q, K, and V vectors.

The process begins with the input embedding matrix. Using the example sentence “The next day is bright”, which has 5 tokens, we start with an input embedding matrix. As illustrated, each token is represented by an embedding of eight dimensions. This configuration results in an input embedding matrix with dimensions of 5 by 8. The goal of a two head attention mechanism is to transform this single input matrix into two distinct context vector matrices, with each one capturing a different perspective.

To establish a baseline, recall the procedure for a single attention head. In this case, the 5 by 8 input embedding matrix is multiplied by three separate trainable weight matrices. As shown in the diagram, these are the Query Weight Matrix (W_q) with dimensions 8 by 4, the Keys Weight Matrix (W_k) with dimensions 8 by 4, and the Values Weight Matrix (W_v) with dimensions 8 by 4. This matrix multiplication operation produces a Query Vectors matrix (5 by 4), a Keys Vectors matrix (5 by 4), and a Values Vectors matrix (5 by 4). This single set of query, key, and value matrices is what the multi head attention mechanism will expand upon.

Figure 1.107: The head dimension: with d_out = 4 split across 2 heads, each head operates on a reduced dimension of 2. Weight matrices W_k1 and W_k2 each have shape (8, 2).

To transition from a single head to a multi head mechanism, such as one with two heads, the first step is to adapt the trainable weight matrices. Instead of one single query weight matrix (W_q), we now initialize two separate matrices, W_q1 and W_q2, one for each head. This same division is applied to the key and value matrices, creating W_k1, W_k2, W_v1, and W_v2}.

Figure 1.108: Each head operates on a reduced subspace. Head 1 produces Q1, K1, V1 of shape (5, 2); Head 2 produces Q2, K2, V2 of shape (5, 2).

The dimensions of these new matrices are determined by the “head dimension”. This value is calculated by dividing the original total output dimension (d_out) by the number of heads. For example, if the original d_out was 4, for a two head system, the head dimension is 4 divided by 2, which equals 2. This means that while the original weight matrices might have been 8 by 4, each new head specific matrix (like W_k1 and W_k2) will have dimensions of 8 by 2. The main idea of this step is to create multiple, smaller copies of the trainable W_q, W_k, and W_v matrices. As a direct consequence, this will naturally produce multiple sets of query vectors, key vectors, and value vectors when multiplied with the input embeddings.

Figure 1.109: Parallel attention computation: the input embeddings are projected through head-specific weight matrices to produce independent Q, K, V matrices for each head

Multi-head attention extends the basic attention mechanism by creating parallel attention computations that can capture different types of relationships simultaneously. Starting with the input embedding matrix of dimensions 5×8 for our five tokens, the process splits attention into multiple heads. For a two-head configuration with output dimension of 4, each head operates on a reduced dimension of 2. The input embeddings are multiplied with separate weight matrices to produce Query, Key, and Value matrices for each head. Head 1 generates Q1, K1, and V1 matrices, while Head 2 produces Q2, K2, and V2 matrices, all with dimensions 5×2 to match the five tokens and head dimension of 2.

Figure 1.110: Each head independently computes its (5, 5) attention score matrix by multiplying Q with K^T , maintaining the ability to capture all pairwise token relationships.

Each head then independently computes its attention scores by multiplying its Query matrix with the transposed Key matrix.

The Query matrix for each head has shape (5,2) and the Key matrix has shape (2,5), where 5 represents the number of tokens in the sequence. When we multiply these matrices (Q times K transpose), we obtain a (5,5) attention score matrix for each head.

Figure 1.111: Two separate (5, 5) attention score matrices from two heads: each can capture different types of relationships in the data.

This is the crucial insight: although each head works with half the original dimension, the resulting attention score matrix maintains the full (5,5) shape, which represents relationships between all token pairs. This means that with 2 heads, we generate two separate (5,5) attention score matrices rather than one. Each matrix can capture different types of relationships in the data.

Figure 1.112: Each head independently applies scaling and softmax to produce its own attention weight matrix, then computes context vectors by multiplying with its Value matrix.

These attention scores then undergo the standard scaling and softmax normalization within each head independently. After computing the context vectors by multiplying the attention weights with the Value matrices, the outputs from all heads are concatenated back together to restore the original output dimension. This architecture allows the model to simultaneously learn and represent multiple perspectives of token relationships without increasing the computational cost compared to having a single attention mechanism with the full dimension.

Figure 1.113: Context vectors from both heads: Head 1 produces a (5, 2) context matrix and Head 2 produces another (5, 2) context matrix, each capturing a different perspective.

Once we have the attention weight matrices for both heads, each of dimension, we proceed to compute the context vectors. This is done by multiplying each head’s attention weights with its corresponding Value matrix.

For Head 1, we multiply the (5,5) attention weight matrix with the (5,2) Value matrix V1, producing a (5,2) context matrix. Similarly, for Head 2, we multiply its (5,5) attention weight matrix with the (5,2) Value matrix V2, yielding another (5,2) context matrix. Each context matrix represents how the tokens should be represented based on the attention patterns learned by that particular head.

The final step in multi-head attention involves concatenating the context matrices from all heads to produce a unified output representation. Each head generates its own context matrix by multiplying its attention weights with its corresponding Value matrix.

Figure 1.114: Concatenation: the (5, 2) outputs from Head 1 and Head 2 are concatenated along the feature dimension to recover the original output dimension, forming a (5, 4) final context matrix.

In our example with two heads processing a sequence of 5 tokens, Head 1 produces a context matrix representing the first perspective on token relationships, while Head 2 independently generates another context matrix capturing a second, distinct perspective. To combine these complementary perspectives, the context matrices are concatenated along the feature dimension. Specifically, the (5,2) matrix from Head 1 is placed alongside the (5,2) matrix from Head 2, forming a single (5,4) output matrix. This concatenation operation merges the outputs horizontally, stacking the feature vectors from each head side by side for every token position. The resulting (5,4) final context matrix maintains the sequence length of 5 tokens while recovering the original output dimension of 4, which equals the head dimension multiplied by the number of heads. This concatenated representation now contains the enriched information from both attention heads, allowing each token’s final representation to simultaneously encode multiple types of relationships and patterns discovered by the different heads. The concatenated output serves as the complete output of the multi-head attention mechanism and is typically passed through a final linear projection layer before proceeding to subsequent layers in the transformer architecture.

The Dimensional Trade-Off in Multi-Head Attention

While multi-head attention offers significant advantages in capturing diverse perspectives, it does involve a fundamental trade-off in its design. When the output dimension is split across multiple heads, each head operates with a reduced dimensionality compared to single-head attention. In our example with an output dimension of 4 split into 2 heads, each head works with only 2 dimensions rather than the full 4 dimensions that would be available in single-head attention. This reduction in per-head dimensionality means that each individual head has a smaller representational capacity and fewer parameters to capture nuanced patterns within its specific perspective. With fewer dimensions to work with, each head may be limited in the complexity and detail of the relationships it can encode. However, this apparent limitation is offset by the increased number of perspectives that can be learned in parallel. The architecture essentially implements a divide-and-conquer strategy: instead of attempting to capture all types of token relationships within a single high-dimensional space, the model distributes this learning task across multiple specialized heads, each focusing on different aspects of the input. While one head might capture syntactic dependencies with its 2 dimensions, another head simultaneously learns semantic relationships with its own 2 dimensions. This parallelization allows the model to explore a broader range of attention patterns across the same computational budget. The concatenation of outputs from all heads ultimately reconstructs the full output dimension, ensuring that the combined representation benefits from multiple complementary perspectives. Thus, while each head operates with reduced capacity, the overall multi-head architecture achieves greater expressiveness through diversification, making this trade-off worthwhile for most applications.

Listing 1.14: Creating the Input Tensor for Multi-Head Attention

import torch
torch.manual_seed(123)
torch.set_printoptions(precision=3, suppress=True)

# b, num_tokens, d_in = (1, 3, 6)
x = torch.tensor([[
    [1.0, 2.0, 3.0, 4.0, 5.0, 6.0],   # “The”
    [6.0, 5.0, 4.0, 3.0, 2.0, 1.0],   # “kid”
    [1.0, 1.0, 1.0, 1.0, 1.0, 1.0],   # “smiles”
]])

print(”x.shape:”, x.shape)

Output

x.shape: torch.Size([1, 3, 6])

The tensor x holds one mini batch of token embeddings. The shape 1, 3, 6 reads as batch size one, three tokens per sequence, and six features per token. The three rows correspond to “The”, “kid”, and “smiles”, and each row has six embedding values.

Listing 1.15: Projecting Input to Query, Key, and Value

b, num_tokens, d_in = x.shape

d_out = 6          # final output dimension we want per token
num_heads = 2
head_dim = d_out // num_heads   # 6 // 2 = 3

W_q = torch.nn.Parameter(torch.randn(d_in, d_out), requires_grad=False)
W_k = torch.nn.Parameter(torch.randn(d_in, d_out), requires_grad=False)
W_v = torch.nn.Parameter(torch.randn(d_in, d_out), requires_grad=False)

q = x @ W_q    # (1, 3, 6)
k = x @ W_k    # (1, 3, 6)
v = x @ W_v    # (1, 3, 6)

print(”q.shape:”, q.shape)
print(”k.shape:”, k.shape)
print(”v.shape:”, v.shape)

Output

q.shape: torch.Size([1, 3, 6])
k.shape: torch.Size([1, 3, 6])
v.shape: torch.Size([1, 3, 6])

Instead of having separate weight matrices for each head, we follow the weight splitting idea. We keep a single large query, key, and value matrix of shape 6, 6 Multiplying the 1, 3, 6 input by a 6, 6 weight gives new 1, 3, 6 tensors for q, k, and v. At this point there is no explicit notion of heads in the tensors; all six output features per token are packed into the last dimension.

Listing 1.16: Splitting Projections into Multiple Heads

# reshape from (b, num_tokens, d_out) to (b, num_tokens, num_heads, head_dim)
q = q.view(b, num_tokens, num_heads, head_dim)
k = k.view(b, num_tokens, num_heads, head_dim)
v = v.view(b, num_tokens, num_heads, head_dim)

print(”q after view:”, q.shape)
print(”k after view:”, k.shape)
print(”v after view:”, v.shape)

Output

q after view: torch.Size([1, 3, 2, 3])
k after view: torch.Size([1, 3, 2, 3])
v after view: torch.Size([1, 3, 2, 3])

The six features that came out of each projection are now interpreted as two heads with three features each. The view operation does not change any values; it only changes how we index them. The new shape 1, 3, 2, 3 can be read as batch size one, three tokens, two heads, three features per head. For a given token position, the last two dimensions now contain the representation for head one and head two.

Listing 1.17: Reordering Dimensions to Group by Head

# move the head dimension in front of the token dimension
# from (b, num_tokens, num_heads, head_dim)
# to   (b, num_heads, num_tokens, head_dim)
q = q.transpose(1, 2)
k = k.transpose(1, 2)
v = v.transpose(1, 2)

print(”q after transpose:”, q.shape)
print(”k after transpose:”, k.shape)
print(”v after transpose:”, v.shape)

Output

q after transpose: torch.Size([1, 2, 3, 3])
k after transpose: torch.Size([1, 2, 3, 3])
v after transpose: torch.Size([1, 2, 3, 3])

After splitting into heads, it is convenient to group all tokens that belong to the same head together. The transpose call swaps the token and head axes. The new shape 1, 2, 3, 3 reads as batch size one, two heads, three tokens, three features per head. If you isolate q[0, 0] you see the three query vectors for head one, while q[0, 1] contains the three query vectors for head two. The same interpretation holds for k and v. This layout allows a single tensor operation to compute attention for all heads in parallel.

Listing 1.18: Computing Per-Head Attention Scores and Context Vectors

import math

# scaled dot product attention for all heads at once
scores = q @ k.transpose(-1, -2)        # (b, num_heads, num_tokens, num_tokens)
print(”scores.shape:”, scores.shape)

scale = math.sqrt(head_dim)
weights = torch.softmax(scores / scale, dim=-1)
print(”weights.shape:”, weights.shape)
print(”weights[0, 0]:”)
print(weights[0, 0])

# context vectors inside each head
context = weights @ v                   # (b, num_heads, num_tokens, head_dim)
print(”context per head shape:”, context.shape)
print(”context[0, 0]:”)
print(context[0, 0])

Output

scores.shape: torch.Size([1, 2, 3, 3])
weights.shape: torch.Size([1, 2, 3, 3])
weights[0, 0]:
tensor([[0.59 , 0.24 , 0.17 ],
        [0.29 , 0.45 , 0.26 ],
        [0.22 , 0.31 , 0.47 ]])
context per head shape: torch.Size([1, 2, 3, 3])
context[0, 0]:
tensor([[ 0.564, -3.817,  2.064],
        [ 3.116,  9.936, 14.649],
        [-2.124, -4.104,  2.056]])

The tensor scores has shape 1, 2, 3, 3. For each batch and head, it contains a full three by three attention score matrix over the three tokens. Dividing by sqrt(head_dim) and applying softmax along the last axis converts these scores into attention weights, again separately for each head. The shape of weights matches that of scores.

Multiplying the weights by v produces the context tensor of shape 1, 2, 3, 3. For each head you now have three context vectors, one per token, each with three features. In the printed slice context[0,0] you can see the three vectors produced by the first head for “The”, “kid”, and “smiles”, which match the output structure.

Listing 1.19: Merging Heads into the Final Context Matrix

# move tokens back in front of heads and merge head and feature dimensions
context = context.transpose(1, 2).contiguous()  # (b, num_tokens, num_heads, head_dim)
context = context.view(b, num_tokens, num_heads * head_dim)

print(”final context.shape:”, context.shape)
print(”final context:”)
print(context)

Output

final context.shape: torch.Size([1, 3, 6])
final context:
tensor([[[ 0.564, -3.817,  2.064,  3.116,  9.936, 14.649],
         [-2.124, -4.104,  2.056,  3.098,  9.814, 14.479],
         [-2.120, -4.099,  2.053,  3.113,  9.915, 14.620]]])

To return to a single context matrix per token, we undo the earlier reordering and then collapse the head dimension and the per head feature dimension back into one. The transpose moves us from 1, 2, 3, 3 to 1, 3, 2, 3 grouping all heads for each token together. The final view then interprets the last two dimensions 2, 3 as a single dimension of size six.

The result is a 1, 3, 6 tensor. Each row is now a six dimensional context vector for one token, built by concatenating the three features from head one and the three features from head two. Compared to single head attention, nothing about the scoring or weighting changed. The difference is that we used reshaping and transposing to let two separate attention heads operate in parallel on smaller subspaces, then merged their outputs to recover the original six dimensional representation per token.

Concluding Multi Head Attention

Multi head attention extends single head attention by running several independent attention mechanisms in parallel, each with its own learned query, key, and value projections on a reduced head dimension. For an input embedding matrix, each head produces its own attention scores over all token pairs and then its own context matrix, so different heads can specialize on different types of relationships in the sequence. These context matrices are then concatenated along the feature dimension to recover the original output size, so each token representation combines multiple complementary views of the same input. While each head has fewer dimensions and therefore lower capacity than a single large attention module, the collection of diverse heads makes the overall representation more expressive, and a final linear projection can further mix and refine these combined features.

1.17 Layer Normalization

Figure 1.115: Layer normalization appears multiple times in the transformer block: before multi-head attention, before the feed-forward network, and often before the output layer.

In the transformer block, layer normalization appears several times. It is applied before the multi head attention sublayer, again before the feed forward network, and often once more outside the block before the final output layer. Because it is used so frequently, it is convenient to implement it as its own reusable module when we code the model.

Figure 1.116: The gradient flow problem in deep networks: without normalization, gradients can explode (very large activations) or vanish (very small activations), making training unstable.

To understand why layer normalization is so important, it helps to step back and look at a standard deep neural network with an input layer, several hidden layers, and an output layer. During the forward pass, activations flow from left to right, and during backpropagation the gradients flow in the opposite direction, from the output layer back through each hidden layer to the input.

Each layer has parameters and therefore receives gradients of the loss with respect to those parameters. The gradients at a given layer depend strongly on the outputs of that layer. If the layer outputs are very large in magnitude, the corresponding gradients tend to become very large as we chain them backward through the network. By the time they reach the earlier layers, they can explode to extremely large values. This is the exploding gradient problem and it leads to unstable updates and divergence during training.

The opposite can also happen. If the outputs of a layer are very small, the gradients that depend on them can quickly shrink as they propagate backward through many layers. Early layers then receive gradients that are almost zero and their parameters barely change. This is the vanishing gradient problem and it makes learning extremely slow or stops it altogether.

In both cases, very large or very small activations in intermediate layers create gradient magnitudes that are either too large or too small. Training then becomes unstable and inefficient. One way to stabilize the gradients is therefore to control the magnitude of the layer outputs themselves. This is exactly what normalization layers are designed to do.

Internal covariate shift

Figure 1.117: Internal covariate shift: as earlier layers update their weights during training, the distribution of activations fed into later layers keeps changing, making learning a moving target.

There is a second issue in deep networks known as internal covariate shift. During training, as earlier layers update their weights, the distribution of activations that they feed into later layers keeps changing. Imagine looking at the inputs to a particular hidden layer at the beginning of training. They may roughly follow one distribution. After a few training iterations, as weights update, the same layer may now see inputs with a different mean or variance, or a skewed shape. The layer is trying to learn a good mapping, but the distribution of its inputs keeps drifting, so the layer is constantly adapting to a moving target. This slows down convergence and makes optimization harder.

If we could keep the mean and variance of the inputs to each layer more stable across training iterations, learning would become easier. Normalization does this by rescaling activations so that their distribution is more consistent over time. This reduces internal covariate shift and helps the model converge faster.

The core idea of layer normalization

Layer normalization is a simple procedure applied to the outputs of a layer. Consider a single training example and focus on the vector of activations produced by some layer for that example.

Figure 1.118: Layer normalization in action: six activations with mean 0.6 and variance 0.07 are centered to mean 0 and rescaled to variance 1, producing standardized activations.

Imagine a single layer that produces six outputs for one training example

These are the values shown in the middle row of the figure. Layer normalization first computes the mean of these activations

Next it computes the variance

When each activation xi is normalized, it is first shifted so its mean becomes zero, and then it is rescaled so its variance becomes one. This process is known as centering and rescaling. The formula

means that for each output value, you subtract the average of all outputs (centering), and then divide by their standard deviation (rescaling). This transformation ensures that the set of normalized activations has an average of zero and a unit variance, making them easier for the next layer of a neural network to process.

If you perform this computation for all six values, you obtain approximately

These are the normalized outputs shown in the top row of the figure. By construction, the mean of these normalized activations is zero and their variance is one, which is why the left side of the illustration reports mean equal to 0.0 and variance equal to 1.00. This example demonstrates how layer normalization transforms a set of layer outputs with mean 0.6 and variance 0.07 into a standardized set of activations that are better behaved numerically, making gradient based training more stable.

In practice, layer normalization is usually followed by a learned scale and shift. After computing the normalized activations x hat i, the layer produces new activations

Here gamma and beta are trainable parameters with the same size as the activation vector. They let the model undo or modify the standardization whenever that helps performance. In other words, the network can learn whatever output distribution it prefers, while still enjoying the stabilizing effect of normalization during training.

The way mean and variance are computed is what distinguishes layer normalization from batch based methods. For layer normalization we normalize across the features of a single example, not across different examples in a mini batch. Each token representation or layer output vector is treated independently and normalized across its dimensions. This makes the procedure independent of batch size and very convenient for transformer models that see variable length sequences and perform autoregressive decoding one token at a time.

Inside a transformer block, layer normalization is applied to the token representations before they enter the multi head attention module. This keeps the scale of the inputs to attention under control and stabilizes its gradients. After attention and its residual shortcut, layer normalization is applied again before the feed forward network so that this network also sees well behaved activations. Many architectures further normalize the final outputs of each block, and often the outputs of the last block before the language modeling head that predicts the next token.

1.18 FeedForward Network

Figure 1.119: The feed-forward network sits between dropout and layer normalization inside each transformer block, processing each token representation independently

The second major component inside a transformer block, after the multi head attention, is the feed forward network. In the block diagram this appears as the Feed forward NN box, sitting between dropout and layer normalization. Conceptually it is an ordinary two layer neural network with an activation function in the middle, applied independently to every token representation. The key is that the same small network, with the same weights, is reused for all tokens and all examples in the batch.

Figure 1.120: The feed-forward network processes each token vector independently using shared weights. For an input tensor of shape (batch, tokens, 768), each 768- dimensional vector is processed in parallel.

2, 3, 768. The first entry is the batch size, the second is the number of tokens in the context window, and the third is the embedding dimension for each token. The important point is that the feed forward network is applied to every token vector of length 768 independently but it uses the same weights for all tokens and all examples in the batch.

Figure 1.121: Internal structure of the feed-forward network: a linear expansion from 768 to 3072 dimensions, followed by GELU activation, then linear contraction back to 768 dimensions.

The internal structure of this network is shown in the detailed diagrams. It consists of two linear layers with an activation function in between. The first linear layer performs an expansion. It takes each 768 dimensional input vector and projects it into a much larger space with 4 × 768 = 3072 hidden units. In matrix form this is a multiplication by a 768 by 3072 weight matrix plus a bias term. Intuitively, this expansion gives the model more capacity to construct rich intermediate features from each token representation before compressing them again. Because every input dimension interacts with every hidden unit, this single layer already introduces dense mixing among the 768 input features.

Figure 1.122: Comparison of ReLU and GELU activation functions. GELU is smooth everywhere, including at zero, and preserves small negative activations rather than collapsing them to exactly zero

After the expansion, the output is passed through a nonlinearity. Early transformer implementations used the ReLU activation function. ReLU simply returns x when x is positive and returns zero for negative x. Graphically, this is a straight line through the origin for positive inputs and a flat line at zero for negative inputs. ReLU is easy to implement and works well in many convolutional and fully connected networks, but it has two important drawbacks in this setting. First, every negative input is collapsed to exactly zero, so any information stored in the magnitude of negative activations is lost. If many units become negative, a large part of the network can effectively stop learning, a problem often described informally as dead neurons. Second, the ReLU curve has a sharp corner at zero and is not differentiable there. In practice we can still compute subgradients and train, but the function is not smooth.

Because of these issues, transformer language models have largely switched to the GELU activation. The GELU curve, shown next to ReLU in the figures, is a smooth S shaped function. For large positive inputs it behaves similarly to ReLU and returns values close to the identity. For large negative inputs it sends activations toward zero, so very negative units are still turned off. The important difference appears around zero. Instead of cutting everything below zero to exactly zero, GELU tapers smoothly and maps small negative inputs to small negative outputs. This has two consequences. First, the function is differentiable everywhere, including at zero, which makes optimization smoother. Second, the network does not discard all information carried by small negative activations. In the region near zero the model can still use their sign and magnitude to encode subtle distinctions. Together with layer normalization, which keeps activations in a moderate range and prevents very large negative or positive values, this leads to more stable training and slightly better performance in practice.

The output tensor of the feed forward network has exactly the same shape as its input, namely batch size, context length, embedding dimension. This design choice is deliberate. Because the main hidden dimension remains constant, we can add a residual connection around the feed forward sublayer and stack an arbitrary number of transformer blocks on top of each other without reshaping tensors. It becomes straightforward to plug in more blocks, remove blocks, or reuse the same block structure in very deep models, since every block expects and returns vectors of size 768 in this example.

It is also helpful to connect this back to the question of how many tokens the network sees at once. While the feed forward network operates on each token vector independently, each of those vectors already encodes information about the entire context window thanks to the preceding multi head attention. In our example with context length 3, the tensor has shape 2, 3, 768, and each of the three vectors for a given sequence summarizes that position in the context of the other positions. The feed forward network then applies a rich non linear transformation to each of these contextual vectors. When generating text autoregressively, the model still predicts one next token at a time, but this prediction is based on the full context representation, which has been refined by both attention and the feed forward expansion, activation and contraction.

In summary, the position wise feed forward network in a transformer block is a powerful per token multilayer perceptron. It expands the embedding dimension to a higher dimensional space, applies a smooth and information preserving GELU nonlinearity, and contracts the representation back to the original size. This structure provides most of the depth and nonlinearity in the model, while attention handles interaction across positions. Together they give transformers the capacity to learn complex patterns in sequences of tokens.

1.19 Shortcut connections

Layer normalization helps stabilize the scale of activations, but on its own it is not enough to reliably train very deep transformer stacks. In the previous section we saw that each block already contains a powerful feed forward network that expands the embedding dimension, applies a GELU nonlinearity, and then contracts it again. If we simply stacked many of these attention plus feed forward blocks, gradients flowing backward through all of those nonlinear layers would quickly become very small. To keep such deep transformers trainable, we rely on another key idea shortcut connections, also called residual connections.