Corpus-ViT

A vision encoder with a runtime-selectable token budget, built to sit between an image and an LLM. A coarse 32px grid covers the whole image; a router spends extra 16px detail only on the regions that need it. Output length is 49 + 3k tokens — 49 at k=0, 121 at k=24 — so you choose the tokens/fidelity trade-off per image, at inference, with no retraining.

24.94M parameters. Trained in 60000 steps at the full k_max=49 budget on 2× Kaggle T4 by distillation from google/siglip-base-patch16-224 (86M).

Read this first: the quality caveat

This is a working, honest baseline, not a competitive encoder. On the On the held-out split it reaches 31.10 mean R@1 (at k=24), against the teacher's 82.40. Chance on 1000 images is 0.10. For reference, an earlier checkpoint of this architecture with a k_max=24 training budget scored 29.15 on the official test split, where this one scores 30.00 at k=49.

Two specific things did not work, and are stated rather than buried:

• The token budget pays, but not dramatically. Retrieval rises monotonically with k (27.25 → 29.15 → 30.00 for 49 → 121 → 196 tokens), so the router does concentrate detail usefully. But 4× the tokens buys only +2.75, and 121 tokens already gets 97% of the 196-token score — the ceiling is the encoder, not the budget. Budgets above k=24 were never supervised during training, so the gain there is inference-time only.
• An in-batch contrastive objective made things worse (9.00 → 7.35 mean R@1, and Imagenette kNN 64.0 → 59.2). The 48-way task saturates early, so it competes with the distillation gradient rather than adding signal. It is kept behind --clip-w, default 0.

If you need competitive absolute numbers, the lever is data and training scale, not architecture: this saw 31k images for 60000 steps.

Measured results

Kaggle kernel gugu888/foveavit v31 (arm A), T4 x2, 60 000 steps, batch 48, fp16, trained at the full k_max=49 budget. Raw JSON: benchmarks_kaggle_v31_kmax49_heldout.json (held-out split; the kernel was cancelled during the second ablation arm before the full benchmark ran).

Flickr30k official test split (1000 images, one caption each)

model	k	tokens	i2t R@1	i2t R@5	t2i R@1	mean R@1
teacher (SigLIP-B/16)	–	196	82.60	95.40	83.90	83.25
Corpus-ViT	0	49	26.60	—	28.70	27.65
Corpus-ViT	4	61	27.50	—	29.30	28.40
Corpus-ViT	8	73	28.20	—	30.80	29.50
Corpus-ViT	16	97	29.40	—	31.50	30.45
Corpus-ViT	24	121	30.10	—	32.10	31.10

Detail subset (captions ≥12 words or containing a digit; 920/1000)

model	k	tokens	i2t R@1	mean R@1
teacher	–	196	83.59	84.08
Corpus-ViT	0	49	15.87	17.23
Corpus-ViT	24	121	18.04	19.78

Imagenette (12044 images, 9 classes)

model	zero-shot	kNN	linear probe
teacher	99.20	99.37	99.34
Corpus-ViT k=0	46.95	69.04	76.01
Corpus-ViT k=24	47.83	69.75	76.32

Probes tell a better story than retrieval (73% linear vs 99% teacher), which is the expected pattern for a model that learns representations faster than it learns the metric head.

Test hygiene: training used the official split and held out the 1000 official test images — 30783 train / 1000 held out. Benchmark images were never trained on.

Usage

import torch
from transformers import AutoModel

model = AutoModel.from_pretrained("Gugu8/Corpus-ViT", trust_remote_code=True).eval()

# ImageNet-style normalisation to [-1, 1], 224x224.
px = torch.randn(1, 3, 224, 224)

# 1) Spatial tokens for an LLM connector: (B, 49 + 3k, 768)
tokens = model(px, k=16).last_hidden_state      # (1, 97, 768)

# 2) One retrieval/classification vector in SigLIP's joint space: (B, 768)
emb, n_tokens = model.embed(px, k=16)           # emb is L2-normalised

The embedding is directly comparable to SigLIP text embeddings:

from transformers import AutoModel, AutoTokenizer

