nanoVLM 230M — cross-document attention masking A/B
Four checkpoints, not one model. This repository is a controlled experiment: two training arms ×
two seeds, 10,000 steps each, identical in every respect except how packed training rows are masked.
The comparison is the artifact. Any single checkpoint here is a weak 230M research model and is not
meant for use.
Each arm lives on its own branch:
| revision |
arm |
lm_attn_packing_impl |
lm_loss_impl |
seed |
main, C_s1 |
fork |
flex_document_causal |
gather |
1 |
C_s0 |
fork |
flex_document_causal |
gather |
0 |
A_s0 |
upstream recipe |
none |
full |
0 |
A_s1 |
upstream recipe |
none |
full |
1 |
from models.vision_language_model import VisionLanguageModel # from the fork, see below
fork = VisionLanguageModel.from_pretrained("Arjun1995/nanoVLM-230M-packing-ab", revision="C_s0")
baseline = VisionLanguageModel.from_pretrained("Arjun1995/nanoVLM-230M-packing-ab", revision="A_s0")
Weights are fp32, exactly as trained — these are the tensors every number below was computed from.
Optimizer state is not included.
What is being tested
nanoVLM packs several unrelated samples into one 4096-token row to avoid padding waste. Upstream's
packing applies a single causal mask over the whole row, so a sample attends into the earlier,
unrelated samples packed beside it, and RoPE positions run continuously across document boundaries.
At inference a sample is seen alone, starting from position 0 — so this is train/inference skew, not
a tuning knob.
Arm A reproduces that behaviour. Arm C masks attention at document boundaries and restarts
positions per document.
Arm A is upstream's recipe run on the fork's code, not a checkout of upstream, which cannot train on
torch 2.14 at all. Both arms carry that compatibility fix; neither touches the loss math, the mask
semantics or the data order.
Results
1× H100, SmolLM2-135M-Instruct + siglip2-base-patch16-512 (228,063,936 params), micro-batch 2 ×
grad-accum 8, cosine over the full 10k steps, FineVision. Within a seed, both arms train on
byte-identical batches.
The decision rule was fixed before the runs: an improvement counts only if the gap between arms
exceeds the spread between seeds within an arm. Both arms are scored under both masks by the same
evaluator, so neither is judged on its own regime.
| metric |
A (upstream recipe) |
C (fork) |
outcome |
| val loss, per-document |
0.9216 |
0.8994 |
gap 0.0222 vs 0.0134 seed spread — passes |
| val loss, unmasked |
0.9201 |
0.9147 |
gap 0.0054, inside the 0.0122 spread — fails |
| peak memory |
48.6 / 46.8 GiB |
38.4 / 36.8 GiB |
−10.1 GiB, replicating to 0.1 GiB |
| throughput |
1.00× |
1.2–1.7× |
varies with which image-tile shapes compile first |
One of the two pre-registered checks did not pass, and it is reported here. That split is the
signature the hypothesis predicts: the fix should help where the metric matches inference, and do
nothing under a metric that rewards reading unrelated neighbours.
Downstream
Seven lmms-eval tasks, cluster-bootstrap CIs, Holm-corrected, with a control measuring what changing
the training seed alone does to each task:
| task |
pooled Δ (points) |
verdict |
| TextVQA |
+4.52 [+3.81, +5.26] |
credible |
| DocVQA |
+2.13 |
inconclusive |
| OCRBench |
+1.70 |
inconclusive |
| ChartQA |
+1.00 |
inconclusive |
| ScienceQA |
+0.32 |
inconclusive |
| InfoVQA |
+0.08 |
inconclusive |
| AI2D |
−1.10 |
inconclusive, trending against |
inconclusive means undetectable at the stated minimum detectable effect, not no effect. ChartQA
moves −2.84 under seed 0 and +4.84 under seed 1, with seed controls larger than either — a
demonstration that single-seed benchmark deltas at this scale cannot be trusted.
Benchmark inference never packs documents for either arm, so the packing mechanism cannot
mechanically explain which task moved. That remains an open question and is not claimed as proven.
Reproduce
Code, run logs, analysis scripts and the full write-ups:
https://github.com/aarjunsrinivasan/nanoVLM
Limitations
- Two seeds per arm. Training-seed variance cannot be properly estimated from two runs; the seed
controls in the downstream table are the only guide to it.
- 230M parameters, 10k steps, single GPU. These are research checkpoints, not useful models.
- Upstream nanoVLM is required to load these; they are not
transformers-compatible.
