Building a Vision Language Model From Scratch: What I Learned Reimplementing PaliGemma in PyTorch

A deep dive into the design decisions, engineering trade-offs, and unexpected lessons from building a production-grade VLM from the ground up — without relying on any existing model implementations.

Introduction

Most engineers who work with Vision Language Models interact with them through an API or a pretrained checkpoint. You pass in an image and a prompt, and text comes back. What happens in between is abstracted away behind several layers of framework code, pretrained weights, and documented-but-rarely-read architecture papers.

I wanted to understand what was actually happening. Not at the level of “this model uses a transformer decoder” — but at the level of individual tensor shapes, design choices made in specific papers, and the engineering constraints that motivated each architectural decision. So I built PaliGemma — Google’s multimodal Vision Language Model — from scratch in PyTorch, without borrowing from any existing implementation.

This post is a record of what I built, what I learned, and what surprised me.

The result is a single-file, fully-annotated implementation that covers every component: the SigLIP contrastive vision encoder, the Gemma language model decoder, the KV-Cache, Rotary Positional Encoding, Grouped Query Attention, and the full autoregressive inference pipeline. The entire model runs end-to-end from a raw JPEG and a text prompt to generated text — loading weights directly from HuggingFace safetensor files, with no conversion required.

Why PaliGemma?

PaliGemma is a compelling choice for a from-scratch implementation for several reasons.

First, it is small enough to be tractable. The base variant uses a 400M-parameter SigLIP vision encoder and a 2B-parameter Gemma language model — large enough to be instructive about real design trade-offs, but small enough to run inference on a consumer GPU or even a MacBook with Apple Silicon.

Second, it is architecturally diverse. Building PaliGemma requires implementing two distinct Transformer architectures — a bidirectional encoder for vision and a causal decoder for language — as well as the interface between them. Each brings its own set of design decisions: different positional encoding schemes, different normalisation strategies, different attention patterns.

Third, and most importantly, PaliGemma represents the current state of the art in VLM design. Almost every architectural decision it makes — SigLIP contrastive pre-training, Grouped Query Attention, RMSNorm, Rotary Positional Encoding — was introduced to solve a specific, well-motivated problem. Building it from scratch forces you to understand not just what each component does but why it exists.

The Architecture

PaliGemma is composed of three parts that pass information from left to right: a vision encoder that reads the image, a projection layer that translates the image’s embedding space into the language model’s embedding space, and a language model decoder that generates text conditioned on both the image and the prompt.

Figure 1. PaliGemma end-to-end architecture. The image is processed by the SigLIP 400M vision encoder, projected by a single linear layer, and concatenated with text tokens before being fed into the Gemma 2B transformer decoder, which generates the output text autoregressively.

What makes this interesting is not the diagram — it is the details hiding inside each box. I’ll work through each one.

The Vision Encoder: Why Contrastive Pre-Training?

The first question to answer is why the vision encoder needs to be contrastively trained rather than just trained as a standard image classifier or autoencoder.

The answer comes down to the downstream use case. In a Vision Language Model, image embeddings and text embeddings are used together. The language model will be asked to attend to image patch embeddings alongside text token embeddings. If the image encoder was trained purely on visual tasks — classifying images into 1,000 ImageNet categories, for example — its embedding space has no natural relationship to the space that text tokens live in. The representations may be informationally rich, but they are not aligned.

Contrastive pre-training solves this by training the vision encoder and a text encoder simultaneously, on a shared objective: image and text descriptions of the same thing should be close in embedding space; images and unrelated descriptions should be far apart. After billions of (image, alt-text) pairs from the internet — a dataset that is essentially free, since every HTML image tag has an alt attribute — the encoder learns a representation that already speaks the same geometric language as text.

From CLIP to SigLIP

The seminal contrastive model is CLIP (OpenAI, 2021). CLIP uses a softmax cross-entropy loss computed over the full N×N similarity matrix — all dot products between N image embeddings and N text embeddings in a batch. The diagonal entries (matched pairs) should be maximised; all off-diagonal entries should be minimised.