siglip = AutoModel.from_pretrained("google/siglip-base-patch16-224").eval()
tok = AutoTokenizer.from_pretrained("google/siglip-base-patch16-224")

with torch.no_grad():
    t = siglip.get_text_features(**tok(["a dog on a beach"], return_tensors="pt",
                                       padding="max_length", max_length=64,
                                       truncation=True))
    t = torch.nn.functional.normalize(t.float(), dim=-1)
score = (emb @ t.T).squeeze()      # higher = better match

Content-driven sharpening

k uses the learned router. To pick cells by image content instead (e.g. detail that matters for a document, a chart, or a face), use sharpen with a budget fraction. Cells are ranked by how much a coarse cell disagrees with its four fine children:

emb, n = model.sharpen(px, budget=0.25)   # ~25% of cells sharpened, content-driven

Token budget

k	0	4	8	16	24	32	40	49
tokens out	49	61	73	97	121	145	169	196

tokens_out(k) returns the count without running the model.

How it works

input 224x224
  ├─ coarse conv, stride 32  -> 7x7  =  49 tokens      (always)
  └─ fine   conv, stride 16  -> 14x14 = 196 tokens      (used where needed)
                                    |
       router scores all 49 cells for "busyness"
                                    |
       top-k cells are REPLACED by their 4 fine children each
                                    |
       output = (49 - k) + 4k = 49 + 3k spatial tokens   (+4 internal registers)

• Registers are learned tokens kept out of the LLM-facing output.
• RoPE is relative and 2-D; a parent cell sits at 2i+1, the geometric centre of its four children, so mixed grids share one consistent frame.
• embed() uses a learned pool head (LayerNorm + Linear over the mean final token state), trained against the teacher's projected image embedding. This matters: mean-pooling the raw token features is not in SigLIP's text space, and retrieving with it scores at exactly chance (0.10 R@1). Getting this wrong is the single biggest trap in this repo.

Distillation setup

• Teacher google/siglip-base-patch16-224, frozen, fp16.
• SigLIP's 14x14 token grid groups exactly into 7x7 coarse blocks, so targets need no resampling.
• Per-token targets: mean of each coarse block's four child tokens.
• Pooled target: the teacher's joint-space image embedding, one tower forward.
• Router target: detail = cosine disagreement between a coarse cell and its four children, i.e. the teacher's own busyness.
• Loss = 1.0*token + 0.5*router + 1.0*pool.

Reproduce

KAGGLE_USERNAME=<user> KAGGLE_KEY=<key> \
KAGGLE_BIN=/Users/everybody/Downloads/foveal-env/bin/kaggle \
PYTHON=/Users/everybody/Downloads/foveavit-env/bin/python \
./scripts/push_foveavit.sh --push

The Kaggle notebook is self-contained — it writes the package inline and pulls Flickr30k from HuggingFace, so no dataset attachment is required.

Limitations

• ~10% of teacher retrieval; not competitive with CLIP/SigLIP at scale.
• The token-budget payoff is unproven (0.5 R@1 for 2.5x tokens).
• Trained at 224px on one GPU pair; multi-crop and higher resolution untested.
• Eval uses one caption per image (the first). Scoring against all five inflates i2t R@1 and is not comparable to published numbers.
• The detail subset threshold (12 words) is a heuristic.

Provenance

Weights come from Kaggle kernel gugu888/foveavit v24. model.safetensors was verified bit-identical to the trained checkpoint, and the modeling code in this repo reproduces the training implementation to 0.00e+00 max absolute difference across k = 0, 4, 8, 16, 24.

Note for contributors: loading with trust_remote_code

transformers' dynamic-module loader inspects each uploaded file for top-level imports and treats every non-relative from X import Y as a third-party package it must be able to importlib.import_module. Sibling imports therefore have to be relative:

from .configuration_corpusvit import CorpusVitConfig   # correct
from configuration_corpusvit import CorpusVitConfig    # breaks loading

The try/except ImportError fallback does not help — the loader's regex strips such blocks unreliably across lines, so the absolute import is still seen and the load fails with This modeling file requires the following packages that were not found in your environment: configuration_corpusvit.

Verified by loading this repo from the Hub and diffing against the training checkpoint: max|delta| = 0.00e+00.