HyperAIHyperAI

Command Palette

Search for a command to run...

DistilVDR: 双方向生徒蒸留によるコンパクトなエンドツーエンド視覚文書検索モデル

Zhuchenyang Liu Ziyi Wang Yao Zhang Yu Xiao

概要

視覚文書検索(VDR)は数十億パラメータ規模のモデルが主流であるが、コーパス全体へのインデックス作成が遅く、サービングコストも高い。従来の圧縮手法は、小規模なマルチベクトルエンコーダをゼロから学習するか、クエリ側のみを蒸留するものであり、コンパクトなシングルベクトル検索モデルをエンドツーエンドで実現するものではなかった。本研究では、DistilVDRを提案する。これは、80億パラメータの単一の視覚言語モデルを教師とし、ポイントワイズなコサイン類似度アライメント損失の下で双方向に蒸留した、5億2400万パラメータのエンドツーエンドVDRシステムである。全ての教師信号は、関連性に基づく教師あり学習済みの凍結された教師モデルの埋め込み空間から得られるため、生徒モデルの学習には関連性ラベル、ネガティブサンプリング、対照学習の項を一切必要としない。我々は、VDRにおけるテキストクエリと画像文書の入力非対称性に、非対称なエンコーダのみの生徒モデルで対応し、文書側に視覚処理能力を集中させ、クエリ側は7000万パラメータに抑えた。同一のエンコーダと学習手法を用い、文書エンコーダの視覚トークン量のみが異なる2つのモデルを公開する。DistilVDR-HiResは、ViDoRe v1+v2+v3において平均NDCG@5で61.74を達成し(80億パラメータ教師比86.9%)、高解像度に敏感なv3ベンチマークでは、再現した全ての10億パラメータ未満のベースラインを上回る。DistilVDR-Fastは、視覚トークン量を3分の1に削減しつつ、59.98を達成する。両モデルとも、100万文書のインデックスサイズは最強の10億パラメータ未満マルチベクトルベースラインと比較して15.6分の1であり、コーパスのインデックス作成速度は1桁高速である。コードはhttps://github.com/Ryenhails/NanoVDRで公開されている

One-sentence Summary

Aalto University and an independent researcher present DistilVDR, a 524M end-to-end visual document retriever bilaterally distilled from an 8B teacher via pointwise cosine alignment loss, using an asymmetric encoder-only student to match text-image asymmetry and achieving 61.74 NDCG@5 on ViDoRe v1+v2+v3 without relevance labels, with an index 15.6 times smaller than the best sub-1B multi-vector baseline.

Key Contributions

  • DistilVDR is trained via bilateral cosine alignment distillation, where student embeddings regress onto frozen embeddings from a single 8B vision-language teacher, requiring no relevance labels, negative sampling, or contrastive terms.
  • An asymmetric encoder-only student architecture concentrates visual capacity on the document side while keeping the query side at 70M parameters, matching the text-query/image-document input asymmetry of visual document retrieval.
  • A unified reproduction and profiling of twelve retrievers shows that the two released variants, DistilVDR-HiRes and DistilVDR-Fast, achieve up to 61.74 NDCG@5 on ViDoRe v1+v2+v3, outperform all reproduced sub-1B baselines, and deliver a 15.6× smaller index and an order of magnitude faster indexing.

Introduction

Visual document retrieval (VDR) encodes entire document pages as images, sidestepping OCR errors and preserving layout, tables, and figures for tasks like enterprise search and retrieval-augmented generation. Current top-performing VDR systems are heavy, with 2–8B parameters, and incur high GPU memory and indexing costs. Prior work attempted to reduce this cost through two routes: (1) training small from-scratch multi-vector models that compress the encoder but inflate the index size and scoring latency by orders of magnitude, and (2) distillation, which shrinks the query encoder but leaves the large teacher in the document path, keeping indexing and deployment expensive. Neither approach delivers a fully compact single-vector retriever. The authors close this gap with DistilVDR, a 524M end-to-end single-vector system that bilaterally distills both document and query encoders from a single 8B vision-language teacher under an asymmetric encoder-only design, achieving strong retrieval quality while dramatically reducing index footprint and processing time.

Method

