From CNNs to LLMs: How Vision Transformers Changed Everything
~35 min read
17 diagrams
LaTeX + PyTorch code
The trajectory from convolutional neural networks to large language models is not a random walk through research papers. It is a story about one idea — attention — escaping its original domain and rewriting the rules of deep learning. If you build ML systems in production, this arc directly shapes the infrastructure you need, the hardware you provision, and the trade-offs you make every day.
I have worked on both sides of this transition: building CNN-based anomaly detection and object tracking systems, and now deploying LLM inference infrastructure with vLLM. This post traces the architectural path between those worlds.
Table of Contents
Part I — The CNN Era
- How Convolution Works — kernel, stride, padding, receptive field
- Hierarchical Features — edges → textures → parts → objects
- CNN Architecture — full pipeline with dimensions
- The Receptive Field Problem — CNN vs Transformer comparison
- How Many Layers Until Ear Meets Paw? — why CNNs need depth for global context
Part II — ViT: From Images to Sequences
- Patch Embedding — flatten + linear projection (e = Wx + b)
- The CNN ↔ ViT Equivalence — Conv2d ≡ Linear (same math!)
- Positional Embeddings — why we need them, add vs concatenate
- ViT Architecture — full pipeline diagram
- ViT vs CNN — comparison table
Part III — Self-Attention: How Patches Talk
- The Goal (Without Q, K, V) — ear patch gathering information
- Q, K, V Step by Step — projections, scores, scaling
- Why Multiple Heads? — parallel attention patterns
- Concat + Output Projection — combining head outputs
- Single Head Alternative — the associativity trick
- Complete Transformer Block — residuals + MLP
- Classification — CLS token vs average pooling
- Logits, Softmax, Loss — the final step
Part IV — From ViT to LLMs
- Same Architecture, Different Domain — side-by-side comparison
- Token Embeddings in LLMs — patches → words
- Causal Masking — why LLMs can't look ahead
- MLP Variations — GPT-3 vs Qwen3/LLaMA
- Next-Token Prediction — autoregressive generation
- The Language Model Head — Linear → Softmax → sample
- What This Means for Infrastructure — CNN vs LLM serving
The CNN Era: Locality as Inductive Bias
Convolutional neural networks dominated computer vision for nearly a decade. The core idea was elegant: spatial locality matters. A pixel's meaning depends on its neighbors, not on pixels across the image. The convolutional kernel encodes this assumption directly — it slides across the image, processing small local patches.
How Convolution Works
The kernel is a small matrix of learnable parameters (e.g. 3×3 = 9 weights). It slides across the input, computing a weighted sum at each position to produce one value in the output feature map.
Each position: the kernel multiplies its 9 weights with the 9 input values under it, sums them up, and produces one output value. The kernel weights are learned during training — the network discovers which patterns (edges, textures, shapes) matter for the task.
Convolution output size:
where = input size, = kernel size, = padding, = stride.
Example: Input , kernel , padding 1, stride 1:
Max pooling: pool with stride 2 halves dimensions:
Hierarchical Features: Wider → Deeper
As you go deeper in the network, feature maps get spatially smaller (less width/height) but deeper (more channels). Each level captures increasingly abstract patterns:
Each pooling layer halves the spatial dimensions while each convolution layer can increase the number of channels. The network trades spatial resolution for semantic depth — from raw pixels to high-level concepts.
CNN Architecture
This inductive bias was a gift:
- Parameter efficiency: A 3×3 kernel has 9 parameters regardless of image size. A fully connected layer over a 224×224 image would need ~150 million parameters for a single layer.
- Translation equivariance: A cat in the top-left corner activates the same filters as a cat in the bottom-right. The network doesn't need to re-learn features at every position.
- Hierarchical features: Stacked convolutional layers build increasingly abstract representations — edges → textures → parts → objects.
Architectures like ResNet, VGG, and EfficientNet refined this recipe. Residual connections solved gradient vanishing. Batch normalization stabilized training. Depthwise separable convolutions reduced compute. The field was optimizing within a well-understood paradigm.
But CNNs had a ceiling.
The local receptive field — the very thing that made CNNs efficient — also limited them. To understand relationships between distant parts of an image, you need the receptive field to grow large enough to cover both regions. With 3×3 kernels, that means stacking many layers:
The Receptive Field Problem
How Many Layers Until Ear Meets Paw?
Each 3×3 conv layer expands the receptive field by 1 pixel in each direction. Two features far apart need many stacked layers before their receptive fields overlap:
This is the fundamental limitation: depth is the price of global context in CNNs. A Transformer skips all of this — self-attention connects everything in one step.
A CNN needs 50+ layers of 3×3 convolutions just to let distant pixels "see" each other. A Transformer connects everything in one layer through self-attention. This is the fundamental trade-off: CNNs are efficient but local; Transformers are expensive but global.
ViT: The Bridge Between Worlds
In 2020, Google Research published "An Image Is Worth 16×16 Words" — the Vision Transformer (ViT). The idea was almost provocatively simple: treat image patches as tokens and feed them into a standard Transformer.
Patch Embedding: How Images Become Tokens
The key step is linear projection of flattened patches. Each 16×16 patch of the image is flattened into a vector and projected into the Transformer's embedding dimension through a learnable linear layer:
where is the flattened patch (16×16×3 = 768), and are learnable projection weights and bias, is a learnable position embedding, and is the embedding dimension: (Base), (Large), (Huge) — a hyperparameter, not derived from input size.
The CNN ↔ ViT Equivalence
Here's something elegant: a convolution with kernel size 16×16, stride 16, and D output filters is mathematically identical to ViT's linear projection. The two formulations produce exactly the same result.
For a single 16×16 RGB patch, one convolution filter computes:
Where:
- — pixel value at row , column , channel
- — corresponding learned weight (768 weights total)
- — one learned bias
You multiply matching values, then add all 768 products together — the same as flattening and doing a dot product. Expanding across the 3 RGB channels:
For the whole image with stride 16 and output filter :
The and move the filter 16 pixels at a time (stride 16). Every patch uses the same weights — this IS the shared linear projection.
The equivalence:
This is why ViT is not a complete departure from CNNs — the patch embedding is a convolution. The difference is that ViT uses this single convolution only for tokenization, then switches entirely to attention for all feature learning.
Positional Embeddings: Teaching Position to a Sequence Model
In a CNN, pixels have spatial relationships because convolutions operate on a 2D grid — the kernel physically slides across spatial positions. But after flattening image patches into a sequence of vectors, the Transformer doesn't inherently know which patch came from which image location. Patch 1 (top-left) and patch 196 (bottom-right) are just vectors in a list — the model has no notion of "above," "below," "left," or "right."
Why Do We Need Positional Embeddings?
Imagine two images containing exactly the same patches, but arranged in different positions:
Without positional information, a standard Transformer has no built-in way to distinguish their spatial arrangement — self-attention is permutation-invariant, meaning it produces the same output regardless of input order. The model would classify both images identically, even though they represent completely different scenes.
We solve this by adding (not concatenating) a learned positional vector to each patch embedding:
Where:
- = patch embedding (from linear projection)
- = positional embedding (learnable)
- = final input vector — same dimension as
We add the vectors, we do not concatenate them. The embedding dimension remains (e.g. 768). Positional embeddings are like learnable location labels added to your patch features — they don't change the number of patches or the feature dimension.
Why learned, not fixed? The original Transformer (for text) used fixed sinusoidal positional encodings. ViT uses learned positional embeddings — the model discovers the optimal spatial encoding during training. Interestingly, after training, the learned position embeddings show a 2D spatial pattern: nearby patches have similar embeddings, and the model recovers a grid-like structure even though it was never told the patches came from a 2D image.
The full ViT input can be written as:
where patches, each , and is a special learnable class token whose output representation is used for classification.
ViT Architecture
ViT vs CNN: The Trade-Off
| | CNN | ViT | |---|---|---| | Inductive bias | Locality (hardcoded) | None (learned from data) | | Receptive field | Local → grows with depth | Global from layer 1 | | Data efficiency | Good (works on small datasets) | Poor (needs millions of images) | | Scaling | Diminishing returns | Near-linear improvement | | Parameters (large) | ~25M (ResNet-50) | ~307M (ViT-Large) | | Compute pattern | Compute-bound, regular memory | Memory-bandwidth bound |
ViT proved something profound: you don't need domain-specific architectural biases if you have enough data. With sufficient training data, the Transformer learned spatial relationships purely from data rather than having them hardwired.
This had immediate implications:
- Unified architecture: The same Transformer block processes both images (as patch sequences) and text (as token sequences). One architecture, multiple modalities.
- Scaling behavior: ViT scaled better than CNNs with increasing compute and data. Larger models consistently improved.
- Transfer learning: Pre-trained ViT representations transferred more effectively than CNN features.
Self-Attention: How Patches Talk to Each Other
Now that we understand how ViT splits an image into 196 patch embeddings and feeds them as a sequence into a Transformer, let's look at the core mechanism inside the Transformer Encoder: self-attention. This is how each patch gathers information from every other patch to build a richer, context-aware representation.
The Transformer architecture (Vaswani et al., 2017 — "Attention Is All You Need") originally proposed self-attention for NLP, where it let every token in a sentence attend to every other token in parallel. ViT applies the exact same mechanism — but instead of words, the tokens are our 196 patch embeddings.
First, Understand the Goal (Without Q, K, V)
Before diving into the mechanics, let's understand what self-attention is trying to do. Imagine we have only 4 image patches:
The new representation of the ear patch is a weighted combination of information from all patches:
After self-attention, the ear patch doesn't just contain "ear pixels" — it contains information about the eye, body, and their spatial relationships. It now understands context. The sky contributes almost nothing because it's semantically irrelevant to the ear.
This is the fundamental idea: every patch updates its representation by gathering relevant information from all other patches. The weights (15%, 35%, 45%, 5%) determine how much each patch contributes.
The question is: how does the model learn these weights? That's where Q, K, V come in.
Self-Attention: Step by Step (the Q, K, V mechanism)
After patch projection and positional encoding, we have:
That's 196 vectors (one per patch), each with 768 features. Now self-attention lets every patch look at every other patch to decide what's relevant.
Step 1: Project X into three different representations
Each patch vector gets projected through three separate learned weight matrices into Query, Key, and Value:
Think of it as asking each patch three questions:
- Query (Q): "What am I looking for?" — what information does this patch need from others?
- Key (K): "What do I contain?" — what information can this patch offer to others?
- Value (V): "What information do I carry?" — the actual content to share
Step 2: Compute attention scores
How much should patch attend to patch ? Take the dot product of patch 's query with patch 's key:
What about the diagonal? Notice that when (a patch attending to itself), the same appears twice:
This is not simply the square of . The matrix sits between the two copies and can mix features with positive or negative weights. For a tiny example, suppose the projections produce:
Even though both came from the same patch, the score is zero! If another patch has , then . So using the same input twice does not guarantee a larger self-score — the two transformations ( and ) are different learned matrices.
Why divide by ? Recall the dot product formula:
where is the vector magnitude. As grows, dot products tend to grow proportionally (more terms summed), pushing softmax into saturated regions with near-zero gradients. Dividing by normalizes this:
What is ? It is the number of features in each key/query vector within one attention head — not the vector's magnitude:
- With 1 head: , so we divide by
- With 12 heads: , so we divide by
Step 3: Convert scores to weights with softmax
The raw scores can be any real number — we need to convert them into weights that sum to 1 for each query patch. The division by helps keep the scores at a manageable scale. Then apply softmax:
So for each query patch , its attention weights across all 196 patches sum to exactly 1:
This is a probability distribution — the ear patch "spends" 100% of its attention budget across all other patches. Higher scores get exponentially more weight (that's what does), making the distribution peaky: a few patches get most of the attention, the rest get almost none.
This gives us a attention matrix — every patch's relationship to every other patch:
Step 4: Multiply by V — gather the actual information
Attention weights tell us where to gather information, but contains the information itself. Without , we only have relationships between patches — not their content.
For the ear patch, its attention row might be:
Those four numbers only say how much to take from each patch. They don't contain the eye's shape or the body's appearance. Multiplying by gathers that information:
The result is a new feature vector for the ear patch — it no longer contains just "ear pixels." It now holds a weighted mixture of eye features, body features, and spatial context. The ear patch now understands that it's part of a cat face.
For the full input:
Keeping only the weights () gives relationships between patches. Multiplying by turns those relationships into updated patch features — each patch's representation is enriched with context from the entire image.
Putting it all together:
Why Do We Need Multiple Heads?
Imagine a ViT processing a photograph of a dog. Each head has different trainable weights (), allowing the model to learn different ways of comparing and aggregating patches. A single head produces one attention distribution per patch. Twelve heads produce twelve potentially different distributions.
Different attention heads could learn different relationships:
What Happens After the 12 Heads?
Each head produces its own output . We concatenate them:
Then we apply another fully connected layer (output projection):
Why this extra FC layer? Because concatenation only places the head outputs next to each other — head 1's 64 features, then head 2's 64 features, etc. The output projection allows the model to learn combinations of information from all 12 heads. Without it, each head's features stay isolated.
This is analogous to a Conv2d in CNNs — it mixes channels from different convolution filters into richer representations.
The output has the same shape as the input () — this is critical because it means Transformer blocks can be stacked. The output of one block feeds directly into the next.
Single Head Alternative: A × X Then FC
With a single head, there's an interesting equivalence. Since matrix multiplication is associative:
These give the same result:
- Project with , then apply attention
- Apply attention to , then use a fully connected layer with the same
So a single-head architecture can be written as: first mix patches (attention), then project features (linear layer):
class SingleHeadAttention(nn.Module):
def __init__(self, dim=768):
super().__init__()
self.dim = dim
self.q = nn.Linear(dim, dim)
self.k = nn.Linear(dim, dim)
self.v = nn.Linear(dim, dim)
self.out = nn.Linear(dim, dim)
def forward(self, x):
# x: [B, N, 768]
Q = self.q(x)
K = self.k(x)
# scores: [B, N, N]
scores = Q @ K.transpose(-2, -1)
scores = scores / (self.dim ** 0.5)
weights = F.softmax(scores, dim=-1)
# First mix patches, then project features
mixed = weights @ x # [B, N, 768]
attended = self.v(mixed) # [B, N, 768]
return self.out(attended) # [B, N, 768]This works because every row of the attention weights sums to 1 (from softmax), which also makes the equivalence hold for bias terms in nn.Linear — apart from floating-point rounding.
Important: This fully connected layer () is a linear projection — not the full Transformer MLP, which includes a nonlinear activation (GELU). The distinction matters.
One Head vs Multi-Head: Full Comparison
The winner: Multi-head attention is used in every modern Transformer (ViT, GPT, LLaMA) because the diversity of 12 parallel perspectives consistently outperforms one large perspective — even though each individual head sees fewer features. The sweet spot is typically with heads.
Complexity: Self-attention has complexity — the attention matrix means every patch attends to every other patch. For 196 patches it's fine. For 1024 patches (smaller patch size), it starts to hurt. This quadratic cost is why techniques like FlashAttention and sparse attention matter at scale.
The Complete Transformer Block: Residuals + MLP
Once attention produces an output of shape , the next steps are the same whether you use one or multiple heads.
Step 1: Add the residual connection
Add the attention output back to the block's original input:
Both have the same shape, so this is element-by-element addition. It lets the block keep the original patch information while adding information gathered through attention.
Why residual? Without it, attention can dilute patch-specific information. Imagine the ear patch attends mostly to the body and eye:
Without the residual, the MLP receives only this mixture — it has no direct access to the original ear embedding. With the residual:
The original ear features contribute directly, alongside the gathered context. If the attention update is not useful, it can learn , and the input passes through unchanged:
This also makes deep networks trainable — gradients have a direct path backward through the addition.
Why add, not concatenate?
Concatenation would double the dimension at every block — after 12 blocks you'd have . Addition keeps the same size and provides a direct gradient path without a learned projection.
Step 2: Apply the MLP with another residual
A typical ViT MLP is:
mlp = nn.Sequential(
nn.Linear(768, 3072), # expand 4×
nn.GELU(), # nonlinear activation
nn.Linear(3072, 768), # project back
)Why the activation matters: Without GELU, two linear layers collapse into one (). The nonlinearity lets the MLP learn more complex transformations that a single FC layer cannot.
The two components have different jobs:
- Attention: mixes information between patches (inter-patch)
- MLP: transforms features within each patch (intra-patch), applied identically to every patch
The complete block outputs — same shape as input. This is the pre-normalization layout used in modern ViTs. Stack 12 of these blocks, and the final output goes to the classification head.
Key insight: If both transformations () are zero, the data flows directly from initial patch embeddings to the final layer through the residual connections — the network has a direct path to the original information at every depth.
Notice: the MLP has its own separate residual connection, independent from attention:
# Attention block with residual
x = x + self.attn(self.norm1(x)) # [B, N, 768]
# MLP block with separate residual
self.mlp = nn.Sequential(
nn.Linear(768, 3072),
nn.GELU(),
nn.Linear(3072, 768),
)
x = x + self.mlp(self.norm2(x)) # [B, N, 768]Both residual additions have shape — the block preserves dimensions throughout. Each residual independently lets information skip its transformation. The full block has two skip paths: one around attention, one around MLP. This means the original patch embedding can survive through both — the network can selectively use attention, MLP, both, or neither.
From Patches to Classification
After all 12 Transformer blocks, we have enriched patch representations. But we need a single vector representing the whole image for classification. First, apply a final layer norm:
x = self.final_norm(x) # [1, 197, 768]Now there are two approaches:
Option A: Global Average Pooling (without [CLS] token)
Average across all 196 patches, separately for each of the 768 features:
# x: [1, 196, 768]
image_vector = x.mean(dim=1) # [1, 768]Feature 1 of the image vector is the average of feature 1 across all patches. This operation is called global average pooling — you get one 768-dim vector representing the whole image, not individual patches.
Option B: [CLS] Token (the ViT approach)
Recall we prepended a special learnable token at the start. After 12 blocks of self-attention, this token has attended to every patch and aggregated image-wide information:
# x: [1, 197, 768] (196 patches + 1 CLS)
image_vector = x[:, 0, :] # [1, 768] — just take the first tokenThe CLS token's final representation is the image summary — it learned to aggregate during training.
Both approaches produce a single vector. Then classification is a simple linear layer:
self.head = nn.Linear(768, num_classes) # e.g., 1000 for ImageNet
logits = self.head(image_vector) # [1, 1000]Deep Dive: The [CLS] Token
The CLS token is a single learnable vector — an nn.Parameter of shape — that we prepend to the patch sequence before the first Transformer block:
self.cls_token = nn.Parameter(torch.zeros(1, 1, 768))
# In forward:
cls = self.cls_token.expand(x.shape[0], -1, -1) # [B, 1, 768]
x = torch.cat([cls, x], dim=1) # [B, 197, 768]nn.Parameter makes it trainable. It starts the same for every image, but becomes image-specific as it passes through attention — each image's patches fill it with different information.
After all blocks and final normalization:
The image representation is simply:
We drop all 196 patch vectors and keep only the CLS position:
# After all Transformer blocks + final norm
# x: [B, 197, 768]
cls_vector = x[:, 0, :] # [B, 768]
logits = self.classifier(cls_vector) # [B, num_classes]The 196 patch vectors are not passed to the classifier — but they helped build the CLS representation through 12 layers of attention. The CLS vector at the end is not the original nn.Parameter — it's a completely transformed representation containing information about the current image.
Why CLS instead of average pooling?
CLS lets the model learn how to gather information for classification. Average pooling gives every final patch vector equal weight.
For example, if a small animal occupies only a few patches:
- Average pooling: averages all 196 vectors equally, including background patches
- CLS: can learn through attention to gather more information from relevant patches and less from background
Important nuance: The patch vectors being averaged already contain contextual information from 12 attention layers — so average pooling can also work well. Neither method is always better. CLS provides a dedicated learned representation; average pooling is simpler and needs no extra token.
Final Classification: Logits, Softmax, and Loss
The classification head is just a fully connected layer:
self.classifier = nn.Linear(768, num_classes) # e.g., 1000 for ImageNet
logits = self.classifier(image_vector) # [B, num_classes]These logits can be positive or negative — they are raw scores, not probabilities.
Softmax converts logits into probabilities that sum to 1:
where is the logit for class , and is the total number of classes.
Argmax selects the class with the highest score:
Cross-entropy loss measures how wrong the predictions are during training:
where is the correct class. This simplifies because the true label is one-hot — only the correct class contributes to the loss.
criterion = nn.CrossEntropyLoss()
# During training: pass raw logits (CrossEntropyLoss applies softmax internally)
loss = criterion(logits, labels)
# To get probabilities:
probabilities = logits.softmax(dim=-1)
# To choose the predicted class:
predicted_class = logits.argmax(dim=-1)Why not apply softmax before the loss? PyTorch's
nn.CrossEntropyLosscombines softmax + log + negative likelihood in one numerically stable operation. Applying softmax separately can cause floating-point issues with very large or small logits.
This completes the ViT pipeline: image → patches → embed → position → 12× (attention + MLP) → CLS/avg → linear → logits → prediction.
From ViT to LLMs: The Same Architecture, Different Domain
GPT-style language models also stack attention and MLP blocks, with normalization and residual connections. You now understand much of their foundation.
The main differences are:
| | Vision Transformer (classification) | GPT-style language model | |---|---|---| | Input tokens | Image patches | Text tokens | | Token embeddings | Represent patches (linear projection) | Represent text tokens (lookup table) | | Attention | Every patch attends to every patch | Each token attends only to itself and earlier tokens | | Final output | One image-class prediction | A next-token prediction at each position |
The core blocks are identical — same attention, same MLP, same residuals:
The core Transformer block is identical. What changes is what goes in and what comes out:
From Patches to Words: Token Embeddings in LLMs
In ViT, we tokenized images by splitting into patches and linearly projecting them. In LLMs, tokenization is simpler — each word (or subword) maps to a learned embedding vector via a lookup table.
For example, with Qwen3-0.6B (embedding dimension = 1024):
Each row contains 1,024 learned features for one token. Same idea as ViT patch embeddings — just a different input domain.
# ViT: patches → linear projection
embeddings = linear(flatten(patches)) # [196, 768]
# LLM: token IDs → embedding lookup
embeddings = embedding_table[token_ids] # [3, 1024]Both produce a matrix of shape — and from here, the Transformer blocks work identically.
The Key Difference: Causal Masking
In ViT, every patch can attend to every other patch — the attention is bidirectional. The ear can look at the tail and the tail can look at the ear.
In a language model, we don't want the next word to contribute to predicting itself — only previous words should. This is the causal mask.
First, compute scores normally:
Then mask the scores for future tokens. For "the cat sat":
Rows are queries; columns are keys. We keep the diagonal and everything below it (past + current). Everything above (future) is set to .
Why , not zero? Because softmax turns into exactly zero attention weight. A score of zero would still produce , which means the token would still receive some attention.
The implementation is straightforward — Q, K, V are computed normally, but we apply a mask before softmax:
scores = Q @ K.transpose(-2, -1)
scores = scores / (Q.shape[-1] ** 0.5)
# Create upper-triangular mask (future positions)
N = Q.shape[-2]
future = torch.triu(
torch.ones(N, N, device=Q.device, dtype=torch.bool),
diagonal=1,
)
scores = scores.masked_fill(future, float("-inf"))
weights = scores.softmax(dim=-1)
output = weights @ VSo Q, K, and V are computed normally — each token gathers information only from itself and previous tokens. The architecture is the same; only the mask changes.
MLP Variations Across Models
The MLP block also evolves across model generations. The concept stays the same — expand features, apply nonlinearity, project back — but the details differ:
GPT-3 style (classic):
self.mlp = nn.Sequential(
nn.Linear(dim, 4 * dim), # expand 4×
nn.GELU(approximate="tanh"), # nonlinearity
nn.Linear(4 * dim, dim), # project back
)Qwen3 / LLaMA style (gated MLP with SwiGLU):
# Two parallel projections + gating
gate = F.silu(self.gate_proj(x)) # [B, N, 4*dim]
up = self.up_proj(x) # [B, N, 4*dim]
x = self.down_proj(gate * up) # [B, N, dim]The gated variant uses two parallel linear projections: one goes through SiLU activation (the "gate"), then multiplies element-wise with the other ("up"), and the result is projected back down. This gives the model more expressive power at similar compute cost.
| | GPT-3 | Qwen3 / LLaMA | |---|---|---| | MLP | Linear → GELU → Linear | Gate×Up (SwiGLU) → Linear | | Activation | GELU | SiLU (Swish) | | Normalization | LayerNorm | RMSNorm (faster, no mean subtraction) | | Position encoding | Learned absolute | RoPE (rotary, extrapolates to longer sequences) | | Residuals | Same | Same | | Attention + MLP structure | Same | Same |
The core pattern — attention for inter-token mixing, MLP for per-token transformation, residuals for information flow — remains unchanged from ViT through GPT-3 to the latest open-source LLMs.
Dimensions Through the LLM
The MLP output has the same shape as its input. For our Qwen3-0.6B example with 3 tokens:
Internally, the feature dimension expands to 3,072, then returns to 1,024. The batch size and token count stay unchanged. That lets us add the residual:
x = x + self.mlp(self.norm2(x))
# [1, 3, 1024] + [1, 3, 1024] → [1, 3, 1024]We still have one updated vector per token.
From Representations to Next-Token Prediction
After attention and MLP, we pass through all Transformer blocks, then predict the next token:
# "the cat sat": assume 3 tokens
x = self.embedding(token_ids) # [1, 3, 1024]
for block in self.blocks:
x = block(x) # [1, 3, 1024]
x = self.final_norm(x) # [1, 3, 1024]Then we take the last token's representation. Because of causal attention, this vector contains information from "the," "cat," and "sat" — but nothing from the future:
last = x[:, -1, :] # [1, 1024]A final linear layer — the language model head — produces one score for every entry in the vocabulary:
logits = self.lm_head(last) # [1, 151936] (Qwen3 vocab size)We turn these scores into probabilities and select a token:
probs = logits.softmax(dim=-1)
next_token = torch.multinomial(probs, 1) # sample from distributionThen append the new token and repeat:
This is autoregressive generation — the model produces one token at a time, each conditioned on everything before it. This is why LLM inference is fundamentally different from image classification: it requires sequential decode steps, KV-cache management, and is memory-bandwidth bound rather than compute-bound.
The Language Model Head
The LM head is a fully connected linear layer — no activation function after it:
self.lm_head = nn.Linear(1024, 151936, bias=False)
logits = self.lm_head(last) # [1, 151936]
probabilities = logits.softmax(-1) # [1, 151936]No ReLU or GELU after this linear layer — the raw logits go directly to softmax (for inference) or cross-entropy loss (for training):
# Training: logits → CrossEntropyLoss (handles softmax internally)
loss = F.cross_entropy(logits, target_token_id)
# Inference: logits → softmax → sample
probs = logits.softmax(dim=-1)
next_token = torch.multinomial(probs, 1)This is the same pattern as ViT classification — the difference is the vocabulary size. ViT classifies into ~1,000 ImageNet classes. An LLM "classifies" into ~150,000 vocabulary tokens. But the final step is identical: Linear → Softmax → pick the winner.
LLM Architecture (Decoder-Only Transformer)
The Key Difference: Training Objective
- CNN: Supervised classification — predict the image class from labeled data
- ViT: Supervised classification — same objective, but on patch sequences
- LLM: Self-supervised next-token prediction — predict the next word given all previous words
Next-token prediction: the text itself is the supervision signal. No labels needed.
Next-token prediction turned out to be a remarkably powerful objective. By predicting the next token across trillions of tokens of text, the model implicitly learns syntax, semantics, reasoning patterns, world knowledge, and even basic logic.
What This Means for Infrastructure
If you are building ML infrastructure — which is what I do — the CNN → ViT → LLM trajectory has concrete consequences:
Compute Patterns Changed
| | CNN Inference | LLM Inference | |---|---|---| | Pattern | Single forward pass | Autoregressive (token by token) | | Bottleneck | Compute (FLOPs) | Memory bandwidth (KV-cache) | | Latency | Fixed per image | Grows with output length | | Memory | O(1) per request | O(seq_len) KV-cache per request | | Batching | Simple (fixed shape) | Complex (variable length) | | Hardware | CPU or single GPU | Multi-GPU, tensor parallelism | | Serving | Flask + PyTorch | vLLM + PagedAttention |
LLM inference has two phases:
Prefill phase (parallel):
Process entire prompt at once
Similar to a CNN forward pass
Compute-bound
"The cat sat on the" → [hidden states]
Decode phase (sequential):
Generate one token at a time
Each token depends on all previous
Memory-bandwidth bound
[hidden states] → "mat" → "." → [EOS]
↑
KV-cache stores all
previous attention statesThis is where vLLM's PagedAttention becomes critical — it eliminates memory fragmentation in the KV-cache, allowing more concurrent requests per GPU.
The Transformer did not just replace the CNN. It replaced the assumption that different domains need different architectures — and that single idea reshaped all of deep learning.