UltraVBERT

UltraVBERT is a 293M-parameter vision-language encoder built as a base for visual document retrievers. It combines the UltraBERT-150M text encoder with the SigLIP2-B/16-512 vision encoder to read page images and text together, producing a contextual embedding for every token.

Models

Architecture

SigLIP2 encodes each 512 px tile of a page into 1,024 patches. The connector merges every 4 x 4 block of patches into one token, so each tile becomes 64 tokens, and RMS-normalizes them to the scale of UltraBERT's word embeddings. UltraBERT then attends over image and text tokens together.

UltraVBERT extends ModernVBERT (arXiv 2510.01149) with UltraBERT (arXiv 2610.15028) in place of its Ettin-150M text encoder, a normalized connector, pages up to 3,072 px and 27 tiles instead of 2,048 px and 16, and an extra training stage on rendered text pages.

The encoder without the MLM head has 254M parameters. Weights are stored in bfloat16.

The processor resizes a page to a 3,072 px longest edge, shrinking it in 512 px steps until it fits 27 tiles, and adds a downscaled copy of the whole page, so a page uses at most 28 x 64 = 1,792 image tokens. For shorter sequences, set processor.image_processor.size = {"longest_edge": 2048} (at most 1,088 image tokens) or lower processor.image_processor.max_image_tiles.

Usage

Requires transformers>=5.20. The model code ships with the repository, so every load passes trust_remote_code=True.

import torch
from PIL import Image
from transformers import AutoModel, AutoModelForMaskedLM, AutoProcessor

processor = AutoProcessor.from_pretrained("athrael-soju/UltraVBERT", trust_remote_code=True)
page = Image.open("page.png")

# Predict a masked word from the page.
model = AutoModelForMaskedLM.from_pretrained(
    "athrael-soju/UltraVBERT", trust_remote_code=True, dtype=torch.bfloat16, device_map="auto"
)
inputs = processor(text="<image>The table reports revenue in <|mask|>.", images=[page], return_tensors="pt").to(model.device)
with torch.no_grad():
    logits = model(**inputs).logits
masked_logits = logits[inputs.input_ids == processor.tokenizer.mask_token_id]
print(processor.tokenizer.batch_decode(masked_logits.argmax(dim=-1)[:, None]))

# One 768-d embedding per text and image token.
encoder = AutoModel.from_pretrained(
    "athrael-soju/UltraVBERT", trust_remote_code=True, dtype=torch.bfloat16, device_map="auto"
)
inputs = processor(text="<image>", images=[page], return_tensors="pt").to(encoder.device)
with torch.no_grad():
    states = encoder(**inputs).last_hidden_state

AutoModelForSequenceClassification and AutoModelForTokenClassification add task heads on the same encoder.

Training

UltraVBERT is trained with masked language modelling on the answers of image-grounded conversations: half of the answer tokens are masked (80% <|mask|>, 10% random token, 10% unchanged).

• DoRA adapters (rank 64, alpha 16) on both encoders and the MLM head, peak learning rate 1e-4
• Connector trained in full at 1e-3; image tokens start at 10x the RMS of UltraBERT's word embeddings
• The 40 added token embeddings trained directly
• AdamW (beta2 0.95), weight decay 0.1, gradient clipping 1.0, bfloat16 autocast

The mixture stage uses a warmup-stable-decay schedule with a 1 - sqrt decay over the last 20% of steps. Pages train at a 1,024 px longest edge until the decay, then at random edges from 1,024 to 3,072 px.

Datasets: Docmatix, FineVision, The Cauldron and RLHN-680K. Subsets used:

Limits

• English only.
• Trained on inputs of up to 7,000 tokens and 3 images.
• Pages larger than the tile cap are downscaled, which can blur small print.

License

MIT. UltraBERT-150M and SigLIP2-B/16-512 are Apache-2.0; their license text is in LICENSE-APACHE. Training used passages from RLHN-680K (CC-BY-SA-4.0).