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.

23.46M parameters. Trained in 6000 steps (~30 min) 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 Flickr30k official test split it reaches 8.30 mean R@1 against the teacher's 83.25 β€” roughly 10% of its teacher. Chance on 1000 images is 0.10.

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

  • The token budget barely pays. Going from 49 to 121 tokens buys 0.5 mean R@1. The router fits the teacher's busyness signal (loss 1.52 β†’ 0.45), but that does not yet convert into retrieval. Treat the k dial as working infrastructure with an unproven payoff, not a demonstrated win.
  • 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 6000 steps.

Measured results

Kaggle kernel gugu888/foveavit v15, status COMPLETE, T4 x2, 6000 steps, batch 48, fp16. Raw JSON: benchmarks_kaggle_v15.json. Training curve: train_log_v15.jsonl.

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 6.20 21.20 9.40 7.80
Corpus-ViT 8 73 6.00 20.90 9.50 7.75
Corpus-ViT 16 97 6.60 20.80 9.30 7.95
Corpus-ViT 24 121 6.80 20.50 9.80 8.30

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 7.17 8.47
Corpus-ViT 24 121 7.83 9.08

Imagenette (12044 images, 9 classes)

model zero-shot kNN linear probe
teacher 99.20 99.36 99.42
Corpus-ViT k=0 36.87 64.01 73.15
Corpus-ViT k=24 36.37 63.90 73.46

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
tokens out 49 61 73 97 121

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 v15. 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.

Downloads last month
14
Safetensors
Model size
23.5M params
Tensor type
F32
Β·
Inference Providers NEW
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support