Figure 2. Contrastive pre-training: the N×N similarity matrix of all image-text dot products. We want matched pairs (diagonal, highlighted in pink/teal) to have high dot products, and all off-diagonal entries to be low. The cross-entropy loss enforces this — the same mechanism used for next-token prediction in language models.

This works, but it has a fundamental scaling problem. Softmax requires computing a normalisation constant across the entire row or column. Every device that computes part of the similarity matrix must communicate with every other device to finish the softmax. This limits batch sizes and creates synchronisation overhead that becomes painful at scale.

Figure 3. The original CLIP pseudocode from the paper. The key line is logits = np.dot(I_e, T_e.T) np.exp(t) — all possible dot products — followed by a symmetric cross-entropy loss applied both row-wise and column-wise. The hand-written annotations highlight that this computes all possible dot products, then uses cross-entropy to teach the model which item in each row/column should be maximised.*

SigLIP (Google, 2023) eliminates this bottleneck by replacing softmax with sigmoid. Instead of treating each row as a multi-class classification problem — “which of these N texts matches this image?” — SigLIP treats each cell in the similarity matrix as an independent binary classification: “do this image and this text correspond? Yes or no?”

Figure 4. The softmax normalisation factor problem. To compute the normalization constant, you must go through ALL elements of each row AND each column. Because the CLIP similarity matrix is asymmetric, this is done twice. The SigLIP paper’s highlighted excerpt confirms this is why they switched to a sigmoid loss — shown on the right, where each cell is now an independent binary classification with no cross-cell normalisation.

The sigmoid function maps any real number to (0, 1) independently. No normalisation constant. No cross-device communication. The loss is simply:

loss = -mean( log σ(label × dot_product) )
where label = +1 for matched pairs, −1 for unmatched pairs

Figure 5. SigLIP parallel computation. With the sigmoid loss, each device computes its own block of the similarity matrix completely independently — no cross-device communication is needed. This enables training on batch sizes orders of magnitude larger than CLIP.

The practical effect is that SigLIP can train on batch sizes orders of magnitude larger than CLIP, with each device computing its block of the similarity matrix completely independently.

The Vision Transformer

The vision encoder itself is a Vision Transformer (ViT). An image is divided into a grid of 16×16 pixel patches. For a 224×224 image, that gives 196 patches. Each patch is linearly projected into an embedding vector by a Conv2d layer with kernel_size=16 and stride=16 — the stride ensures patches do not overla

self.patch_embedding = nn.Conv2d(
    in_channels=3,          # RGB
    out_channels=embed_dim, # 768
    kernel_size=16,
    stride=16,
    padding="valid"         # no padding — clean patch boundaries
)

The resulting 196 vectors are flattened into a sequence, and a learned positional embedding is added to each one. These positional embeddings are not sinusoidal functions as in the original Transformer — they are trainable parameters, one per patch position, that the model adjusts during training to encode whatever spatial information it finds useful.

The key architectural difference from a language model encoder is that there is no causal mask. Every patch attends to every other patch freely and bidirectionally. This is correct for images: the brightness of a patch in the bottom-right corner depends on the light source in the top-left corner. There is no temporal or sequential order to enforce.

Normalization: Why Modern LLMs Moved Away from LayerNorm

Before diving into the language model, it is worth understanding the normalisation progression — from Batch Norm to Layer Norm to RMSNorm — because each step was motivated by a specific failure of the previous approach.

Figure 6. The covariate shift problem. Top: stable input gives stable output through layers L₁ and L₄. Bottom: when the input magnitude shifts drastically between batches, the intermediate representation x′ changes drastically too — forcing downstream layers to constantly re-adapt. The result is a cascade: big input change → big output change → big loss change → big gradient → big weight update → the network learns slowly.

