A Progressive Evidence Localization Framework Based on Wasserstein Gradient Flows for Document Visual Question Answering
Abstract
Precise localization of evidence regions in Document Visual Question Answering is crucial for improving model interpretability and reliability. However, existing methods predominantly adopt one-step localization strategies, which often fail to effectively distinguish evidence regions from irrelevant content when page semantics are complex or evidence regions are extremely small, leading to ambiguous boundaries and inaccurate localization. To address this issue, we propose a progressive evidence localization framework based on Wasserstein gradient flow, which formulates evidence localization as an optimal transport problem over probability distributions. Since continuous gradient flows are intractable in practice, we employ the Jordan-Kinderlehrer-Otto (JKO) scheme for discrete optimization and further derive an end-to-end trainable loss function that transforms the theoretical formulation into a neural network optimization objective, enabling coarse-to-fine precise characterization of evidence regions.Experimental results demonstrate that the proposed method significantly outperforms existing approaches in both evidence localization and answer generation tasks.
Lay Summary
When answering questions about a document, a trustworthy model should not only produce correct responses but also accurately identify where the supporting evidence appears on the page. Existing approaches typically attempt to localize the relevant region in a single step. However, such strategies often fail when the page contains dense text or when the evidence is confined to a very small area.To address this limitation, we draw inspiration from the principle in physics and mathematics that systems naturally evolve toward stable low-energy states. Specifically, we design the model to progressively refine its attention, analogous to how humans read documents by first surveying the overall layout and then focusing on the most relevant details. As a result, the evidence region gradually converges from the entire page to the precise location that truly contains the answer.We further formulate this process as a practical training framework that enables end-to-end learning within neural networks. Experimental results demonstrate that our method achieves substantially more accurate evidence localization than existing approaches, particularly in challenging scenarios where the supporting evidence occupies only a very small fraction of the page. In addition, the proposed method produces higher-quality answers. These advantages improve both the reliability and interpretability of document question answering systems, highlighting the practical value of our approach for applications such as intelligent retrieval.