Wiki-PRF-3B-Infoseek

Knowledge-based-Visual-Question-Answering-with-Multimodal-Processing-Retrieval-and-Filtering

This repository provides the official PyTorch implementation for Wiki-PRF, a novel three-stage method for Knowledge-Based Visual Question Answering (KB-VQA). Wiki-PRF consists of Processing, Retrieval, and Filtering stages that dynamically extract multimodal cues, perform joint visual-text knowledge retrieval, and filter irrelevant results. The paper has been accepted at NeurIPS 2025.

🪵 TODO List

• ✅ Release core test implementation
• ✅ Complete README documentation
• ✅ Release core train implementation
• ✅ Add More detailed Quick Start.

🔥 What's New

• (2026.2.27) The core training code has been released. Tell me any problem in the issue.
• (2026.2.17) The dataset setting can be found in echosight (https://github.com/Go2Heart/EchoSight)
• (2026.2.02) Model checkpoint released! You can check the Huggingface link to download the model.
• (2025.9.19) 🎉 Our paper (Wiki-PRF) is accepted as Neurlps 2025!
• (2025.10.17) 📄 Paper released on arXiv

🧠 Wiki-PRF: A Three-Stage Framework for Knowledge-Based Visual Question Answering

Official PyTorch implementation of Wiki-PRF, accepted at NeurIPS 2025.

<div align="center"> <img src="assert/guanggao.png" alt="logo" /> </div> *Wiki-PRF achieves state-of-the-art results on KB-VQA benchmarks.*

---

📌 Abstract

Knowledge-based visual question answering (KB-VQA) requires models to combine visual understanding with external knowledge. While retrieval-augmented generation (RAG) helps, it often suffers from poor multimodal queries and noisy retrieved content.

We propose Wiki-PRF, a three-stage framework:

• 🔍 Processing: Dynamically invokes visual tools to extract precise multimodal cues for querying.
• 📚 Retrieval: Integrates visual and text features to retrieve relevant knowledge.
• 🧹 Filtering: Filters out irrelevant or low-quality results using reinforcement learning rewards based on answer accuracy and format consistency.

Our method significantly improves performance on E-VQA and InfoSeek, achieving new state-of-the-art results.

---

🏗️ Architecture

<div align="center"> <img src="assert/main.png" alt="logo"/> </div>

Our framework consists of three main components:

1. 🔍 Processing Module
   Uses vision-language tools to generate accurate, grounded queries for knowledge retrieval.

2. 📚 Multimodal Retrieval Module
   Combines image and text embeddings to retrieve top-k relevant passages from a knowledge base.

3. 🧹 Filtering & Refinement Module
   Applies RL-based filtering to discard noisy context and refine the final answer generation.

---

📊 Results

📊 Main Results on E-VQA and InfoSeek

All values are accuracy (%). Best results in bold, second best <u>underlined</u>.

Zero-shot MLLMs

Retrieval-Augmented Models

Retrieval-Augmented Models with Reinforcement Learning (Ours)

🚀 Get Started

git clone https://github.com/cqu-student/Wiki-PRF.git
cd Wiki-PRF
pip install -r test/requirements.txt

Environment Variables

Some components require API keys. Set them before running:

# Required for GPT-4 based answer generation
export OPENAI_API_KEY=<your-openai-api-key>

# Required for Google PaLM based answer generation
export GOOGLE_APPLICATION_CREDENTIALS=/path/to/credentials.json
export GOOGLE_CLOUD_PROJECT=<your-gcp-project-id>
export GOOGLE_CLOUD_REGION=us-central1  # optional, defaults to us-central1

# Required for training with Weights & Biases logging
export WANDB_API_KEY=<your-wandb-api-key>

🧪 Test (Evaluation)

Run inference and evaluation with the trained model:

cd Wiki-PRF
python test/test.py

Key paths to configure inside test/test.py:

• MODEL_PATH: path to the base model (e.g. Qwen2.5-VL-3B-Instruct)
• PEFT_MODEL_PATH: path to the fine-tuned LoRA checkpoint
• DATA_ROOT: path to the evaluation data config YAML
• IMAGE_ROOT: path to the image directory

🏋️ Train

Launch distributed training with the provided script:

cd Wiki-PRF/train
export WANDB_API_KEY=<your-wandb-api-key>
bash train.sh

Key paths to configure inside train/train.sh:

• --model_name_or_path: path to the base model
• --dataset_name: path to the data config YAML (default: data_config/rag_data.yaml)
• --image_root: path to the image directory
• --output_dir: directory to save checkpoints