DocRouter: Prompt guided vision transformer and Mixture of Experts connector for document understanding

Document Visual Question Answering (DVQA) involves responding to queries based on the contents of document images. The emergence of large visual language models (LVLMs) has produced impressive results on documents with simple layouts, but they still struggle with documents with complex layouts, mainly because these models cannot accurately locate the content related to the prompt. In this work, we propose a two-stage text prompt fusion scheme. In the first stage, an attention-gating mechanism is used to insert the text prompt into each layer of the visual encoder, and a late fusion strategy is adopted to reduce interference with the performance of the original model. This uses the text prompt to instruct the visual encoder to produce visual content that is more relevant to the text information and to achieve fine-grained alignment between visual and textual information. In the second stage, we replace the original MLP projection module with a Mixture-of-experts (MoE) module, which better aligns the visual information with the large language model and provides stronger generalization in multi-task training. The relevant data are collected for pre-training to achieve multi-modal alignment. In addition, to ensure the coarse-grained alignment of vision and text, we propose a gated cross-attention fusion method accordingly. Through these two fusion schemes, the model effectively filters out irrelevant information and enhances accuracy in question answering. The experimental results show that DocRouter can achieve robust results with only a 2B model, and other experimental analyses show the effectiveness of our method.