The progression to address this is clean: Batch Norm computed statistics across the batch dimension — but this means the normalisation for any single example depends on whatever other examples happen to be in the same batch, requiring large batch sizes to be stable. Layer Norm fixed this by computing statistics within each individual example across its own feature dimensions — batch composition no longer matters, and batch size 1 works fine. RMSNorm went one step further: its paper argues that the benefit of LayerNorm comes from its re-scaling effect, not its re-centering effect. If that is true, there is no need to compute the mean at all — only the Root Mean Square is needed.

Figure 7. Root Mean Square Normalisation (Zhang & Sennrich, 2019). The highlighted hypothesis: “re-scaling invariance is the reason for success of LayerNorm, rather than re-centering invariance.” This justifies dropping the mean entirely — one statistic (RMS) instead of two (mean + std). Each example is still treated independently. Most modern LLMs (LLaMA, Gemma, Mistral) have adopted this.

def forward(self, x):
    rms = x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
    return x * rms * self.weight

The Language Model: Multi-Head Attention in Detail

The most conceptually dense part of building this from scratch was the attention mechanism — not because the formula is hard, but because the tensor bookkeeping requires careful thought at every step. The goal is to transform a sequence of uncontextualised embeddings into contextualised ones: in the Vision Transformer every patch attends to every other patch freely, while in the language model each token attends only to itself and past tokens.

Projecting to Query, Key, and Value

The input sequence X of shape (4, 1024) is projected through three separate parameter matrices — WQ, WK, WV. Crucially, the 1024-dimensional output of each projection is conceptually pre-split into 8 groups of 128 — one per attention head — so that the heads can run in parallel, each working on a distinct slice of every token’s embedding.

Figure 8. Step 1 — projecting X (4, 1024) through WQ, WK, WV. Each weight matrix has effective shape (1024, 8×128), producing Q, K, V each of shape (4, 8, 128). After a transpose to (8, 4, 128), each of the 8 heads has a full sequence of 4 tokens, but each token is only 128-dimensional — the slice dedicated to that head. Two reasons to split this way: parallelise the computation, and let each head learn to relate tokens differently.

Computing Attention Scores

For each head in parallel, Q (4, 128) is multiplied by Kᵀ (128, 4) and scaled by 1/√128, producing a (4, 4) score matrix. After softmax, each row sums to 1 — these are the attention weights. Multiplying by V then produces one contextualised output embedding per position: a weighted sum where the weights tell us how much each past token contributes to the current position’s output.

Figure 9. Step 3 — attention scores Q × Kᵀ / √d_head, then softmax. Top: raw scores (13.9, 21.1, −100.3, 17.5 for the first row). Bottom: after softmax, rows sum to 1 (e.g., 0.1, 0.2, 0.5, 0.3). The note “BRO, WHERE IS YOUR MASK?” is a reminder that the Vision Transformer has NO causal mask — all patches attend freely. In the language model the upper triangle is set to −∞ before softmax, zeroing those weights out.

The Output Projection W_O

After all 8 heads compute their outputs independently, they are concatenated back to shape (4, 1024). But at this point the result is just 8 independent streams glued together — Head 1’s 128 dimensions have never interacted with Head 5’s. The output projection W_O (1024, 1024) fixes this by mixing all dimensions together, so every output value is a function of every head’s contribution.

Figure 10. Step 7 — multiply by W_O (1024, 1024). Without W_O each 128-dim group remains independent — a concatenation with no cross-head interaction. W_O mixes ALL dimensions so that each output dimension is a function of every head’s output. This is what makes the heads “talk to each other” rather than remain parallel isolated streams.

The KV-Cache: Making Inference Practical

Figure 18. The KV-Cache in action. Top: the standard attention formula — softmax(Q × Kᵀ / √d_head + MASK) applied to the full sequence “I LOVE PEPPERONI”. The (3, 3) attention weight matrix times the (3, 128) V matrix produces (3, 128) contextualised embeddings. Bottom: to predict the next token, only the LAST ROW of this output is needed. The KV-Cache avoids recomputing the other rows by caching K and V from past tokens.