The authors formulate the retrieval task as a single-vector dense retrieval problem. Given a text query qqq and a corpus of document images {d1,,dN}\{ d_1, \dots, d_N \}{d1,,dN}, the system ranks the corpus by relevance to qqq. A query encoder fqf_qfq maps the query to a vector qRk\mathbf{q} \in \mathbb{R}^kqRk and a document encoder fdf_dfd maps each document image to a vector diRk\mathbf{d}_i \in \mathbb{R}^kdiRk in the same space. Both vectors are L2-normalised, and the retrieval score is the dot product s(q,di)=qdis(q, d_i) = \mathbf{q}^\top \mathbf{d}_is(q,di)=qdi. Document vectors are computed once at indexing time.

The core architecture relies on a dual-student distillation paradigm. A frozen 8B vision-language teacher produces target embeddings of dimension k=4096k = 4096k=4096 for both modalities. Two students learn to reproduce these targets independently. The query student fqf_qfq takes a text query and the document student fdf_dfd takes a document image. Both students project to the teacher's output space and are L2-normalised. Retrieval at deployment uses the students only; the teacher is discarded after training. Refer to the framework diagram for an overview of this architecture and the two distillation objectives.

The document encoder fdf_dfd comprises 454M parameters and proceeds in four stages. First, the system enforces a fixed per-document visual-token budget by capping the per-page tile count at TmaxT_{\mathrm{max}}Tmax and appending a single thumbnail at the visual encoder's native resolution for global context. An aspect-ratio-matched grid layout is chosen based on the page's aspect ratio. Second, every tile is encoded by an InternViT-300M-448 visual encoder into 1024 patch tokens of dimension 768. Patch tokens from all tiles are concatenated into one sequence. Third, the visual sequence is mapped into the embedding space of a ModernBERT-base text backbone by a learned linear projection. ModernBERT re-encodes the projected visual tokens with bidirectional attention, acting as a contextual encoder over visual tokens. Fourth, the contextualised tokens are mean pooled, projected from 768 to 4096 dimensions by a final linear layer, and L2-normalised.

The query encoder fqf_qfq has 70M parameters and operates in three stages. The query text is prefixed by the same instruction string π\piπ used during teacher training. The prefixed sequence is encoded by a DistilBERT-base text encoder and mean pooled over the contextual token representations. Finally, the pooled vector is passed through a linear projection from 768 to 4096 dimensions and L2-normalised.

For the distillation objective, let TTT denote the frozen teacher. For a document image ddd, the teacher produces a target vector T(d)R4096T(d) \in \mathbb{R}^{4096}T(d)R4096. For a text query qqq with instruction prefix π\piπ, it produces a target T(πq)R4096T(\pi \circ q) \in \mathbb{R}^{4096}T(πq)R4096. Both targets are L2-normalised. Each student is trained independently against these cached targets under a cosine alignment loss:

Ld=1fd(d),T(d)\mathcal{L}_d = 1 - \langle f_d(d), T(d) \rangleLd=1fd(d),T(d)⟩ Lq=1fq(πq),T(πq)\mathcal{L}_q = 1 - \langle f_q(\pi \circ q), T(\pi \circ q) \rangleLq=1fq(πq),T(πq)⟩

where ,\langle \cdot, \cdot \rangle, denotes the dot product between L2-normalised vectors. The two students never share a forward pass during training, making the doc-side and query-side distillations fully decoupled. No contrastive term, hard negative, or relevance label enters the student objective.

Experiment

The paper evaluates on the full ViDoRe suite across English, multilingual, and professional-domain benchmarks, comparing two distilled student retrievers against a broad set of baselines. DistilVDR-HiRes and Fast outperform all sub-1B models, with HiRes rivaling several 2–3B retrievers, while the efficiency analysis shows that the distilled students achieve dramatic speed and storage advantages over multi-vector baselines, particularly in document encoding and index size. Ablations confirm that visual tiling benefits hardest benchmarks, data scale saturates, and a 768-dimensional output trades quality for reduced storage, and adding a contrastive loss after distillation does not improve over pure cosine alignment. Overall, the distillation approach yields strong single-vector retrieval that captures most of the teacher's quality at a fraction of the cost.

DistilVDR-HiRes and DistilVDR-Fast, both single-vector models under one billion parameters, surpass every sub-1B baseline on the ViDoRe retrieval benchmarks, including multi-vector alternatives. HiRes delivers the strongest average NDCG@5, leading the next-best sub-1B retriever by 8.73 points, with a substantial advantage on the challenging v3 reports dataset, while Fast achieves similar quality on v1 and v2 with a smaller visual token budget. Both variants remain competitive with 2–3B models and retain over 84% of the 8B teacher’s performance. Against the top sub-1B multi-vector baseline (colSmol-500M, average 53.01), HiRes leads by 8.73 points and Fast by 6.97 points. On the hardest v3 benchmark (long professional reports), HiRes reaches 47.07 NDCG@5, 13.55 points above the best sub-1B competitor, while Fast scores 43.66, more than 10 points ahead of any other sub-1B model.

