NanoVDR-D-HiRes: 457M document tower distilled from Qwen3-VL-Embedding-8B

README.md

@@ -0,0 +1,189 @@
+---
+license: mit
+library_name: transformers
+pipeline_tag: feature-extraction
+tags:
+  - visual-document-retrieval
+  - document-ai
+  - nanovdr
+  - knowledge-distillation
+  - retrieval
+base_model:
+  - OpenGVLab/InternViT-300M-448px-V2_5
+  - answerdotai/ModernBERT-base
+---
+
+<p align="center">
+  <img width="440" src="https://huggingface.co/nanovdr/NanoVDR-S-Multi/resolve/main/banner.png" alt="NanoVDR"/>
+</p>
+
+<h3 align="center">NanoVDR-D-HiRes — the document tower</h3>
+
+<p align="center">
+  <a href="https://huggingface.co/nanovdr">Models</a> &nbsp;|&nbsp;
+  <a href="https://github.com/Ryenhails/NanoVDR">Code</a> &nbsp;|&nbsp;
+  <a href="https://huggingface.co/spaces/nanovdr/NanoVDR-Demo">Demo</a> &nbsp;|&nbsp;
+  <a href="#citation">Paper</a>
+</p>
+
+---
+
+A **457M** encoder that turns a document page image into **one 4096-d vector**.
+It is distilled from a frozen 8B vision-language teacher
+([Qwen3-VL-Embedding-8B](https://huggingface.co/Qwen/Qwen3-VL-Embedding-8B))
+by direct representation alignment: the student is trained only to reproduce
+the teacher's page embedding, with no relevance labels, no negative mining and
+no contrastive term.
+
+Every other NanoVDR release replaces the *query* side and still needs the
+teacher to build the index. This one replaces the *document* side, so paired
+with a NanoVDR query tower the teacher is not needed at all — at indexing time
+or at query time.
+
+| | |
+|---|---|
+| Input | page image (PIL / any resolution) |
+| Output | one L2-normalised 4096-d vector |
+| Scoring | dot product |
+| Index footprint | 16.4 GB per million pages (float32) |
+| Teacher | Qwen3-VL-Embedding-8B, frozen |
+| Visual tokens | up to 6 tiles at 448x448 + one whole-page thumbnail |
+
+A `-Fast` variant with a 2-tile budget trades v3 quality for 3x fewer visual
+tokens; see the [model index](https://huggingface.co/nanovdr).
+
+## How it works
+
+```
+page ──► dynamic tiling ──► InternViT-300M-448 ──► linear ──► ModernBERT-base ──► mean pool ──► Linear(768→4096) ──► L2-norm
+         (≤6 tiles +          per-tile patch        1024→768    bidirectional,
+          thumbnail)          tokens, concatenated              no text tokens
+```
+
+The tiling rule matches the page's aspect ratio to a grid at the encoder's
+native 448px resolution and appends a thumbnail for global context, so small
+text and table cells survive without the patch sequence outrunning
+ModernBERT's 8192-token window.
+
+## Usage
+
+```python
+import torch
+from PIL import Image
+from transformers import AutoModel, AutoImageProcessor
+
+model = AutoModel.from_pretrained("nanovdr/NanoVDR-D-HiRes", trust_remote_code=True).eval()
+processor = AutoImageProcessor.from_pretrained("nanovdr/NanoVDR-D-HiRes", trust_remote_code=True)
+
+pages = [Image.open("page1.png"), Image.open("page2.png")]
+doc_emb = model.encode(pages, processor, batch_size=4)   # (2, 4096), L2-normalised
+```
+
+Pair it with a query tower and retrieval is a dot product:
+
+```python
+scores = query_emb @ doc_emb.T          # (n_queries, n_pages)
+```
+
+On Ampere or newer, switch the fused kernels on:
+
+```python
+model = AutoModel.from_pretrained(
+    "nanovdr/NanoVDR-D-HiRes",
+    trust_remote_code=True,
+    use_flash_attn=True,                       # InternViT
+    text_attn_implementation="flash_attention_2",
+).to("cuda", torch.bfloat16)
+```
+
+Weights ship in float32 so the model loads on any device; cast to bfloat16 for
+deployment.
+
+## Results
+
+NDCG@5 on the full 22-dataset ViDoRe suite.
+
+**End to end**, paired with the 70M text-only query tower distilled from the
+same teacher — the teacher runs nowhere at deployment:
+
+| | v1 | v2 | v3 | Avg |
+|---|---|---|---|---|
+| NanoVDR-D-HiRes + 70M query tower | 82.81 | 55.34 | 47.07 | **61.74** |
+| Qwen3-VL-Embedding-8B (teacher) | 87.31 | 69.76 | 56.07 | 71.05 |
+| | | | | **86.9% of the teacher** |
+
+**Document side in isolation**, scored against teacher-encoded queries, which
+isolates this model from query-side error:
+
+| | v1 | v2 | v3 | Avg |
+|---|---|---|---|---|
+| NanoVDR-D-HiRes | 83.72 | 60.92 | 50.43 | 65.02 |
+
+At the sub-1B scale this beats every retriever we reproduced, including the
+multi-vector ones, while storing 15.6x less index than the strongest of them.
+
+### Deployment cost
+
+Single H200, batch size 8, bfloat16:
+
+| | |
+|---|---|
+| Document throughput | 36.82 pages/s |
+| Peak VRAM | 3.07 GB |
+| Index | 16.4 GB per million pages |
+| Scoring 10K candidates | 9.6 ms (single CPU thread) |
+
+## Training
+
+| | |
+|---|---|
+| Objective | `1 - cos(student, teacher)` — nothing else |
+| Data | 1.20M unique page images, public and permissively licensed |
+| Dedup | perceptual hash against all three ViDoRe evaluation corpora |
+| Optimizer | AdamW, one-cycle, peak LR 1e-3, 3% warmup |
+| Effective batch | 256 |
+| Epochs | 3 |
+| Hardware | 2x H200 |
+
+Teacher targets are cached once before training, so the document and query
+towers train independently and in parallel.
+
+## Limitations
+
+- **Bounded by the teacher.** The student only reproduces
+  Qwen3-VL-Embedding-8B's embedding space; any systematic weakness of that
+  teacher carries over, and nothing in the objective lets the student exceed it.
+- **Fixed tile budget.** Tiling adapts to page aspect ratio but not to page
+  content, so a dense small-text page and a single-figure page get the same
+  visual-token budget.
+- **Evaluation is ViDoRe-only.** Three difficulty levels, six languages, eight
+  professional domains — but not in-house enterprise layouts, scanned or
+  OCR-degraded pages, or production query distributions.
+- **Uncompressed index.** One 4096-d float32 vector per page. Quantisation and
+  product quantisation are untested here, so 16.4 GB per million is an upper
+  bound.
+
+## Citation
+
+A preprint describing this model and the full ablation study is on arXiv.
+**The arXiv identifier is being added here shortly**; until then please cite:
+
+```bibtex
+@misc{nanovdr-doc-2026,
+  title        = {Toward Efficient Representation Alignment for the Visual Tower
+                  in Visual Document Retrieval},
+  author       = {Liu, Zhuchenyang and Wang, Ziyi and Zhang, Yao and Xiao, Yu},
+  year         = {2026},
+  howpublished = {\url{https://huggingface.co/nanovdr/NanoVDR-D-HiRes}},
+  note         = {arXiv preprint, identifier to be assigned}
+}
+```
+
+The companion write-up covering motivation, architecture and every ablation
+behind this model is [here](https://huggingface.co/spaces/nanovdr/NanoVDR-D-Report).
+
+## License
+
+MIT. The visual and text backbones follow the licences of
+[InternViT-300M-448px-V2_5](https://huggingface.co/OpenGVLab/InternViT-300M-448px-V2_5)
+and [ModernBERT-base](https://huggingface.co/answerdotai/ModernBERT-base).

config.json

@@ -0,0 +1,24 @@
+{
+  "architectures": [
+    "NanoVDRDocModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "modeling_nanovdr_doc.NanoVDRDocConfig",
+    "AutoModel": "modeling_nanovdr_doc.NanoVDRDocModel"
+  },
+  "dtype": "float32",
+  "embed_dim": 4096,
+  "image_size": 448,
+  "model_type": "nanovdr_doc",
+  "pixel_shuffle_r": 1,
+  "pool_type": "mean",
+  "text_attn_implementation": "sdpa",
+  "text_backbone_name": "answerdotai/ModernBERT-base",
+  "tile_max_num": 6,
+  "tile_max_total": 7,
+  "tile_min_num": 1,
+  "tile_use_thumbnail": true,
+  "transformers_version": "4.57.6",
+  "use_flash_attn": false,
+  "visual_encoder_name": "OpenGVLab/InternViT-300M-448px-V2_5"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4cffc062320b7cdb609937283c17dff15d8d302c04458fe810acb7b4a3276f6d
+size 1827909656

modeling_nanovdr_doc.py

@@ -0,0 +1,291 @@
+"""NanoVDR document tower — standalone modeling file for the Hub.
+
+A page image is tiled, every tile is encoded by a vision transformer, the patch
+tokens of all tiles are concatenated into one sequence, a bidirectional text
+backbone re-encodes that sequence (no text tokens involved), and a mean pool
+plus linear projection produces one L2-normalised vector in the teacher's
+embedding space.
+
+Retrieval is a dot product against a query vector from the matching NanoVDR
+query tower, or against the teacher itself.
+
+This file is self-contained: it is the only source of truth for the forward
+pass and does not import from the training repository.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import List, Optional, Sequence, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+
+IMAGE_SIZE = 448
+
+
+class NanoVDRDocConfig(PretrainedConfig):
+    model_type = "nanovdr_doc"
+
+    def __init__(
+        self,
+        visual_encoder_name: str = "OpenGVLab/InternViT-300M-448px-V2_5",
+        text_backbone_name: str = "answerdotai/ModernBERT-base",
+        embed_dim: int = 4096,
+        pool_type: str = "mean",
+        pixel_shuffle_r: int = 1,
+        tile_min_num: int = 1,
+        tile_max_num: int = 6,
+        tile_max_total: int = 7,
+        tile_use_thumbnail: bool = True,
+        image_size: int = IMAGE_SIZE,
+        use_flash_attn: bool = False,
+        text_attn_implementation: str = "sdpa",
+        **kwargs,
+    ):
+        self.visual_encoder_name = visual_encoder_name
+        self.text_backbone_name = text_backbone_name
+        self.embed_dim = embed_dim
+        self.pool_type = pool_type
+        self.pixel_shuffle_r = pixel_shuffle_r
+        self.tile_min_num = tile_min_num
+        self.tile_max_num = tile_max_num
+        self.tile_max_total = tile_max_total
+        self.tile_use_thumbnail = tile_use_thumbnail
+        self.image_size = image_size
+        # InternViT ships a native flash-attention path and ModernBERT picks
+        # flash-attention-2 when the package is importable. Both need Ampere or
+        # newer, so the defaults here stay portable: plain attention for the
+        # vision tower and sdpa for the text tower (sdpa dispatches to a fused
+        # kernel on capable hardware anyway). Set use_flash_attn=True and
+        # text_attn_implementation="flash_attention_2" on Ampere+ for speed.
+        self.use_flash_attn = use_flash_attn
+        self.text_attn_implementation = text_attn_implementation
+        super().__init__(**kwargs)
+
+
+# --------------------------------------------------------------------------
+# Dynamic tiling (InternVL-V2 partition rule)
+# --------------------------------------------------------------------------
+def _closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
+    best_diff, best = float("inf"), (1, 1)
+    area = width * height
+    for ratio in target_ratios:
+        target = ratio[0] / ratio[1]
+        diff = abs(aspect_ratio - target)
+        if diff < best_diff:
+            best_diff, best = diff, ratio
+        elif diff == best_diff:
+            if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
+                best = ratio
+    return best
+
+
+def dynamic_tile(
+    image,
+    min_num: int = 1,
+    max_num: int = 6,
+    image_size: int = IMAGE_SIZE,
+    use_thumbnail: bool = True,
+) -> list:
+    """Split a page into aspect-ratio-matched ``image_size`` crops.
+
+    The final element is a whole-page thumbnail when ``use_thumbnail`` is set
+    and more than one crop was produced. Returns a list of PIL images.
+    """
+    orig_w, orig_h = image.size
+    aspect_ratio = orig_w / orig_h
+    target_ratios = sorted(
+        {
+            (i, j)
+            for n in range(min_num, max_num + 1)
+            for i in range(1, n + 1)
+            for j in range(1, n + 1)
+            if i * j <= max_num and i * j >= min_num
+        },
+        key=lambda x: x[0] * x[1],
+    )
+    cols, rows = _closest_aspect_ratio(aspect_ratio, target_ratios, orig_w, orig_h, image_size)
+    resized = image.resize((image_size * cols, image_size * rows))
+    tiles = []
+    for i in range(cols * rows):
+        c, r = i % cols, i // cols
+        box = (c * image_size, r * image_size, (c + 1) * image_size, (r + 1) * image_size)
+        tiles.append(resized.crop(box))
+    if use_thumbnail and len(tiles) > 1:
+        tiles.append(image.resize((image_size, image_size)))
+    return tiles
+
+
+@dataclass
+class NanoVDRDocOutput(ModelOutput):
+    embedding: Optional[torch.FloatTensor] = None
+
+
+# --------------------------------------------------------------------------
+# Model
+# --------------------------------------------------------------------------
+class NanoVDRDocModel(PreTrainedModel):
+    config_class = NanoVDRDocConfig
+    base_model_prefix = "nanovdr_doc"
+    _no_split_modules = ["InternVisionEncoderLayer", "ModernBertEncoderLayer"]
+
+    def __init__(self, config: NanoVDRDocConfig):
+        super().__init__(config)
+        self.cfg = config
+
+        vis_cfg = AutoConfig.from_pretrained(config.visual_encoder_name, trust_remote_code=True)
+        if hasattr(vis_cfg, "use_flash_attn"):
+            vis_cfg.use_flash_attn = config.use_flash_attn
+        visual = AutoModel.from_config(vis_cfg, trust_remote_code=True)
+        # Some releases wrap the tower; take the vision half when they do.
+        self.visual_encoder = getattr(visual, "vision_model", visual)
+        self.vis_dim = getattr(vis_cfg, "hidden_size", None) or vis_cfg.vision_config.hidden_size
+
+        txt_cfg = AutoConfig.from_pretrained(config.text_backbone_name)
+        if hasattr(txt_cfg, "reference_compile"):
+            txt_cfg.reference_compile = False
+        self.text_backbone = AutoModel.from_config(
+            txt_cfg, attn_implementation=config.text_attn_implementation
+        )
+        text_dim = txt_cfg.hidden_size
+
+        r = int(config.pixel_shuffle_r)
+        self.connect = nn.Linear(self.vis_dim * r * r, text_dim)
+        self.proj = nn.Linear(text_dim, config.embed_dim)
+        if config.pool_type == "eos":
+            self.eos_token = nn.Parameter(torch.randn(1, 1, text_dim) * 0.02)
+
+        self.post_init()
+
+    # -- internals ---------------------------------------------------------
+    def _pixel_shuffle(self, patches: torch.Tensor) -> torch.Tensor:
+        B, N, D = patches.shape
+        hw = int(round(math.sqrt(N)))
+        r = int(self.cfg.pixel_shuffle_r)
+        if hw * hw != N or hw % r != 0:
+            return patches
+        x = patches.reshape(B, hw, hw, D).reshape(B, hw // r, r, hw // r, r, D)
+        x = x.permute(0, 1, 3, 2, 4, 5).contiguous()
+        return x.reshape(B, (hw // r) ** 2, r * r * D)
+
+    # -- forward -----------------------------------------------------------
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        tile_mask: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+        **kwargs,
+    ):
+        """Encode pages into L2-normalised embeddings.
+
+        Parameters
+        ----------
+        pixel_values : (B, T, 3, H, W) or (B, 3, H, W)
+            Tiled pages, zero-padded to ``T`` tiles, or a single view per page.
+        tile_mask : (B, T) bool
+            True for real tiles. Required when ``pixel_values`` is 5-D.
+        """
+        tiled = pixel_values.dim() == 5
+        if tiled:
+            B, T, C, H, W = pixel_values.shape
+            pv = pixel_values.reshape(B * T, C, H, W)
+            if tile_mask is None:
+                raise ValueError("tile_mask is required for tiled (5-D) pixel_values.")
+        else:
+            B, T = pixel_values.size(0), 1
+            pv = pixel_values
+
+        vis_dtype = next(self.visual_encoder.parameters()).dtype
+        vis_out = self.visual_encoder(pixel_values=pv.to(vis_dtype))
+        patches = vis_out.last_hidden_state if hasattr(vis_out, "last_hidden_state") else vis_out[0]
+        patches = patches.float()
+
+        if patches.dim() == 4:  # (B, C, H, W) spatial output
+            patches = F.adaptive_avg_pool2d(patches, (16, 16)).flatten(2).transpose(1, 2)
+
+        if int(self.cfg.pixel_shuffle_r) > 1:
+            patches = self._pixel_shuffle(patches)
+
+        n_per_tile, vis_dim = patches.size(1), patches.size(2)
+        if tiled:
+            patches = patches.reshape(B, T * n_per_tile, vis_dim)
+
+        pseudo = self.connect(patches)
+        max_pos = self.text_backbone.config.max_position_embeddings
+
+        if tiled:
+            attn = tile_mask.unsqueeze(-1).expand(B, T, n_per_tile).reshape(B, T * n_per_tile)
+            attn = attn.to(pseudo.device).long()
+        else:
+            attn = torch.ones(B, pseudo.size(1), device=pseudo.device, dtype=torch.long)
+
+        if self.cfg.pool_type == "eos":
+            if pseudo.size(1) > max_pos - 1:
+                pseudo, attn = pseudo[:, : max_pos - 1], attn[:, : max_pos - 1]
+            eos = self.eos_token.expand(B, 1, -1).to(pseudo.dtype)
+            pseudo = torch.cat([pseudo, eos], dim=1)
+            attn = torch.cat([attn, torch.ones(B, 1, device=attn.device, dtype=torch.long)], dim=1)
+        elif pseudo.size(1) > max_pos:
+            pseudo, attn = pseudo[:, :max_pos], attn[:, :max_pos]
+
+        text_out = self.text_backbone(inputs_embeds=pseudo, attention_mask=attn)
+        if self.cfg.pool_type == "eos":
+            pooled = text_out.last_hidden_state[:, -1, :]
+        else:
+            m = attn.unsqueeze(-1).float()
+            pooled = (text_out.last_hidden_state * m).sum(1) / m.sum(1).clamp(min=1)
+
+        emb = F.normalize(self.proj(pooled), p=2, dim=-1)
+        return NanoVDRDocOutput(embedding=emb) if return_dict else (emb,)
+
+    # -- convenience -------------------------------------------------------
+    @torch.no_grad()
+    def encode(
+        self,
+        images: Union["Image.Image", Sequence["Image.Image"]],  # noqa: F821
+        processor,
+        batch_size: int = 8,
+        device: Optional[Union[str, torch.device]] = None,
+    ) -> torch.Tensor:
+        """Tile, preprocess and encode PIL pages. Returns (N, embed_dim) on CPU."""
+        from PIL import Image  # local import; PIL is not a hard dependency of import
+
+        if isinstance(images, Image.Image):
+            images = [images]
+        device = device or next(self.parameters()).device
+        cfg, out = self.cfg, []
+        self.eval()
+
+        for start in range(0, len(images), batch_size):
+            chunk = images[start : start + batch_size]
+            px_batch, mask_batch = [], []
+            for img in chunk:
+                tiles = dynamic_tile(
+                    img.convert("RGB"),
+                    min_num=cfg.tile_min_num,
+                    max_num=cfg.tile_max_num,
+                    image_size=cfg.image_size,
+                    use_thumbnail=cfg.tile_use_thumbnail,
+                )[: cfg.tile_max_total]
+                px = processor(images=tiles, return_tensors="pt")["pixel_values"]
+                padded = torch.zeros(cfg.tile_max_total, 3, cfg.image_size, cfg.image_size, dtype=px.dtype)
+                padded[: px.size(0)] = px
+                m = torch.zeros(cfg.tile_max_total, dtype=torch.bool)
+                m[: px.size(0)] = True
+                px_batch.append(padded)
+                mask_batch.append(m)
+            emb = self(
+                pixel_values=torch.stack(px_batch).to(device),
+                tile_mask=torch.stack(mask_batch).to(device),
+            ).embedding
+            out.append(emb.float().cpu())
+        return torch.cat(out, dim=0)
+
+
+AutoConfig.register("nanovdr_doc", NanoVDRDocConfig)
+AutoModel.register(NanoVDRDocConfig, NanoVDRDocModel)

preprocessor_config.json

@@ -0,0 +1,27 @@
+{
+  "crop_size": {
+    "height": 448,
+    "width": 448
+  },
+  "do_center_crop": true,
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.485,
+    0.456,
+    0.406
+  ],
+  "image_processor_type": "CLIPImageProcessor",
+  "image_std": [
+    0.229,
+    0.224,
+    0.225
+  ],
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "shortest_edge": 448
+  }
+}