CAB-Diff: Reliable Answer Generation with Causal-Bidirectional Block Diffusion for Medical Visual Question Answering

Medical Visual Question Answering (Med-VQA) requires answers to closely follow visual evidence and clinical context. While modern Vision-Language Models (VLMs) leverage autoregressive generation for this task, they often suffer from error accumulation (exposure bias)—where a single incorrect token early in the sequence can derail the entire clinical reasoning chain. In medicine, minor deviations regarding negation, location, or lesion attributes can completely alter the diagnosis.

CAB-Diff is a novel method that addresses this by reorganizing the generation process into a block-wise diffusion process specifically applied to the answer region. By leveraging large vision-language models (e.g., Qwen2.5-VL-7B-Instruct, MedGemma-4B-IT) as backbones, CAB-Diff introduces a hybrid approach. It maintains a causal autoregressive branch to preserve language fluency, while introducing a bidirectional masked diffusion branch that learns to recover missing tokens. This bidirectional interaction allows for self-correction, effectively reducing accumulated errors and improving clinical reliability.

🧠 Framework Architecture

Our proposed CAB-Diff framework operates by conditioning the diffusion process on robust visual and textual embeddings, progressively removing noise to construct the final clinical answer.

---

📂 Datasets

To run the experiments, you need to download the Med-VQA datasets and place them into the datasets folder.

By default, the configuration files expect the following folder structure:

CabDiff/
└── datasets/
    ├── kvasir-vqa/
    │   ├── images/
    │   ├── train.json
    │   └── test.json
    └── slake/
        ├── imgs/
        ├── train.json
        └── test.json

(Note: If you place them in a different folder, you must update the data_root property in the YAML configuration files accordingly)

---

🚀 Getting Started

The training and inference pipelines are managed via SLURM scripts and YAML configuration files.

1. Baseline Model (Qwen / MedGemma)

Before training the CABdiff diffusion wrapper, you need to fine-tune the baseline Vision-Language Model (LoRA).

Train Baseline:

export CONFIG=configs/baseline/train_baseline.yaml
sbatch scripts/train.slurm

Infer Baseline:

export CONFIG=configs/baseline/infer_baseline.yaml
sbatch scripts/infer.slurm

2. CAB-Diff Model

Once the baseline is trained, its checkpoint is used as a warm-start (resume_from) for the CAB-Diff model.

Train CAB-Diff: Ensure your CAB-Diff config file (e.g., configs/cabdiff/kvasir/train_qwen_cabdiff.yaml) has resume_from pointing to the baseline checkpoint.

export CONFIG=configs/cabdiff/kvasir/train_qwen_cabdiff.yaml
sbatch scripts/train.slurm

Infer CAB-Diff:

export CONFIG=configs/cabdiff/kvasir/infer_qwen_cabdiff.yaml
sbatch scripts/infer.slurm

---

🤝 Contributing

For inquiries, feature requests, or benchmarking ideas, please visit our repository: https://github.com/SharkTien/CabDiff and open an Issue or Pull Request.