Hand-Writing a Vision-Language Model, Part 5: Stage 1 Training, The Projector Actually Learns Something

Dataset: Flickr8k

Confirmed via the HuggingFace Hub before writing any code, same discipline as inspecting model checkpoints. Single train split, 40,455 rows (approximately 8,000 images with 5 captions each), columns image and caption, approximately 1.16GB, public, no gating. Downloaded into a repo-local flickr8k_data/ folder rather than the default global HuggingFace cache, consistent with how the Qwen and SigLIP checkpoints are stored.

One quirk worth knowing: asking datasets for split='train[:5]' still generates and downloads the entire 40,455-row split first, then slices. Non-streaming Arrow datasets do not lazily fetch just the requested rows. Not a problem in practice, one-time cost, then it is cached locally and every future slice is instant, just not what "give me 5 rows" suggests at first.

Hand-Rolled Image Preprocessing

def preprocess_image(image):
    image = image.convert("RGB").resize((224, 224), Image.Resampling.BICUBIC)
    array = np.array(image, dtype=np.float32) / 255.0
    array = (array - 0.5) / 0.5
    return torch.from_numpy(array).permute(2, 0, 1)

Values (224x224, mean/std 0.5) come directly from siglip/preprocessor_config.json. Checked against the real SiglipImageProcessor on an actual dataset image, max absolute difference: 0.0, exact match, before trusting it inside the training pipeline.

Converting to grayscale was considered and rejected. The frozen patch-embedding conv is hard-shaped for exactly 3 input channels ([768, 3, 16, 16], confirmed in Part 3). That would throw away real colour signal the pretrained model already learned to use, for zero speed benefit — the 12-layer transformer stack dominates FLOPs regardless of input channels.

The Dataset Class

class ImageCaptionDataset(torch.utils.data.Dataset):
    def __init__(self, tokenizer_dir, num_examples=1000, cache_dir="flickr8k_data"):
        self.dataset = load_dataset("ariG23498/flickr8k",
                                    split=f"train[:{num_examples}]",
                                    cache_dir=cache_dir)
        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir)

    def __getitem__(self, idx):
        example = self.dataset[idx]
        pixel_values = preprocess_image(example["image"])
        input_ids = torch.tensor(self.tokenizer(example["caption"])["input_ids"])
        return pixel_values, input_ids

Tokenisation uses Qwen's real tokenizer directly, never in scope to hand-write. Verified round-trip: decoding the tokenized caption back gives the exact original sentence.

Two Model-Level Decisions Worth Calling Out

torch.no_grad() around the ViT call, baked into VLM.forward itself:

def forward(self, pixel_values, input_ids):
    with torch.no_grad():
        patch_embeds = self.vision_encoder(pixel_values)
    vision_embeds = self.projector(patch_embeds)
    ...

Nothing upstream of the projector's input ever needs a gradient. The projector's weights depend only on the ViT's output values, not its weights. Since vision_encoder is frozen in every context this class will ever be used in, this belongs in the model's own definition rather than re-implemented per training script. Saves all of the ViT's 12-layer activation memory during backward.

Loss computed only over text-predicting positions, correctly aligned:

logits = vlm(pixel_values, input_ids)                   # (batch, num_patches + T, vocab)
num_patches = logits.shape[1] - input_ids.shape[1]
predicted = logits[:, num_patches - 1 : -1, :]           # (batch, T, vocab)
loss = F.cross_entropy(predicted.reshape(-1, vocab_size), input_ids.reshape(-1))

The model's output at position p predicts whatever sits at p+1. The last vision-patch position already predicts the first caption token, a genuinely useful signal ("given only the image, what is the first word?"), so starting the slice at num_patches - 1 aligns predicted-logit k with input_ids[k] directly. No separate "shift the targets" step needed.

MLflow Instead of Matplotlib

Asked for a quick loss-curve plot via Matplotlib first. Before that got built, reconsidered and swapped to MLflow, the more appropriate tool for actual experiment tracking. Metrics and parameters logged per step, browsable later via mlflow ui, and it scales naturally to comparing future runs (stage 2, the projector-architecture experiment) rather than producing one disposable PNG.

scripts/train_stage1.py now logs loss and perplexity (exp(loss), the standard companion metric in language modelling) per step, plus run parameters. One gotcha: MLflow 3.14's default backend creates a mlflow.db SQLite file at the project root, not just inside mlruns/. Had to gitignore that specifically alongside mlruns/, mlartifacts/, and checkpoints/.

The First Real Training Run

500 examples, batch size 1, 1 epoch, AdamW on vlm.projector.parameters() only, learning rate 1e-4:

step    0 | loss 14.8563
step  100 | loss 3.8803
step  200 | loss 3.0624
step  300 | loss 3.2037
step  400 | loss 4.0056
step  480 | loss 5.8309

Started right around ln(151936) ≈ 11.9, the loss you would expect from guessing uniformly at random over the vocabulary. Actually slightly above that at step 0, then dropped fast, settling mostly in the 2.5 to 5 range. Noisy step-to-step (batch size 1, wildly different caption lengths and content each step), but the overall trend is unmistakable.

The perplexity chart tells the same story in a more intuitive unit: "how many tokens was the model effectively choosing uniformly between" at each position. It spikes to the millions in the first few steps (loss still near-random) then collapses fast once the projector starts producing usable embeddings. Matching the loss chart exactly, just on a scale that is easier to reason about than raw cross-entropy.

What Is Next

Stage 2: LoRA fine-tune Qwen on the procedurally-generated maze dataset (DFS-traced paths, interleaved reasoning and <|point|> coordinates), the stage where the model actually learns to "think with visual primitives," not just describe images. The MLP to attention-based projector swap stays deliberately last, after stage 2 works.

part 4

Hand-Writing a Vision-Language Model, Part 4: Wiring the VLM Together, and What SigLIP's Training Actually Means

part 3

Hand-Writing a Vision-Language Model, Part 3: The ViT (SigLIP) and the First Bridge to Qwen

part 1

Hand-Writing a Vision-Language Model, Part 1: Qwen3-4B's Attention Stack

related read

Going Below PyTorch: Learning Raw CUDA Kernel Development