liguo12/REVA_PO_Weights

REVA-PO Weights

Pretrained checkpoints for REVA-PO: Stabilizing Reinforcement Learning for Chest X-ray Report Generation, accepted to ECCV 2026.

Overview

REVA-PO is a reinforcement learning framework for chest X-ray report generation. It combines:

• Response-Weighted Regularization, which adjusts response-level KL regularization using the response advantage and reference-policy entropy.
• Validation-Anchored Policy Reset, which periodically synchronizes the current and reference policies with the best validation checkpoint.
• A three-stage training pipeline consisting of warm-up training, classifier-guided supervised fine-tuning, and reinforcement learning.

This repository contains the PyTorch and DeepSpeed checkpoints used to reproduce the IU-Xray and MIMIC-CXR experiments.

These files are project-specific .pth checkpoints. They are not standalone Hugging Face Transformers checkpoints and cannot be loaded directly with AutoModel.from_pretrained(). Use them with the released REVA-PO code and configuration files.

Repository Contents

REVA_PO_Weights/
├── Medical_MAE/
│   └── vit-b_CXR_0.5M_mae_CheXpert.pth
├── iuxray/
│   ├── iuxray_cls_ckpt.pth
│   ├── iuxray_stage1_ckpt.pth
│   ├── iuxray_stage2_ckpt.pth
│   ├── iuxray_stage3_ckpt.pth
│   └── mimic_stage1_ckpt.pth
└── mimic/
    ├── chexbert.pth
    ├── mimic_cls_ckpt.pth
    ├── mimic_stage1_ckpt.pth
    ├── mimic_stage2_ckpt.pth
    └── mimic_stage3_ckpt.pth

Checkpoint Guide

Shared and Auxiliary Checkpoints

File	Purpose
Medical_MAE/vit-b_CXR_0.5M_mae_CheXpert.pth	Medical MAE initialization used when fine-tuning the image classifier.
mimic/chexbert.pth	CheXbert checkpoint used by the MIMIC-CXR training and evaluation pipeline.

MIMIC-CXR Checkpoints

File	Purpose
mimic/mimic_cls_ckpt.pth	Image classifier fine-tuned for MIMIC-CXR.
mimic/mimic_stage1_ckpt.pth	Stage 1 warm-up checkpoint.
mimic/mimic_stage2_ckpt.pth	Stage 2 classifier-guided supervised fine-tuning checkpoint.
mimic/mimic_stage3_ckpt.pth	Final Stage 3 REVA-PO checkpoint used for evaluation.

IU-Xray Checkpoints

File	Purpose
iuxray/mimic_stage1_ckpt.pth	MIMIC-CXR Stage 1 initialization used for IU-Xray Stage 1 training.
iuxray/iuxray_cls_ckpt.pth	Image classifier fine-tuned for IU-Xray.
iuxray/iuxray_stage1_ckpt.pth	Stage 1 warm-up checkpoint.
iuxray/iuxray_stage2_ckpt.pth	Stage 2 classifier-guided supervised fine-tuning checkpoint.
iuxray/iuxray_stage3_ckpt.pth	Final Stage 3 REVA-PO checkpoint used for evaluation.

Download

Install the Hugging Face Hub client:

pip install -U huggingface_hub

Download all checkpoints:

from huggingface_hub import snapshot_download

snapshot_download(
    repo_id="liguo12/REVA_PO_Weights",
    local_dir="REVA_PO_Weights",
)

Download only the final checkpoints and required auxiliary weights:

from huggingface_hub import hf_hub_download

files = [
    "iuxray/iuxray_cls_ckpt.pth",
    "iuxray/iuxray_stage2_ckpt.pth",
    "iuxray/iuxray_stage3_ckpt.pth",
    "mimic/chexbert.pth",
    "mimic/mimic_cls_ckpt.pth",
    "mimic/mimic_stage2_ckpt.pth",
    "mimic/mimic_stage3_ckpt.pth",
]

for filename in files:
    hf_hub_download(
        repo_id="liguo12/REVA_PO_Weights",
        filename=filename,
        local_dir="REVA_PO_Weights",
    )

Installation

Clone the code repository and install its requirements:

git clone https://github.com/LiGuo12/REVA_PO.git
cd REVA_PO
pip install -r requirements.txt

Download and prepare the datasets from the REVA-PO dataset repository.

IU-Xray Evaluation

