AntonV HF Staff

Duplicate from athrael-soju/UltraVBERT

214d44d 19 days ago

8.04 kB

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

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