Blame for README.md

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

<img src="ultravbert-hero.png" width="100%" alt="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

<table style="display: table; width: 100%;">
<thead>
<tr><th align="left">Model</th><th align="left">Role</th><th align="left">Status</th></tr>
</thead>
<tbody>
<tr><td><a href="https://huggingface.co/athrael-soju/UltraVBERT">UltraVBERT</a></td><td>Base encoder aligned to page images</td><td>Available</td></tr>
<tr><td><a href="https://huggingface.co/athrael-soju/BiUltraVBERT">BiUltraVBERT</a></td><td>Single-vector retriever for visual documents</td><td>Available</td></tr>
<tr><td><a href="https://huggingface.co/athrael-soju/ColUltraVBERT">ColUltraVBERT</a></td><td>Late-interaction retriever for visual documents</td><td>Available · <a href="https://huggingface.co/spaces/athrael-soju/ColUltraVBERT">Demo</a></td></tr>
</tbody>
</table>

## 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="<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](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).