blog

Hand-Writing a Vision-Language Model, Part 9: Testing What Stage 2 Actually Learned

Building the Test

scripts/test_stage2_generation.py reconstructs the trained model exactly: loads the stage 2 checkpoint's projector state, re-wraps q_proj and v_proj with apply_lora and loads each LoRALinear's saved state, rebuilds the tokenizer with <|point|> and <|/point|> added, and reloads the tiny point-token overlay. Then calls VLM.generate() with an empty prompt, matching how stage 2 was actually trained, since compute_loss never included instruction text: only [vision_patches..., reasoning_and_point_tokens...].

Tested on maze seed 0, which was in the 500-example training set. Deliberately. A "did it learn anything at all" check before ever asking about generalisation to unseen mazes.

What the First Run Actually Produced

Genuinely mixed, and worth recording exactly as it came out rather than smoothing it over.

What worked: perfectly-formed syntax the entire way through. Correct <|point|>[[x,y]]<|/point|> tags, correct "Then moving X for N cells" grammar, zero repetition-loop garbage: a real contrast against the untrained "before" case from Part 6, which degenerated into Tab.Tab.Tab... and BLOCK BLOCK BLOCK. The model also reached (915,915), this maze's actual true exit coordinate, partway through generation, correctly phrased as "to reach the exit." Real evidence that the LoRA and point-token training did something.

What did not work: two separate problems.

Problem 1: It Never Learned to Stop

After correctly reaching the exit, generation kept going. Repeating "to reach the exit" several more times, revisiting coordinates, trailing off into nonsense well past the actual answer.

Root cause, obvious in hindsight: MazeDataset never appended an end-of-sequence token to the label text. Training data just ended with a period, so there was never a training signal for "you are done, stop here." The model learned to produce the right path but was never told what comes after it.

Problem 2: Generic Shape vs. Specific Maze

The predicted path looked more like a generic "zigzag toward the bottom-right" pattern than a route grounded in this specific maze's actual walls. The true solution winds according to the real wall layout. The model's output was a suspiciously clean right-down-right-down staircase.

Plausible explanation: with only 500 examples, 1 epoch, and rank 8, the model may have picked up "the general shape most training solutions look like" rather than truly reading this specific image's walls per example. Not yet confirmed either way. Flagged as something to check across multiple seeds once the EOS fix is in, rather than concluded from a single example.

The Fix: Teach It to Stop

data/maze_dataset.py's MazeDataset.__getitem__ now appends the tokenizer's real EOS token (<|im_end|>, id 151645 for this Qwen checkpoint) to every label sequence:

token_ids = self.tokenizer(text)["input_ids"] + [self.tokenizer.eos_token_id]
input_ids = torch.tensor(token_ids)

No changes needed anywhere else. compute_loss's masked-slice logic already trains the model to predict every position in input_ids from the position before it, so this just adds one more thing to predict: after the final point tag and period, the next token is end-of-sequence.

Verified directly: decoding a dataset example now ends ...<|/point|> to reach the exit.<|im_end|>, exactly as intended.

What Retraining Will Tell Us

Retraining now with the same setup (500 examples, 6x6 mazes, LoRA rank 8). The plan once it finishes:

Test on seed 0 again to confirm it now actually stops at the exit rather than continuing past it. This tests the EOS fix directly.

Test on several held-out seeds outside the training range. This is the more important question. If the model is genuinely reading each maze's wall layout, held-out mazes should produce different paths that respect those walls. If it is just reproducing a generic shape, the paths will look suspiciously similar regardless of which maze is shown. That is the real diagnostic.

The Honest Assessment

The first run showed something real and something missing. The right syntax, the right coordinates, the right endpoint: but no sense of when to stop, and possibly no deep grounding in the specific image's walls.

Both are fixable. The EOS problem has a one-line fix. The grounding question is a data and compute question: more examples, more epochs, potentially larger rank. And it can only be answered by running the held-out test.

This is the part of a project that does not show up on the loss chart. Loss went to 0.08. Whether 0.08 means "learned the task" or "memorised the format" is only answerable by looking at what the model actually produces on inputs it has not seen.

What Is Next

Retrain with EOS appended, then run the held-out seed test to get a real read on whether the model is grounding on each maze's actual walls or reproducing a general shape. That result determines what comes next: more data, higher rank, or both.

part 8
Hand-Writing a Vision-Language Model, Part 8: Hand-Writing LoRA, and a Real Memory Bug That Was Not What It Looked Like
part 7
Hand-Writing a Vision-Language Model, Part 7: Stage 2 Kickoff, Mazes, Coordinates, and a Reserved Vocab Surprise
part 6
Hand-Writing a Vision-Language Model, Part 6: Teaching the VLM to Actually Talk
part 5
Hand-Writing a Vision-Language Model, Part 5: Stage 1 Training, the Projector Actually Learns Something
Python PyTorch LoRA HuggingFace Transformers GTX 1080 Pascal CC 6.1 CUDA Toolkit 12.5 Ubuntu 24.04