A from-scratch PyTorch implementation of this model, bit-exact with transformers, plus an interactive explainer

#32
by melihgenel - opened

Hi! I rebuilt PaliGemma (SigLIP + linear projector + Gemma) from scratch in plain PyTorch and loaded these weights into it, to
understand the architecture end to end.

With paligemma-3b-pt-224 in bf16 it matches transformers exactly: logits are bit-identical, layer-by-layer error is 0, and greedy
generation is token-for-token the same, with and without a KV cache.

Two details that might help others comparing implementations:

  • HF only makes the prefix (image + prompt) bidirectional when token_type_ids is passed. Without it the mask is fully causal, which
    differs from the paper and gives very different logits.
  • Only text embeddings are scaled by √d. Dividing image features by √d and letting Gemma multiply everything is the same math, but
    in bf16 it adds two roundings and shows up as a ~0.2 mean logit difference.

Code: https://github.com/mlihgenel/paligemma-from-scratch
Interactive handbook (English/Turkish): https://mlihgenel.github.io/paligemma-from-scratch/

Feedback welcome!

Sign up or log in to comment