Tiny-VLM OCR Study

Three tiny vision-language OCR models that share one LLM backbone and differ only in the vision frontend + attention mask, plus one extension arm on the image-position axis. Task: document image → markdown.

Experiment	Vision frontend	Attention over image tokens	Image position
llava16	SigLIP2 (frozen) + pixel-unshuffle(×2) + MLP, AnyRes	causal	SigLIP internal
gemma4_bi	encoder-free 32×32 raw-patch projector (Gemma-4-flavored)	bidirectional within each image block	learned 2D tables
gemma4_causal	same as gemma4_bi	causal (ablation)	learned 2D tables
gemma4_bi_rope	same as gemma4_bi (image_pos: rope2d)	bidirectional	2D M-RoPE in attention

gemma4_bi_rope is an extension arm (config flag image_pos: learned | rope2d): identical to gemma4_bi but image tokens get their (row, col) position via 2D M-RoPE inside gemma-3's attention (Qwen2-VL style, vlm_ocr/mrope.py) instead of learned additive tables — the clean learned-vs-rotary ablation. Text positions collapse to bit-for-bit stock 1D RoPE, so only image tokens are affected. A rope2d × causal cell is a staged follow-up.

The two frontends are footprint-matched: a llava16 token covers a 32px region (SigLIP 16px patch × pixel-unshuffle 2) and so does a gemma token (a 32px raw patch), so the llava16-vs-gemma comparison isolates encoder vs no-encoder rather than token granularity. The gemma arm is a generic encoder-free projector based on Gemma-4 (keeps its op order, no-activation, factorized-2D-pos, RMSNorm→Linear), with the patch relaxed from Gemma's 48px to 32px for that match.

• LLM backbone (shared): google/gemma-3-270m-it (Gemma-lineage: sandwich RMSNorm, QK-norm, GeGLU/gelu_pytorch_tanh, RoPE, GQA 4/1, hidden 640, 262k tied vocab, √640-scaled input embeddings, 512-tok sliding window on 5-of-6 layers). Frozen in Stage 1. Gated — needs HF auth to download.
• Vision (llava16 only): google/siglip2-base-patch16-384 (ViT-B/16, hidden 768, 384²→576 patches), always frozen.
• Data: DBlake-BoxedLogic/Image-2-Markdown (100k pages, streamed). Splits are hashed by url/pdf_relpath so pages never leak across train/val/test.
• Eval: CER / WER / normalized edit-distance (not BLEU) and TEDS for tables (metrics.py). The driver's in-loop eval reports teacher-forced loss + CER/WER; TEDS and per-category breakdowns are provided as utilities for offline eval, not wired into the driver.

Architecture