Update train_configs/stage3/iuxray_stage3.yaml with local paths:

use_chexbert: true
chexbert_ckpt: /path/to/REVA_PO_Weights/mimic/chexbert.pth
pretrained_cls_ckp: /path/to/REVA_PO_Weights/iuxray/iuxray_cls_ckpt.pth
pretrained_stage2: /path/to/REVA_PO_Weights/iuxray/iuxray_stage2_ckpt.pth
pretrained_stage3: /path/to/REVA_PO_Weights/iuxray/iuxray_stage3_ckpt.pth
storage: /path/to/iuxray_dataset
ann_file: /path/to/iuxray_dataset/annotation_with_categories.json

Adjust world_size and the DeepSpeed batch size for the available hardware, then run:

deepspeed --num_gpus 8 train.py \
  --cfg-path train_configs/stage3/iuxray_stage3.yaml \
  --use_zero_optimizer \
  --deepspeed_config train_configs/stage3/zero_iuxray_stage3.json

MIMIC-CXR Evaluation

Update train_configs/stage3/mimic_stage3.yaml with local paths:

chexbert_ckpt: /path/to/REVA_PO_Weights/mimic/chexbert.pth
pretrained_cls_ckp: /path/to/REVA_PO_Weights/mimic/mimic_cls_ckpt.pth
pretrained_stage2: /path/to/REVA_PO_Weights/mimic/mimic_stage2_ckpt.pth
pretrained_stage3: /path/to/REVA_PO_Weights/mimic/mimic_stage3_ckpt.pth
storage: /path/to/mimic_dataset
ann_file: /path/to/mimic_dataset/mimic_with_categories_sampled_10k.json

Adjust world_size and the DeepSpeed batch size for the available hardware, then run:

deepspeed --num_gpus 8 train.py \
  --cfg-path train_configs/stage3/mimic_stage3.yaml \
  --use_zero_optimizer \
  --deepspeed_config train_configs/stage3/zero_mimic_stage3.json

See the GitHub README for Stage 1, Stage 2, classifier fine-tuning, and full training instructions.

Training Data

The checkpoints were trained and evaluated using:

• IU-Xray
• MIMIC-CXR-JPG
• Processed annotations in liguo12/REVA_PO_Datasets

MIMIC-CXR images require credentialed access through PhysioNet.

Intended Use

These checkpoints are intended for research on:

• Chest X-ray report generation
• Reinforcement learning for clinical language generation
• Medical vision-language models
• Evaluation of report quality and clinical accuracy
• Reproduction and extension of the REVA-PO experiments

They are not intended for direct clinical diagnosis, patient management, or deployment without independent medical and technical validation.

Limitations

The checkpoints inherit limitations from their training datasets, including institution-specific patterns, reporting bias, class imbalance, and annotation noise. Generated reports can contain omissions or incorrect findings. Results on IU-Xray and MIMIC-CXR do not establish performance on other institutions, patient groups, image devices, or clinical workflows.

Checkpoint Safety

The repository stores PyTorch .pth files, which use pickle-based serialization. Only load checkpoint files from sources you trust. The provided checkpoints should be used through the released REVA-PO code and compatible PyTorch and DeepSpeed versions.

License and Access Notes

The REVA-PO checkpoint files and other original materials released in this repository are licensed under the Creative Commons Attribution 4.0 International License. You may use, share, and adapt these materials for research or other purposes, provided that appropriate credit is given to the REVA-PO authors.

This license applies only to materials created and released by the REVA-PO authors. It does not replace the terms of the source datasets or upstream pretrained models:

• The REVA-PO source code is released under the license stated in the GitHub repository.
• MIMIC-CXR requires credentialed access and compliance with the PhysioNet data use requirements.
• Upstream pretrained models and auxiliary checkpoints may be subject to their own licenses.
• Users are responsible for checking and following all applicable terms before redistribution or deployment.

Citation

If you use these checkpoints or the REVA-PO method, please cite:

@misc{guo2026revapo,
  title        = {REVA-PO: Stabilizing Reinforcement Learning for Chest X-ray Report Generation},
  author       = {Li Guo and Anas M. Tahir and Z. Jane Wang},
  year         = {2026},
  eprint       = {2607.10147},
  archivePrefix = {arXiv},
  primaryClass = {cs.CV},
  url          = {https://arxiv.org/abs/2607.10147}
}