MedCausalX: Adaptive Causal Reasoning for Trustworthy Medical Vision-Language Models

When Models Learn to Ask Why: Adaptive Causal Reasoning for Trustworthy Medical Vision-Language Models

Dependencies

Core requirements:

• transformers >= 4.45.0
• trl >= 0.22.0
• torch >= 2.0.0
• peft >= 0.11.0
• flash-attn >= 2.3.0 (optional)

See requirements.txt for full list.

Hardware Requirements

• Training: 6× NVIDIA A100 (40GB) with DeepSpeed ZeRO-3
• Inference: 1× A100 (40GB) or 2× RTX 4090
• Storage: ~100GB for model + data

Setup

# Clone repository
git clone https://github.com/your-org/MedCausalX.git
cd MedCausalX

# Create environment
conda create -n medcausalx python=3.10
conda activate medcausalx

# Install dependencies
pip install -r requirements.txt

# Install package
pip install -e .

# Install Flash Attention (optional, for faster training)
pip install flash-attn --no-build-isolation

---

🚀 Quick Start

1. Data Preparation:

   python scripts/prepare_crmed_dataset.py \
       --sa_med_path /path/to/SA-Med2D-20M \
       --mimic_path /path/to/mimic-cxr \
       --output_dir ./data/crmed
   

2. Training:

   bash scripts/run_full_pipeline.sh
   # Or run stages individually:
   python scripts/train_sft.py
   python scripts/train_dpo.py
   python scripts/train_grpo.py
   

3. Evaluation:

   python scripts/evaluate.py \
       --checkpoint outputs/grpo/final \
       --benchmarks slake vqa-rad pathvqa
   

4. Inference:

   python scripts/inference.py \
       --checkpoint outputs/grpo/final \
       --image path/to/xray.jpg \
       --query "What abnormalities are present?"
   

Data Preparation

# From SA-Med2D-20M
python scripts/prepare_crmed_dataset.py \
    --source_dir /path/to/SA-Med2D-20M \
    --output_dir ./data/crmed \
    --dataset_type samed2d

# From MIMIC-CXR
python scripts/prepare_crmed_dataset.py \
    --source_dir /path/to/mimic-cxr \
    --output_dir ./data/crmed \
    --dataset_type mimic-cxr

MedCausalX Project Structure

MedCausalX/
│
├── README.md                           # Project overview and quick start
├── ARCHITECTURE.md                     # Detailed architecture documentation
├── LICENSE                             # Apache 2.0 license
├── setup.py                            # Package installation script
├── requirements.txt                    # Python dependencies
├── .gitignore                          # Git ignore patterns
├── __init__.py                         # Package initialization
│
├── models/                             # Core model implementations
│   ├── __init__.py
│   ├── medcausalx.py                   # Main MedCausalX model class
│   ├── causal_tokenizer.py             # Tokenizer with reflective tokens
│   └── qwen_vl_wrapper.py              # Qwen2.5-VL wrapper utilities
│
├── datasets/                           # Dataset processing
│   ├── __init__.py
│   └── crmed_dataset.py                # CRMed dataset loader and collator
│
├── training/                           # Training modules
│   ├── __init__.py
│   ├── sft_trainer.py                  # Causal SFT trainer (Stage I)
│   ├── dpo_trainer.py                  # Error-attributed DPO (Stage II-A)
│   ├── grpo_trainer.py                 # On-policy GRPO (Stage II-B)
│   └── rewards.py                      # Composite reward functions
│
├── utils/                              # Utility functions
│   ├── __init__.py
│   └── metrics.py                      # Evaluation metrics
│
├── scripts/                            # Executable scripts
│   ├── prepare_crmed_dataset.py        # Data preparation
│   ├── train_sft.py                    # SFT training script
│   ├── train_dpo.py                    # DPO training script
│   ├── train_grpo.py                   # GRPO training script
│   ├── evaluate.py                     # Multi-benchmark evaluation
│   ├── inference.py                    # Interactive inference
│   └── run_full_pipeline.sh            # Complete training pipeline
│
└── configs/                            # Configuration files
    ├── sft_config.yaml                 # SFT hyperparameters
    ├── dpo_config.yaml                 # DPO hyperparameters
    ├── grpo_config.yaml                # GRPO hyperparameters
    └── ds_config_zero3.json            # DeepSpeed ZeRO-3 config