ICML2026
A Progressive Evidence Localization Framework Based on Wasserstein Gradient Flows for Document Visual Question Answering
Haosen Wang, Jing Xiao, Mengqiao Li, Xuanze Wang, Mingzheng Zhang, Xiaowang Zhang, Zhiyong Feng
摘要
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.