blog

Hand-Writing a Vision-Language Model, Part 7: Stage 2 Kickoff, Mazes, Coordinates, and a Reserved Vocab Surprise

What Stage 2 Is

Stage 1 taught the projector to map vision embeddings into a space Qwen could interpret. The result was coherent English about an image. Stage 2 is a different task entirely: teaching the model to reason about spatial position and emit structured coordinates interleaved with its reasoning text.

The mechanism comes directly from the DeepSeek "Thinking with Visual Primitives" paper. Rather than emitting coordinates as a final answer, the model learns to produce them as part of its chain of thought: spatial markers woven into the reasoning trajectory itself. The training domain here is maze navigation. Given an image of a maze, produce a step-by-step path from entrance to exit with <|point|> coordinates marking each turn.

Building the Maze Generator

data/maze.py's Maze class generates perfect mazes via recursive backtracker (DFS, iterative with a stack): visit a random unvisited neighbour, knock down the wall, backtrack on dead ends. A perfect maze is a spanning tree, so there is exactly one path between any two cells. solve() finds it via DFS with backtracking: no shortest-path algorithm needed. render() prints an ASCII debug view at double resolution, where cells and the walls between them each get their own character position.

The Honest Account of Understanding the ASCII Grid

This part took several real rounds of confusion, not one clean explanation. The sequence of misunderstandings in order.

First: thought a wall flanking a doorway in the same printed row would block movement through it. It does not. Wall passability is checked per axis. Whether you can move vertically depends only on the character directly above or below. The left and right neighbours in the same row belong to a completely unrelated pair of cells.

Second: needed the whole grid shrunk to 4x4 and every step traced by hand with explicit direction words before "walk from entry to exit through the one open doorway between them" actually clicked.

Third: once path-tracing clicked, the question of why point coordinates like (124, 624) did not look like "0 to 8" numbers for an 8x8 grid was a good one to have out loud. Grid coordinates (cell indices, small numbers, what generate and solve work in) and normalised point coordinates (0 to 999, what the model predicts) are two unrelated coordinate systems. The model never sees "cell 3 of 8", it only sees a rendered image, so labels describe position in the image rescaled to a fixed 0 to 999 range. (124, 624) means "12.4% across, 62.4% down" the image. Multiplied by 999 so the model never has to generate a decimal point: next-token predictors are much better at picking from a fixed set of discrete symbols than generating arbitrary-precision decimal text character by character.

Worth keeping for later spatial work: shrink to a small concrete example, trace every step out loud, and isolate one axis at a time the moment an axis mix-up appears.

From Raw Path to Training Label

Two functions convert a solved path into training data.

path_to_points() converts cell coordinates to pixel-centre coordinates normalised to 0 to 999, matching the paper's format regardless of maze size.

path_to_reasoning_text() groups the cell-by-cell path into straight-line runs and narrates each one, tagging the endpoint of each run with its normalised coordinate:

Starting at the entrance, moving down for 2 cells <|point|>[[124,624]]<|/point|>.
Then moving right for 1 cell <|point|>[[374,624]]<|/point|>.
Then moving up for 1 cell <|point|>[[374,374]]<|/point|>.
... to reach the exit.

This is the actual DeepSeek paper mechanism at toy scale. Spatial markers interleaved directly into the reasoning trajectory, not emitted as a final answer after all reasoning is complete.

The Reserved Vocab Discovery

Before training, <|point|> and <|/point|> need to exist as real tokens. Adding special tokens to a model normally means growing the embedding table: embed_tokens gets new rows, the LM head gets new columns, and the new rows need to be initialised somehow.

Except Qwen3's tokenizer only assigns 151,669 tokens, but the checkpoint's embedding table has 151,936 rows: 267 pre-reserved, never-trained slots. Their embedding norms all cluster around 0.3, versus 1.0+ for a real trained token. Consistently small, clearly untouched by training.

Looking at what is already named near that boundary explains why: <|object_ref_start|>, <|box_start|>, <|quad_start|>. Qwen ships pre-reserved space for vision and pointing tokens even in this text-only checkpoint, presumably shared tokenizer infrastructure across the whole Qwen family.

tokenizer.add_special_tokens() assigns <|point|> and <|/point|> directly into two of those existing slots. The embedding table never needs to actually grow.

Only those two claimed rows should update during stage 2. The other 151,934 rows stay frozen. qwen3.model.Model.make_token_embeddings_trainable() does this via a gradient hook that zeroes every row's gradient except the claimed ones. Verified with a real forward and backward pass: claimed rows get nonzero gradient, every other row gets exactly zero.

Scope Pinned for Later

An idea to restructure the project into separate Docker and Kubernetes microservices was floated and explicitly pinned until after stage 2 and the MLP-to-attention projector swap are both done. Three open questions to answer before that becomes active: how services would talk once vlm calls qwen3 and vit over a network instead of importing them (tensor serialisation cost), how a single 8GB GPU gets shared across pods, and what a "data" microservice would actually serve at runtime given it is currently just training-time preprocessing code.

What Is Next

Hand-write LoRA, then the stage 2 training script. Same skeleton as train_stage1.py, but unfreezing the projector (continuing from projector_stage1.pt) alongside LoRA parameters and the two claimed embedding rows, training on maze image and reasoning-and-point-label pairs.

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
part 3
Hand-Writing a Vision-Language Model, Part 3: The ViT (SigLIP) and the First Bridge to Qwen
related read
Going Below PyTorch: Learning Raw CUDA Kernel Development
Python PyTorch HuggingFace Transformers GTX 1080 Pascal CC 6.1 CUDA Toolkit 12.5 Ubuntu 24.04