Binesh Sadanandan PhD Dissertation Companion
Read thesis

Vision-language model internals

How a medical vision-language model works, end to end

MedGemma-4B is too large to instrument layer by layer, so the dissertation uses a smaller model with the same architecture. FlipLens pairs the same frozen MedSigLIP-448 image encoder with a 14.1M-parameter Gemma-3 decoder, small enough that every token and gradient can be inspected directly. This page covers the architecture, the training objective, the finetuning setup, and inference.

  • 429M parameters, frozen encoder
  • 14.1M parameters, trained decoder
  • 256 image tokens per X-ray
  • 141 words in the vocabulary
  • 74.8% accuracy

FlipLens is a research probe for mechanism. It is not a medical question-answering system, and it is not for clinical use.

Architecture

X-ray and question become one sequence

A chest X-ray goes in as pixels. The frozen encoder splits it into 4,096 patches on a 64 by 64 grid, reads them all, and pools them to 256 image tokens. The question is tokenized, and the token embedding layer maps each token id to a vector. The image tokens are prepended to the embedded question tokens, and the decoder reads the result as one sequence: image first, then words. The decoder is a small Gemma-3 stack: root-mean-square normalization (RMSNorm), rotary position embeddings (RoPE), and grouped-query attention (GQA) with six query heads over two key-value heads, under a causal mask. The dashed box with the snowflake is frozen; training never changes it.

Input image: posteroanterior chest X-ray, public domain (CC0), via Wikimedia Commons. Not a dataset image.

Full scale

The same layout at MedGemma-4B scale

FlipLens keeps the deployed model's structure at reduced size. In MedGemma-4B, the vision tower runs patch embedding and 27 encoder blocks, and a multimodal projector maps the pooled features into the language space as 256 image tokens. Text passes through the token embedding layer, is concatenated with the image tokens, and is read by 34 decoder blocks that interleave local sliding-window and global attention at a 5-to-1 ratio, with 8 query and 4 key-value heads, RoPE, and QK-norm, followed by a final RMSNorm and a linear output layer.

Architecture diagram of MedGemma3 4B. An image enters the vision tower: patch embedding, then 27 encoder blocks of LayerNorm, multi-head attention, LayerNorm, and MLP, then a final LayerNorm and a multimodal projector that outputs 256 image tokens. Text is tokenized and passed through a token embedding layer. The image tokens and embedded text tokens are concatenated and read by 34 decoder blocks, each with pre and post RMSNorm around 5-to-1 local-to-global sliding-window grouped-query attention with rotary position embeddings and QK-norm, 8 query and 4 key-value heads, and a feed-forward stage, followed by a last RMSNorm and a linear output layer.
MedGemma-4B end to end: the vision tower (bottom right) feeds the multimodal projector; its 256 image tokens are concatenated with the output of the token embedding layer (bottom left) and read by the 34-block decoder.
Training

Next-token prediction on question-answer pairs

Training adjusts only the decoder; the encoder stays frozen, so the image features do not change across paraphrases of a question. Each training example pairs a chest X-ray with a presence question and its answer, and the objective is standard next-token prediction: at every position, predict the token that comes next. Anything the model learns about wording, it learns in the decoder.

Schematic. The decoder scores the vocabulary at the answer position; the higher of yes and no is the answer. Training moves those scores toward the label.

The probe isolates the effect of the training-phrasing distribution. Training the same model three ways on 86,288 questions from two hospitals, with everything else fixed, changes the flip rate:

Canonicalone fixed phrasing 67.1%
Augmentedevery paraphrase 4.8%
Adversarialregister tied to the answer 88.4%

Flip rates over 8 seeds per regime. Every split is balanced per finding, so a text-only model scores exactly 0.500 and no regime can exploit an answer prior. Augmented separates from both narrow regimes at the maximum effect size (Mann-Whitney p = 1.6e-4, Cliff's delta 1.00), and the two narrow regimes also separate from each other (p = 3.1e-4, delta 0.97): tying a phrasing register to the answer causes more damage than narrow coverage alone. Accuracy follows the same ordering (75.3%, 66.8%, 58.1%), so the shortcut costs diagnostic performance, not only consistency. Scored only on the 24 phrasings withheld from training, the augmented regime's flip rate rises to 26.6%, against 65.9% and 87.4% for the narrow regimes.

Finetuning

Only a small fraction of the weights train

Of the probe's roughly 443M parameters, only the 14.1M in the decoder train. The deployed model takes the same idea further with a low-rank adapter (LoRA): every original weight stays frozen and a small trainable correction is added at layers 15 to 19, 0.1% of parameters. On the clean patient-disjoint split, that adapter cuts pairwise flips from 8.5% to 3.5%, a reduction of about 59%.

A low-rank update: the frozen weight matrix W plus the product of two small trainable matrices A and B gives the correction. A and B together hold far fewer numbers than W, so the update costs little and the base model stays untouched. The adapter result
Inference

A single forward pass, read at one position

Answering is a single forward pass. The X-ray is encoded once, the question is tokenized and embedded, the two are fused, and the decoder reads the sequence once. Everything is decided at the answer position, where the decoder puts a score on each vocabulary token.

  1. 1

    Tokenize

    The question becomes token ids from a 141-piece vocabulary built from chest X-ray presence questions.

  2. 2

    Embed

    The token embedding layer maps each id to a vector. These vectors, not the ids, are what the decoder reads.

  3. 3

    Encode once

    The frozen encoder turns the X-ray into 256 image tokens, computed once per image and reused.

  4. 4

    Fuse

    Image tokens go first, embedded question tokens after. One sequence, read left to right.

  5. 5

    Decode

    A short stack of blocks: grouped-query attention with 6 query and 2 key-value heads and rotary positions, a feed-forward stage, RMSNorm throughout, and a causal mask that hides the future.

  6. 6

    Read

    At the answer position the decoder scores the vocabulary. Whichever of yes and no scores higher is the answer.

Stage 1 of 6: Tokenize

Because the pass is a stack of layers, each layer's contribution can be read separately. A Jacobian lens estimates the yes-or-no margin each layer would produce; for two phrasings of one question the readouts track together in early layers, then commit to opposite answers. In the probe, flipping pairs diverge about 2.2 times more than stable ones at layer 0 (point-biserial correlation with the flip 0.26), and a rank-1 patch at layers 0 to 4 restores the flipped answer with net recovery near 1.0 while norm-matched controls do nothing (disruption 0.000). The deployed model shows the same shape at layer 16. Read every layer, word by word · The experiments that localize the cause

Scope and limitations

FlipLens is a controlled probe, not a medical question-answering system. It grounds weakly, its 141-piece vocabulary covers presence questions only, and it is not for clinical use. What it adds is the one experiment a frozen deployed model cannot run: vary the training-phrasing distribution and measure the effect. The claim it supports is sufficiency plus localization. It does not prove that MedGemma-4B has this exact origin, and the probe's early-layer locus does not map onto the deployed model's layer 16. One measurement caveat belongs next to the 74.8% accuracy above: on a balanced set a blind model, a shuffled-grounding model, and a genuinely grounded one all score about 0.50 accuracy while their areas under the curve separate (0.500, 0.502, 0.604), so grounding must be read from the margin, not from a binary answer.

Source on GitHub Weights on Hugging Face What the probe established

Cite


  

  

     Permanent link
    
  




  

    
    
    
    
  

  

    

    
Type to search. Press Escape to close.