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 and produces a contextual embedding for every text and image token.

Models

Model	Role	Status
UltraVBERT	Base encoder aligned to page images	Available
BiUltraVBERT	Single-vector retriever for visual documents	Available
ColUltraVBERT	Late-interaction retriever for visual documents	Available · Demo

Architecture

SigLIP2 encodes each 512 px tile of a page into 1,024 patches. The connector merges each 4 x 4 block of patches into one token, giving 64 tokens per tile, then projects each token to UltraBERT's hidden size and RMS-normalizes it. UltraBERT then attends over image and text tokens together.

UltraVBERT builds on ModernVBERT (arXiv 2510.01149). It replaces the Ettin-150M text encoder with UltraBERT (arXiv 2610.15028), normalizes the connector output, raises the page limit from 2,048 px and 16 tiles to 3,072 px and 27 tiles, and adds a training stage on rendered text pages.

Component	Configuration	Parameters
Text encoder	UltraBERT-150M: 24 layers, hidden size 768, 6 heads; local attention over ±128 tokens in 2 of every 3 layers; 40 added special tokens	151,130,880
Vision encoder	SigLIP2-B/16-512: 12 layers, 16 px patches, 512 px tiles	93,520,128
Connector	4 x 4 pixel shuffle, linear 12,288 to 768, RMSNorm	9,437,952
MLM head	UltraBERT's released head: two 768 x 768 layers and an output layer extended to 49,280 tokens	39,026,688
Total		293,115,648

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

The processor scales a page to 3,072 px on its longest edge, reduces that in 512 px steps until the page fits in 27 tiles, and adds a downscaled view of the whole page, for at most 28 x 64 = 1,792 image tokens per page. 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.

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 in image-grounded conversations, with half of the answer tokens 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

Stage	Steps	Text tokens	Data
Reading	3,600	0.36B	Docmatix pages and rendered RLHN-680K passages, half each
Mixture	34,206	8.9B	The mix below

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

Category	Share of text tokens	Sources
Text instructions	48.18%	16 FineVision text subsets
OCR and documents	16.21%	Docmatix and 9 FineVision subsets
Screenshot to code	16.71%	4 FineVision subsets
Rendered text pages	5.00%	RLHN-680K passages
Charts	5.00%	8 FineVision subsets
Captioning	3.13%	2 FineVision subsets
Tables	3.00%	6 FineVision subsets
Reasoning and math	1.61%	7 FineVision subsets
Real-world VQA	1.16%	2 FineVision subsets and The Cauldron's TallyQA

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

Category	Subsets
Text instructions	text_code_feedback, text_codefeedback_filtered_instruction, text_infinitymath, text_mathinstruct, text_mathqa, text_mathstepdpo10k, text_numinamath_cot, text_openhermes_2_5, text_OpenMathInstruct-2, text_openorca, text_orcamath, text_pythoncode25k, text_pythoncodealpaca, text_ruozhiba, text_theoremqa, text_wizardlm_evol
OCR and documents	Docmatix, olmOCR-mix-0225-documents, CoSyn_400k_document, CoSyn_400k_diagram, SynthFormulaNet, latexformulas, k12_printing, tal_ocr_eng, ocrvqa, DoclingMatix
Screenshot to code	SynthCodeNet, datik, datikz, websight
Charts	CoSyn_400k_chart, Unichart, SynthChartNet, figureqa, mapqa, chart2text, mmc_instruct, plotqa
Captioning	densefusion_1m, localized_narratives
Tables	CoSyn_400k_table, vqaonbd, robut_wikisql, robut_wtq, finqa, robut_sqa
Reasoning and math	CoSyn_400k_math, mavis_math_metagen, geomverse, geo170k(qa), geo170k(align), CoSyn_400k_circuit, CoSyn_400k_chemical
Real-world VQA	CoSyn_400k_graphic, mmevol, The Cauldron tallyqa

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. The UltraBERT-150M and SigLIP2-B/16-512 components are Apache-2.0, with their license text in LICENSE-APACHE. Training used passages from RLHN-680K (CC-BY-SA-4.0).