Image-conditioned discrete-diffusion CXR report drafting

Code for finetuning and evaluating an image-conditioned discrete-diffusion language model (DiffusionGemma) against an autoregressive baseline (same backbone family) on chest-X-ray VQA and any-order infill. Anonymized for double-blind review.

Setup

pip install torch transformers>=5.11 peft accelerate numpy pillow pyyaml
# optional: scipy (paired significance in eval), anthropic / openai / google-genai (baselines + judge)

Backbones are loaded from the Hugging Face IDs in the configs (model_name). Requires a GPU (bf16, ~80 GB for the diffusion backbone).

Data

Configs point to JSON indices of the public datasets:

• VQA (openmedq_json): a flat list of {dataset, split, image, question, answer} records (VQA-RAD / SLAKE / VQA-Med).
• Report / infill (report_sources, sdr_json): a flat list of {image, Findings, Impression, split} records (MIMIC-CXR).

Set the paths in the configs (or via key=value overrides) to your local copies. synthetic: {n: 16} runs a tiny smoke test with random images.

Method

• Diffusion (backbone: diff): the report is the diffused decoder canvas; the image conditions it through the native vision tower. Training = uniform-state corruption + flat masked cross-entropy over report tokens (models/dllm.py), LoRA on the text/decoder linears (scripts/train.py).
• Autoregressive (backbone: ar): the same backbone family via the image-text-to-text head, next-token CE over report tokens (the controlled baseline).
• Any-order infill (models/infill.py): a sampler hook clamps arbitrary canvas positions to fixed tokens each denoising step, so the model fills the rest bidirectionally.

Usage

Train (LoRA finetune):

accelerate launch --num_processes=1 --mixed_precision=bf16 \
  scripts/train.py --config configs/vqarad_diff.yaml     # diffusion
accelerate launch --num_processes=1 --mixed_precision=bf16 \
  scripts/train.py --config configs/vqarad_ar.yaml       # AR baseline

VQA inference + eval:

python scripts/infer.py --config configs/vqarad_diff.yaml --adapter <adapter_dir> \
  --mode vqa --n 350 --save preds.json
python scripts/eval.py preds.json                        # accuracy (closed/open)
python scripts/llm_judge_vqa.py                           # semantic LLM-judge (needs an API key)

Infill ablation (mid-sentence, {diffusion, AR} × {both-context, left-only}):

# base backbones (no --adapter); run once per backbone, join by path
python scripts/infer.py --config configs/mimic_diff.yaml --mode infill_ablation --n 800 --save abl_diff.json
python scripts/infer.py --config configs/mimic_ar.yaml   --mode infill_ablation --n 800 --save abl_ar.json
python scripts/eval.py abl_diff.json --vs abl_ar.json    # 2x2 token-F1 + paired significance
python scripts/llm_judge_infill.py                        # semantic cross-check

Flagship VLM baselines and inference-speed benchmark:

python scripts/api_vqa_baseline.py --models gemini-3.5-flash,gpt-4.1-mini,claude-sonnet-4-6
python scripts/bench_speed.py --config configs/mimic_diff.yaml

API keys for the judge / baseline scripts are read from the environment (or a local .env).

Layout

config.py, data.py, utils.py     config loader, dataset indexing/loading, logging
models/collator.py               diffusion canvas + AR sequence collators
models/dllm.py                   uniform-state corruption + masked-CE objective
models/generate.py               generation (diffusion / AR) + infill decode
models/infill.py                 fixed-position clamping sampler hook
models/report_format.py          CoT report formatting
scripts/train.py                 LoRA finetuning loop
scripts/infer.py                 VQA + infill-ablation inference
scripts/eval.py                  VQA accuracy + infill 2x2 + significance
scripts/llm_judge_{vqa,infill}.py  semantic LLM-judge scoring
scripts/api_vqa_baseline.py      zero-shot flagship VLM baselines
scripts/bench_speed.py           latency / throughput benchmark
configs/                         per-dataset diffusion + AR configs