PRD — Smallest from-scratch vision-language toy on consumer hardware
Goal
Train a LLaVA-style VL model from scratch on M5 Pro: pre-trained ViT encoder (frozen first cut) + small projector + small-from-scratch TinyGPT decoder + cross-attention or projector. Evaluate on TextVQA / VQAv2-mini. Publish “smallest from-scratch VL model runnable on a Mac” with a reproducible recipe.
Tier-5 research arc. ~2 weeks. The publishable artifact is the same shape as 5.1 / 5.2: smallest-X-on-consumer-hardware. The practical artifact is a vision-conditioned text-generation path that future Tier-B work (vision specialist, screen-understanding agent) can reuse.
Why now
- The Pace planner / agent layer wants a vision route for
screen-understanding (
factory-vision-specialist.md), and the cleanest unblock for that is: do we know how to fuse a vision encoder into our decoder at all from scratch? V1 doesn’t have to be production-grade, just functional. - The smallest VL toys in the literature are ~100M (MobileVLM, etc.). “Smallest TinyGPT-class from-scratch VL” is novel enough to publish — and ahead of TTS/diffusion in the Tier-5 ordering because the cross-attention infra is closer to shipping.
- TinyGPT’s
CrossAttention.swiftalready exists (used for YOCO); fusing into a vision-encoder K/V source is a small extension.
Scope — in
- Vision encoder: load a pretrained ViT-S/16 (any reasonable CLIP-class encoder) via swift-transformers; frozen for V1.
- Projector: 2-layer MLP from ViT hidden dim → TinyGPT d_model. Trained.
- Decoder: flagship-small or flagship-huge (TinyGPT from-scratch architecture). Token-soup the projector output before the text tokens, à la LLaVA-1.5.
- New file:
Sources/TinyGPTModel/VLForward.swift— projector- token-soup forward path; routes image embedding to the prefix.
- Training:
Sources/TinyGPT/TrainVL.swift— extend SFT with VL data loading. Stage 1: train projector only with frozen decoder. Stage 2: unfreeze decoder and continue. Standard LLaVA recipe. - Data: LLaVA-Instruct (~80K), TextVQA (~30K), MS-COCO captions (~80K) — all open. Total ~200K samples, fits a 1-day training budget at 22M.
- Eval:
tinygpt eval-vqa(new) on VQAv2-mini + TextVQA-test subsets. Compare to “decoder-only with image-as-text-description” baseline (CLIP caption → decoder). - Artifact:
docs/research/vl-toy-results.md— paper-shaped.
Scope — out
- Unfrozen vision encoder. V1 frozen. Joint training is V2 research.
- Multi-image inputs. One image per prompt for V1.
- Vision agent loop (screen + tools). That’s the production
arc (
factory-vision-specialist.md); this PRD is research. - High-resolution images. 224×224 fixed for V1, like the early LLaVA papers.
Files to touch
| File | Change |
|---|---|
Sources/TinyGPTModel/VLForward.swift | new — projector + token-soup |
Sources/TinyGPTModel/VitLoader.swift | new — loader for HF ViT weights |
Sources/TinyGPT/TrainVL.swift | new — two-stage training recipe |
Sources/TinyGPT/EvalVQA.swift | new — accuracy on VQA suites |
Sources/TinyGPT/TinyGPT.swift | case "train-vl", case "eval-vqa" |
docs/research/vl-toy-results.md | new — results |
evals/vl-smoke.sh | new — load ViT, project 1 image, decoder generates a non-empty caption |
docs/PLAN.md | 5.3 ⬜ → ✅ on ship |
Don’t touch
- The existing
CrossAttention.swift(used by YOCO) — V1 LLaVA-style uses token-soup, not cross-attention. CrossAttention path is the V2 fusion option.
Acceptance criteria
-
tinygpt train-vl --vit <vit-s.safetensors> --decoder flagship-small.tinygpt --data llava+textvqa+coco --stage 1 --steps 5000 --out vl-stage1.tinygptruns and produces a checkpoint. - Stage 2 (unfrozen decoder) continues; final model captions COCO images coherently (manual rubric inspection on 20 samples).
-
tinygpt eval-vqa <model> --task vqav2-mini --limit 500scores ≥ the caption-baseline by ≥ 5pp. - Total model size on disk ≤ 500 MB (decoder + projector; ViT loaded separately).
-
docs/research/vl-toy-results.mdreports the recipe + numbers- “smallest from-scratch VL on M5 Pro” claim with the size comparison to MobileVLM / TinyLLaVA.
Reference patterns
- LLaVA-1.5 paper — the recipe template. Two-stage projector-then-decoder training.
Sources/TinyGPTModel/HFModel.swift— HF safetensors reader; reuse for ViT.Sources/TinyGPT/Sft.swift— the training-loop template. VL training is SFT with image inputs.
Open questions
- Which ViT to import. Recommendation: ViT-S/16 (DINO) — small, works for VQA; SigLIP is the modern fancy option but larger. Document the choice.
- Whether to ship the trained model on HF / R2 alongside the docs. Recommendation: yes — the artifact is the proof.