Arjun1995/nanoVLM-230M-packing-ab

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

• 10k A/B
• downstream benchmarks and mechanism
• attention masking

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.