Without a KV-Cache, autoregressive generation is prohibitively expensive. At generation step t, you would feed all t tokens to the model, compute a full t × t attention matrix, and then use only the last row to predict the next token. The other t − 1 rows are computed and thrown away.

The KV-Cache avoids this by storing the key and value tensors for every past token, in every layer, as the sequence grows. At each new step, only the new token needs to be processed as a query. It attends to all cached past keys and values, producing exactly the last row of the attention matrix — which is all that is needed.

def update(self, key_states, value_states, layer_idx):
    if len(self.key_cache) <= layer_idx:
        self.key_cache.append(key_states)
        self.value_cache.append(value_states)
    else:
        self.key_cache[layer_idx] = torch.cat(
            [self.key_cache[layer_idx], key_states], dim=-2
        )
        self.value_cache[layer_idx] = torch.cat(
            [self.value_cache[layer_idx], value_states], dim=-2
        )
    return self.key_cache[layer_idx], self.value_cache[layer_idx]

The two phases of inference are:

Prefilling. The entire prompt (256 image tokens + text prompt tokens) is fed in a single forward pass. The GPU processes all tokens in parallel, and the KV-Cache is populated for every position at once.

Token generation. Each subsequent step feeds only the single most recently predicted token. It is appended to the KV-Cache, and attention is computed using 1 query against all cached keys and values — O(t) per step rather than O(t²).

Rotary Positional Encoding (RoPE)

Standard positional encodings add a fixed vector to each token embedding before the model sees it. This encodes absolute position, but makes it difficult for the attention mechanism to reason about relative distances between tokens.

RoPE (Su et al., 2021) takes a fundamentally different approach. Instead of modifying the token embeddings, it modifies the attention mechanism itself. Query and key vectors are rotated in pairs of dimensions, by an angle proportional to the token’s position and the dimension’s frequency.

The mathematical consequence is elegant: when two rotated vectors are dot-producted, the result depends only on the difference in their positions, not their absolute positions. And as a bonus, this dot product naturally decays as the relative distance grows, giving the model an inductive bias toward attending to nearby tokens

inv_freq = 1.0 / (rope_theta ** (torch.arange(0, dim, 2) / dim))
freqs    = torch.outer(position_ids, inv_freq)
emb      = torch.cat([freqs, freqs], dim=-1)
q_rot    = q * cos(emb) + rotate_half(q) * sin(emb)

The Attention Mask: A PaliGemma-Specific Design Choice

Standard causal language models apply a triangular causal mask to the entire input: every token can attend only to itself and earlier tokens. PaliGemma makes a different choice. The 256 image tokens and the text prompt tokens — collectively called the prefix — are not subject to any causal mask. Every prefix token can attend to every other prefix token, including tokens that appear later in the prompt. The causal constraint only begins at the first generated token.

if kv_cache is None or kv_cache.num_items() == 0:
    # Prefill phase: no masking for prefix tokens
    causal_mask = torch.full(
        (batch_size, q_len, q_len), fill_value=0, dtype=dtype, device=device
    )

The reasoning is that the prefix is a condition, not something the model is learning to generate. Allowing full bidirectional attention within the prefix gives each token a richer context, which the authors found to be beneficial.

What I Found Surprising

Building this from scratch produced a few results I did not fully anticipate.

The projection layer matters more than it looks. A single linear layer connecting a 768-dimensional vision encoder to a 2048-dimensional language model sounds trivial. In practice, the quality of the alignment it learns is critical — and it depends heavily on the vision encoder having been trained contrastively. A vision encoder trained purely on classification tasks would require a much more complex bridge, if it worked at all.

Numerical precision issues are real. RoPE requires computing cosines and sines of position-frequency products. These computations need to happen in float32 regardless of the model’s overall precision, because small errors in the rotation angles accumulate across layers and degrade generation quality. The implementation uses torch.autocast(enabled=False) to force float32 for this computation even when the rest of the model runs in bfloat16.

