--- license: mit language: - en base_model: - dmis-lab/UltraBERT-150M base_model_relation: finetune datasets: - HuggingFaceM4/Docmatix - HuggingFaceM4/FineVision - HuggingFaceM4/the_cauldron - rlhn/rlhn-680K tags: - ultravbert - ultrabert - vision-language - multimodal - document-retrieval - colpali - vidore pipeline_tag: visual-document-retrieval library_name: transformers --- UltraVBERT # UltraVBERT UltraVBERT is a 293M-parameter vision-language encoder built as a base for visual document retrievers. It combines the [UltraBERT-150M](https://huggingface.co/dmis-lab/UltraBERT-150M) text encoder with the [SigLIP2-B/16-512](https://huggingface.co/google/siglip2-base-patch16-512) vision encoder and produces a contextual embedding for every text and image token. ## Models
ModelRoleStatus
UltraVBERTBase encoder aligned to page imagesAvailable
BiUltraVBERTSingle-vector retriever for visual documentsAvailable
ColUltraVBERTLate-interaction retriever for visual documentsAvailable · 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](https://huggingface.co/ModernVBERT/modernvbert) ([arXiv 2510.01149](https://arxiv.org/abs/2510.01149)). It replaces the Ettin-150M text encoder with UltraBERT ([arXiv 2610.15028](https://arxiv.org/abs/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. A grid holds at most six tiles per side, so the largest that fits is 5 x 5, for at most 26 x 64 = 1,664 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`. ```python 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="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="", 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](https://huggingface.co/datasets/HuggingFaceM4/Docmatix), [FineVision](https://huggingface.co/datasets/HuggingFaceM4/FineVision), [The Cauldron](https://huggingface.co/datasets/HuggingFaceM4/the_cauldron) and [RLHN-680K](https://huggingface.co/datasets/rlhn/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).