[Paper Review] DreaMR: Diffusion-driven Counterfactual Explanation for Functional MRI

Hasan A. Bedel, Tolga Çukur|arXiv (Cornell University)|Jul 18, 2023

Advanced Neuroimaging Techniques and Applications4 citations

TL;DR

DreaMR is a novel diffusion-driven counterfactual explanation method for functional MRI that generates high-fidelity, plausible counterfactual fMRI samples by leveraging a fractional multi-phase-distilled diffusion prior and a transformer-based architecture to model long-range spatiotemporal dependencies. It achieves superior specificity, fidelity, and sampling efficiency compared to state-of-the-art counterfactual methods in fMRI interpretation, enabling more reliable and interpretable deep learning analysis of brain responses.

ABSTRACT

Deep learning analyses have offered sensitivity leaps in detection of cognitive states from functional MRI (fMRI) measurements across the brain. Yet, as deep models perform hierarchical nonlinear transformations on their input, interpreting the association between brain responses and cognitive states is challenging. Among common explanation approaches for deep fMRI classifiers, attribution methods show poor specificity and perturbation methods show limited plausibility. While counterfactual generation promises to address these limitations, previous methods use variational or adversarial priors that yield suboptimal sample fidelity. Here, we introduce the first diffusion-driven counterfactual method, DreaMR, to enable fMRI interpretation with high specificity, plausibility and fidelity. DreaMR performs diffusion-based resampling of an input fMRI sample to alter the decision of a downstream classifier, and then computes the minimal difference between the original and counterfactual samples for explanation. Unlike conventional diffusion methods, DreaMR leverages a novel fractional multi-phase-distilled diffusion prior to improve sampling efficiency without compromising fidelity, and it employs a transformer architecture to account for long-range spatiotemporal context in fMRI scans. Comprehensive experiments on neuroimaging datasets demonstrate the superior specificity, fidelity and efficiency of DreaMR in sample generation over state-of-the-art counterfactual methods for fMRI interpretation.

Motivation & Objective

• To address the challenge of interpreting deep learning models in fMRI by generating counterfactual samples that are both minimally altered and highly plausible.
• To overcome the limitations of existing counterfactual methods that rely on variational or adversarial priors, which often yield low-fidelity or unstable samples.
• To improve sampling efficiency and fidelity in counterfactual generation for fMRI by introducing a novel fractional multi-phase-distilled diffusion prior (FMD).
• To enable high-specificity, high-plausibility explanations of deep fMRI classifiers by combining diffusion-based generation with classifier guidance and transformer modeling of spatiotemporal context.
• To establish a model-agnostic framework that can be applied to various deep learning architectures and future multi-modal or regression-based fMRI models.

Proposed method

• DreaMR employs a diffusion process that begins with forward diffusion of an input fMRI sample, adding noise over time steps.
• It uses a novel fractional multi-phase-distilled (FMD) diffusion prior that splits the diffusion process into F uniform fractions, each with dedicated denoising networks, enabling efficient and high-fidelity reverse sampling.
• Classifier guidance is injected at the end of each fraction using the conditional score of the downstream classifier to flip the predicted cognitive state.
• A transformer encoder is integrated into the diffusion model to capture long-range spatiotemporal dependencies across brain regions and time points in fMRI data.
• The method generates counterfactual samples by minimizing the L2 difference between the original and generated samples while ensuring the classifier's decision changes.
• The FMD prior enables faster sampling by distilling knowledge from multiple phases, reducing the number of reverse steps required without sacrificing sample quality.

Experimental results

Research questions

• RQ1 Can a diffusion-based generative prior produce counterfactual fMRI samples with higher fidelity and plausibility than variational or adversarial priors?
• RQ2 How can sampling efficiency be improved in diffusion-based counterfactual generation for high-dimensional fMRI data without compromising sample quality?
• RQ3 To what extent does incorporating a transformer architecture enhance the modeling of long-range spatiotemporal dependencies in fMRI counterfactuals?
• RQ4 How does the proposed fractional multi-phase distillation strategy affect the trade-off between sampling speed and generation fidelity in fMRI counterfactuals?
• RQ5 Can the method achieve superior specificity and plausibility in explanation compared to attribution and perturbation-based methods in deep fMRI classification?

Key findings

• DreaMR achieved a proximity score of 38.7 ± 1.9, specificity of 10.6 ± 0.8, and FID of 21.5 ± 0.4, outperforming ablated variants and state-of-the-art methods in counterfactual generation for fMRI.
• Removing the transformer component increased the proximity score to 47.5 ± 2.3 and specificity to 14.0 ± 0.9, demonstrating its critical role in modeling spatiotemporal context.
• Omitting Langevin sampling led to a 44.7 ± 2.1 proximity score and 13.0 ± 0.8 specificity, indicating that stochastic sampling improves generation quality.
• The multi-phase distillation strategy significantly enhanced performance, with F=4 and P=7 yielding the best FID (21.5 ± 0.4), while increasing fractions beyond 8 showed diminishing returns.
• The FMD prior with F=8 and P=7 achieved the lowest FID (21.5 ± 0.4), confirming that fractional distillation improves sampling efficiency and fidelity.
• DreaMR demonstrated superior performance across all metrics compared to ablated models, confirming the necessity of all components for optimal counterfactual explanation in fMRI.