with torch.autocast(device_type=device_type, enabled=False):
    freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)

The causal mask for PaliGemma is different from what most tutorials describe. Every tutorial on transformer causal masking shows a lower-triangular matrix of −∞ values. PaliGemma uses an all-zeros mask during prefilling — meaning no masking at all for the prefix. Getting this wrong produces subtly degraded output that is hard to diagnose because the model still generates coherent text, just less grounded in the image.

Weight tying cannot be done naively. Simply writing lm_head.weight = embed_tokens.weight sets a Python reference, not a PyTorch-registered parameter. The weights must be tied after loading from the checkpoint — otherwise loading the checkpoint overwrites the LM head weights, breaking the tie. The tie_weights() method is called explicitly after load_state_dict().

Results

Running inference on a test image of the Bell Tower of Xi’an with the prompt "this building is " produces:

this building is the Bell Tower located in the center of Xi'an, China.
It is one of the best-preserved examples of ancient Chinese architecture,
dating back to the Ming Dynasty...

The output is correctly grounded in the image content, factually accurate, and stylistically natural. It demonstrates that the model has successfully merged the visual context from the image with the semantic knowledge encoded in the language model’s weights.

Broader Observations

Contrastive pre-training is load-bearing, not decorative. The alignment between image and text embedding spaces is not something the small projection layer can learn on its own — it inherits it from the contrastive training objective. This is what makes the architecture work at all.

Architecture decisions encode domain knowledge. RoPE, RMSNorm, Grouped Query Attention, and SigLIP are not arbitrary choices. Each was introduced to solve a specific, documented problem. Reading the original paper for each component before implementing it paid dividends in understanding why the code is structured the way it is.

The gap between “I understand the paper” and “I can implement it” is significant. The multi-head attention mechanism described in “Attention Is All You Need” can be explained in a paragraph. Actually implementing it correctly — including the right transposition order, the correct scaling factor, the causal mask shape, and the output projection — requires working through a dozen edge cases that the paper does not spell out. The gap closes when you have to write code that runs.

Conclusion

Building PaliGemma from scratch took longer than reading the papers. It also produced a qualitatively different kind of understanding — one grounded in the specific tensor shapes, specific PyTorch APIs, and specific engineering constraints that bring these models to life.

The implementation covers 1,145 lines across five files, 21 classes, and every component needed for full end-to-end inference: SigLIP contrastive vision encoding, Vision Transformer patch extraction, linear modal projection, Gemma decoder with RMSNorm and RoPE and Grouped Query Attention and a gated FFN, KV-Cache, PaliGemma prefix-LM attention masking, and Top-P nucleus sampling.

The most durable lesson is one that applies beyond this specific model: the architectural choices in modern large language models and Vision Language Models are not arbitrary. They form a coherent design lineage in which each choice was made to solve a specific problem identified in a specific paper. The path from the original Transformer to PaliGemma is a sequence of well-motivated engineering decisions, and understanding that sequence is, ultimately, the point.

Acknowledgements

This project would not have been possible without the exceptional teaching of Umar Jamil, whose six-hour YouTube tutorial *“Coding a Multimodal (Vision) Language Model from scratch in PyTorch” served as the primary reference for this implementation. Umar’s approach of drawing every tensor operation by hand and explaining the why* behind every design decision is what made the concepts stick. The hand-drawn diagrams used in this article are taken from his teaching materials. If you want to follow along with the code, his video is the best place to start.

The implementation is based on the HuggingFace Transformers implementation of PaliGemma, which was used as a reference for layer naming conventions to enable direct weight loading from HuggingFace safetensor files.

The full implementation is available on GitHub at VLM-from-Scratch. The code is annotated with shape comments at every tensor operation, making it easy to follow the data flow from input image and text prompt to generated output.