TL;DR: I rebuilt Google’s PaliGemma (SigLIP vision encoder + Gemma decoder) from scratch in PyTorch, and the most interesting part was the attention mask.
Early in my engineering journey, I loved dissecting and understanding the mathematics and logic behind core AI/ML models, but over the past year and a half, most of my learning, work, and freelance projects have revolved around agentic AI systems and pipelines.
So during this exam season, I did what every student does: found something far more interesting than my syllabus. And thus, I set out to implement the entire architecture of PaliGemma, the multimodal version of the original Gemma model, by hand in PyTorch. I finished the project right before the last exam, and now that vacations are over, I have had the time to clean up a couple of bugs and document the project.
Here’s my attempt at a brief overview of what went into it. PaliGemma, being a vision-language model, takes an image plus a text prompt and generates text.
The vision side (SigLIP ViT). A 2D convolution splits the image into patches and embeds them, learned positional embeddings are added, and the result goes through stacked layers of multi-head attention, LayerNorm, and GELU MLPs with residual connections.
The language side (Gemma decoder). This was the harder half, where I had to revisit and implement concepts like the KV cache, RMS normalization, Grouped Query Attention, and RoPE (Rotary Position Embedding).
The glue. Image embeddings are projected into the decoder’s embedding space and placed before the text tokens, so the model treats the image as part of the prompt.
One interesting detail: PaliGemma uses a prefix-LM mask. Image tokens and the prompt attend to each other bidirectionally, and only the generated answer is causal, unlike most other vision LLMs, which apply a plain causal mask over the whole sequence.
The code is on my GitHub, and it should be relatively easy to follow thanks to the comments explaining each block. The complete architecture diagram is below (it’s a Canva whiteboard, so you can pan and zoom, or open it in Canva), and the PaliGemma write-up on Hugging Face is great for visualising the whole thing.
A huge thank you to Umar Jamil. His PaliGemma breakdown, and honestly all his videos, are the best resource I know for the math, logic and code behind AI.
P.S. If you’re into model internals, multimodal models or agentic systems, or want to discuss PaliGemma, my DMs are open. I’d love to chat.