Hand-Writing a Vision-Language Model, Part 6

The Generation Loop

Added directly to VLM in vlm/model.py, since it needs the same vision-encoder to projector to Qwen wiring as forward, just called repeatedly:

@torch.no_grad()
def generate(self, pixel_values, input_ids, max_new_tokens=50, eos_token_id=None):
    self.eval()
    for _ in range(max_new_tokens):
        logits = self.forward(pixel_values, input_ids)
        next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)
        input_ids = torch.cat([input_ids, next_token], dim=1)
        if eos_token_id is not None and torch.all(next_token == eos_token_id):
            break
    return input_ids

Why Greedy Decoding Specifically

Sampling with temperature is the more interesting choice on paper. Greedy (always take the argmax token) was chosen deliberately for a specific reason: sampling introduces randomness, which is actively bad for debugging whether the projector learned anything.

If a sampled generation comes out poorly, there is no way to tell whether that is because the projector's embeddings are genuinely useless to Qwen, or just an unlucky dice roll on that particular run. Greedy removes that confound entirely. Same image in, same tokens out, every time, which matters a lot for a before/after comparison meant to be trusted by a reader.

The trade-off is real and worth naming: greedy decoding has a well-known failure mode, getting stuck in repetition loops ("the the the..."), because it can never explore a lower-probability-but-better continuation once it has committed to a token. This showed up directly in testing below. It is not a bug to fix. It is the expected signature of greedy decoding on an undertrained model.

No KV Cache

Every step re-runs the entire forward pass (vision encoder plus the full growing sequence through Qwen) rather than reusing previous steps' cached attention keys and values. That is the standard real-world optimisation for generation and it is missing on purpose. Nothing in this codebase has a cache mechanism yet, and at this scale (short captions, a 0.6B model) recomputing is cheap enough not to matter. A real optimisation to revisit later, not a bug now.

The FastAPI Endpoint

api/app.py loads the VLM once at startup, loading the frozen ViT and Qwen weights is the expensive part. It then keeps both the fresh (randomly-initialised) and the trained (checkpoints/projector_stage1.pt) projector states in memory, swapping between them per-request via vlm.projector.load_state_dict(...). That avoids reconstructing the whole frozen backbone just to compare before/after.

@app.post("/generate")
async def generate(
    image: UploadFile = File(...),
    prompt: str = Form("Describe this image:"),
    use_trained: bool = Form(True),
    max_new_tokens: int = Form(30),
):
    vlm.projector.load_state_dict(
        trained_projector_state if use_trained else fresh_projector_state
    )
    ...

One subtlety worth noting: the prompt field is optional, and testing with it empty is actually more faithful to how stage 1 was trained. Looking back at compute_loss in train_stage1.py, the sequence fed to Qwen during training was always [vision_patches..., caption_tokens...], no instruction text in between, ever. So "Describe this image:" is something the model never saw during training. An empty prompt tests exactly the thing the projector was actually trained to do.

The Before/After

Same image (dog_park.jpg, two dogs running in a field), empty prompt, max_new_tokens=20, only use_trained flipped between requests.

Before (fresh, untrained projector):

Postman response showing noisy output from the untrained projector

Output: "agnaantaartaartaainanto.Tab.Tab.Tab.TabTabTabTabTabTab BLOCK BLOCK BLOCK BLOCK BLOCK"

Pure noise, plus the exact repetition-loop pathology predicted above. Tab loops, then BLOCK loops. Consistent with an untrained projector feeding Qwen embeddings it has no way to interpret.

After (trained projector, checkpoints/projector_stage1.pt):

Postman response showing coherent output from the trained projector

Output: "I need to get to the other side of the bridge" (the dog) . . . .

Grammatically coherent English. And it says "the dog" for an image of dogs. Not a clean caption. But real, legible evidence that the projector learned to point Qwen at the right concept. A genuinely different kind of result than a loss number on a chart.

What This Moment Means

This is the first point in the whole project where "the projector is learning" stopped being a number going down and became an actual sentence a person can read and judge.

Loss curves are convincing in the abstract. Watching noise turn into "the dog" for an image of dogs is a different, more visceral kind of convincing. Five parts of architecture work, RMSNorm, RoPE, GQA, SwiGLU, patch embeddings, the projector, the glue class, all converging into one intelligible output.

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 5

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

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