VisionFrontend (ABC) → forward(image) -> (N, 640). The VLM wrapper embeds text, splices visual tokens into inputs_embeds at the image positions (scaled by √640 to match Gemma-3's scaled text embeddings), builds the attention mask, runs gemma-3, and computes loss only on assistant tokens (image + prompt labels are -100). The LLM is never modified (no resized embeddings, no added tokens) — image positions use a filler id whose embedding is overwritten. The custom 4D attention mask is honored on every Gemma-3 layer, which also overrides (neutralizes) its sliding window — so the bi-vs-causal comparison controls attention uniformly across all layers.

vlm_ocr/
  frontends.py   VisionFrontend ABC, SiglipFrontend, EncoderFreeFrontend, pixel_unshuffle
  model.py       VLM wrapper + builders (load_llm, load_siglip_frontend, build_vlm)
  masking.py     build_attention_mask (causal / bidirectional-within-image-block)
  mrope.py       2D M-RoPE for image tokens (Gemma2DRotaryEmbedding, build_mrope_positions)
  anyres.py      tile-grid selection, normalization, patchify
  data.py        streaming, hashed splits, chat formatting, collation
  metrics.py     CER / WER / edit-distance / TEDS
  train.py       two-stage full fine-tune helpers (set_stage, optimizer, train_step)
  config.py      YAML loader (ExperimentConfig + per-stage settings)
configs/         llava16.yaml, gemma4_bi.yaml, gemma4_causal.yaml, gemma4_bi_rope.yaml  (stage1 + stage2)
scripts/         train.py (driver), smoketest.py, faithfulness.py
tests/           pytest suite (hermetic, tiny random configs)

Training (2-stage, full fine-tune, no LoRA)

1. Stage 1 — align: train the frontend only; SigLIP + gemma-3 frozen.
2. Stage 2 — adapt: train projector + full gemma-3; SigLIP frozen; AnyRes ON (llava16).

Shared regularization: weight decay 0.01–0.1, warmup+cosine, grad-clip, bf16, grad-checkpointing, projector LR ≥ LLM LR. Early-stop is manual — the driver logs the val curve and checkpoints per stage (no automatic best-model selection); watch it and stop (or raise stage2.steps) accordingly.

Run

python -m venv --system-site-packages .venv && .venv/bin/pip install -r requirements.txt

.venv/bin/python -m pytest tests -q          # unit tests (hermetic, CPU)
.venv/bin/python scripts/smoketest.py        # all 4 experiment configs, real models, streamed data
.venv/bin/python scripts/faithfulness.py     # architecture claims vs live config + docs

GPU-ready when: pytest green + smoketest PASS (all 4) + faithfulness GO (all 4). None of the above launch a training run.

Train

.venv/bin/python scripts/train.py configs/llava16.yaml               # full run (config steps)
.venv/bin/python scripts/train.py configs/gemma4_bi.yaml --max-steps 50   # short trial

scripts/train.py loads the config, streams the leakage-free split (buffered shuffle, since pages cluster by document — different seed per stage, fixed seed for a stable val set), runs stage 1 → stage 2, logs training loss, evaluates on val (loss + CER/WER), and checkpoints each stage to --out. Uses GPU + bf16 + gradient-checkpointing automatically when CUDA is present (else CPU/fp32).

Note: at the default step counts each run sees ~56k pages (~0.6 epoch of the ~90k train split); raise stage2.steps for more coverage. Watch the val curve to decide when to stop.

Documented faithfulness trade-offs

• The gemma arm is a generic encoder-free projector based on Gemma-4, not a faithful reproduction. It keeps Gemma-4's structural pipeline (raw [0,1] pixels, no activation, LN→Dense→LN→+factorized-2D-pos→LN→RMSNorm→Linear) but relaxes the 48px model patch to 32px so a visual token spans the same 32px area as the llava16 arm — a footprint-matched encoder-vs-no-encoder comparison, not strict Gemma-4.
• Resolution curriculum: like llava16's AnyRes off→on, the encoder-free arm raises max_tokens stage1→stage2 (144 ≈ 384px align → 1296 ≈ 1152px adapt), so both arms follow the same low-res-align → high-res-adapt schedule. Max image size is 1152×1152 (gemma 36×36 = 1296 tokens; llava AnyRes up to 3×3 tiles).
• Gemma-4's internal embedder width is 3840; gemma-3-270m hidden is 640, so visual tokens are 640-dim (proj_dim is the internal-width knob).
• The encoder-free projector ends in RMSNorm → Linear — this is Gemma-4's real shared multimodal embedder (verified against the transformers gemma4_unified doc), not an added layer; only the final Linear is resized 3840→640.
• The backbone gemma-3-270m is Gemma-lineage — but "gemma4 style" refers only to the vision frontend; the backbone is identical across all three runs. Two Gemma-3 backbone specifics handled in code: (1) its embed_tokens scales text embeddings by √640, so model.py scales spliced visual tokens by the same factor; (2) its 512-tok sliding window (5-of-6 layers) is overridden by our custom 4D block mask (transformers returns an already-4D mask as-is on every layer), so attention control is uniform for the bi-vs-causal experiment.
• Compute note: ~168M of gemma-3-270m's 270M params is the 262k×640 tied embedding, so stage-2 "full fine-tune" is dominated by the embedding table (memory + trainable params).
• AnyRes uses one image_newline per tile row — a simplification of LLaVA-NeXT's exact scheme.
• pixel-unshuffle 576→144 is InternVL-style space-to-depth compression (sanctioned).

License

Released under the MIT License.