Decomposing the gap between the student and teacher shows that swapping the document encoder costs 6.03 NDCG@5 on average, while swapping the query encoder costs 4.69 points. The full student–student system trails the teacher–teacher oracle by 9.31 points, meaning the two per-side losses are not additive and an interaction effect accounts for the residual. The document-side substitution (T×S) reduces average NDCG@5 by 6.03 points, slightly more than the query-side substitution (S×T) at 4.69 points. The combined student system (S×S) suffers a 9.31-point gap from the teacher oracle, showing that the per-side losses do not simply add together.

DistilVDR's single-vector design achieves a dramatic efficiency advantage over multi-vector baselines: its Fast variant delivers the highest document throughput and lowest VRAM usage among sub-1B models, while both Fast and HiRes encode queries faster than any profiled system. Compared to multi-vector alternatives, DistilVDR requires roughly 16× less index storage and performs scoring two orders of magnitude faster, making end-to-end deployment far more practical. DistilVDR-Fast achieves the highest document throughput (99 docs/s) and uses only 2.1 GB peak VRAM, an order of magnitude faster than every multi-vector baseline. DistilVDR-HiRes recovers part of that throughput while still running 7× faster than the 8B teacher. Query encoding takes just 3.4 ms via the text-only DistilBERT path, faster than all other profiled systems. Multi-vector sub-1B models store 256 GB per million documents, whereas DistilVDR stores only 16.4 GB, a 16× reduction. Scoring 10,000 documents takes 9.6 ms for DistilVDR compared to 1.1–3.2 seconds for multi-vector baselines, roughly a 100× latency improvement.

Ablating DistilVDR's design reveals clear trade-offs: increasing the visual tile budget from 2 to 6 tiles improves average NDCG@5 by 1.76 points, with a 3.41-point gain on the hardest benchmark, at the cost of tripling visual tokens. Training-data scale shows monotonic improvement, with the full dataset yielding about 5 points over a quarter-scale subset and saturating beyond 75%. Reducing the output dimension from 4096 to 768 shrinks the index by 5.3× but loses 3.19 average points, while swapping to a light query encoder cuts latency by 5× at a cost of 4.69 points. HiRes (6 tiles) reaches 61.74 average NDCG@5, gaining 1.76 points over Fast (2 tiles), with the benefit concentrated on the complex v3 benchmark (3.41 points). Scaling training data from 25% to 100% of the 1.20M image mixture adds about 5 average NDCG@5 points, with quality saturating above 75% scale. A 768-dimensional output reduces the index size by 5.3× (to 3.07 GB per million documents) but incurs a 3.19 average point drop, larger on v3 (4.77 points). Replacing the 524M query encoder with a 105M backbone reduces query latency by 5× while sacrificing 4.69 average NDCG@5.

Adding a contrastive term to the joint refinement loss does not improve retrieval quality over cosine-only alignment. InfoNCE causes a mild performance decline as its weight increases, while KL divergence remains flat and nearly redundant with the existing cosine losses. The cosine-only refinement (γ=0) achieves the highest average NDCG@5 of 58.45, outperforming all variants with added contrastive supervision. Increasing the weight of InfoNCE from 0.5 to 2.0 reduces average NDCG@5 from 58.33 to 57.95, a consistent downward drift. KL divergence variants yield NDCG@5 scores within ±0.14 of the cosine-only baseline, showing no meaningful gain.

Evaluated on ViDoRe retrieval benchmarks, the single-vector DistilVDR models under one billion parameters outperform all sub-1B baselines, including multi-vector systems, while dramatically reducing index storage and scoring latency and remaining competitive with much larger models. The HiRes variant excels on the most challenging long-document reports, and the Fast variant achieves the highest throughput. Ablations confirm that increasing visual token budget benefits complex queries, data scaling saturates around 75%, and the simple cosine distillation loss is optimal, with no gains from contrastive objectives.


AIでAIを構築

アイデアからローンチまで — 無料のAIコーディング支援、すぐに使える環境、最高のGPU価格でAI開発を加速。

AI コーディング補助
すぐに使える GPU
最適な料金体系

HyperAI Newsletters

最新情報を購読する
北京時間 毎週月曜日の午前9時 に、その週の最新情報をメールでお届けします
メール配信サービスは MailChimp によって提供されています