FiRe: Fixed-Noise Refinement for Visual Counterfactual Explanations

Official PyTorch implementation of FiRe: Fixed-Noise Refinement for Visual Counterfactual Explanations, accepted at BMVC 2026.

FiRe generates localized visual counterfactual explanations by refining a fixed noisy state with classifier guidance and dynamically updated attribution masks, built on a one-step Pixel Mean Flow (pMF) backbone. The repository supports CelebA (128×128) and CelebA-HQ (256×256) experiments for the age and smile attributes, and CheXpert (512×512) for the pacemaker attribute.

Installation

Python 3.10 or newer is recommended. Install the PyTorch build appropriate for your CUDA or ROCm platform if the pinned build in requirements.txt is not suitable.

git clone git@github.com:Yan-Zen9/FiRe.git
cd FiRe

python -m venv .venv
source .venv/bin/activate
pip install --upgrade pip
pip install -r requirements.txt
pip install -e .

Quick start

1. Download pretrained weights

pMF backbones — all checkpoints are hosted on Hugging Face at DapaoZeng/FiRe:

hf download DapaoZeng/FiRe --local-dir checkpoints/pmf

File	Backbone	Dataset	Resolution
pmf-celeba128-age.pt	DiT-B/8	CelebA (young/old)	128×128
pmf-celeba128-smile.pt	DiT-B/8	CelebA (smile/non-smile)	128×128
pmf-celebahq256-age.pt	DiT-B/16	CelebA-HQ (young/old)	256×256
pmf-celebahq256-smile.pt	DiT-B/16	CelebA-HQ (smile/non-smile)	256×256
pmf-chexpert512-pacemaker.pt	DiT-B/32	CheXpert (pacemaker/no pacemaker)	512×512

Classifiers under explanation — for CelebA and CelebA-HQ we use the same pretrained classifiers as ACE:

• CelebA: DenseNet classifier from the DiME release on Hugging Face — guillaumejs2403/DiME.
• CelebA-HQ: download checkpoints_decision_densenet.tar.gz from the STEEX releases and use celebamaskhq/checkpoint.tar.
• CheXpert: our DenseNet-121 pacemaker classifier, hosted in the same Hugging Face repository as chexpert512-classifier.pt.

2. Inference: generate counterfactuals

The input directory must contain one folder per class (e.g. young/ and old/). A minimal CelebA age example:

python CFgenerating/generate_cf.py \
  --config configs/pMF-CelebA128-age.yaml \
  --checkpoint checkpoints/pmf/pmf-celeba128-age.pt \
  --classifier_checkpoint checkpoints/celeba_classifier.pth \
  --classifier_query_label 39 \
  --image_dir data/celeba/age/val \
  --output_dir outputs/counterfactuals/celeba_age \
  --noise_levels 0.4 \
  --solver_steps 1 \
  --guidance_inner_steps 15 \
  --early_stop_on_flip \
  --early_stop_target_prob 0.85 \
  --guidance_scale 0.02 \
  --adaptive_guidance \
  --lambda_cls 1.0 \
  --lambda_perc 0.20 \
  --lambda_tv 0.01 \
  --mask_attribution_method grad \
  --mask_objective target_only \
  --mask_topk_percent 8 \
  --mask_dilate_radius 2 \
  --mask_feather_width 3 \
  --dynamic_mask

Query labels follow the CelebA attribute indices: 39 (Young) for age, 31 (Smiling) for smile. Generated samples, panels, masks, and metadata are written beneath --output_dir. Ready-to-use SLURM examples are provided in scripts/cf_celeba_age.sh and scripts/cf_celeba_smile.sh.

3. Evaluation

Our evaluation protocol (FID, sFID, FVA, FS, MNAC, CD, COUT, flip rate) follows ACE: we use the evaluation code from the ACE repository together with the evaluation model weights linked there:

• CelebA oracle: guillaumejs2403/DiME on Hugging Face.
• CelebA-HQ oracle: checkpoints_oracle_attribute.tar.gz (celebamaskhq/checkpoint.tar) from the STEEX releases.
• FVA / FS backbones: VGGFace2 resnet50_ft and the SimSiam ResNet-50 (batch size 256).

4. Training

pMF backbone — train from scratch with one of the four provided configurations:

# single GPU
python scripts/train.py --config configs/pMF-CelebA128-age.yaml

# 8 GPUs
torchrun --standalone --nproc_per_node=8 \
  scripts/train.py --config configs/pMF-CelebA128-age.yaml

Add --resume /path/to/checkpoint.pt to resume. SLURM examples for all four settings are in scripts/train_*.sh; sampling quality of intermediate checkpoints can be checked with scripts/eval.py.

Classifiers — we do not retrain classifiers. To train your own, follow DiME (CelebA) or STEEX (CelebA-HQ DecisionDensenet) and pass the resulting checkpoint via --classifier_checkpoint.

Datasets

We use CelebA (128×128) and CelebA-HQ (256×256) for the age and smile attributes, and CheXpert (512×512) for the pacemaker attribute. CelebA / CelebA-HQ download and preparation follow ACE / DiME; CheXpert requires registration on the Stanford AIMI website.

License

License information will be added after the release terms for the project and its upstream components have been confirmed.