AXIAL — Scalable 2.5D Attention Framework for Interpretable 3D MRI Classification

AXIAL is a 2.5D deep learning framework that turns a 3D volume into a sequence of 2D slices, extracts features with a pretrained 2D CNN backbone, and learns a volumetric representation via slice-attention aggregation. The same attention mechanism yields voxel-level explanations aligned with model decisions.

This repository contains the original code implementation of the AXIAL framework proposed in:

AXIAL: Attention-based eXplainability for Interpretable Alzheimer's Localized Diagnosis using 2D CNNs on 3D MRI brain scans
Gabriele Lozupone et al. — [Paper]

The paper uses Alzheimer’s Disease diagnosis as a benchmark on ADNI, but the core contribution is a general, reusable approach for efficient 3D scan classification and explainability using 2D backbones.

<img src="./readme_files/summary.png">

Why this repo?

If you build models for high-dimensional data (3D medical imaging, scientific imaging, volumetric sensing), you often face a trade-off:

• 3D CNNs: strong but expensive + harder to interpret
• 2D slice pipelines: efficient but weak global reasoning + brittle explainability

AXIAL bridges this gap:

• Scalability: use fast pretrained 2D backbones (transfer learning friendly)
• 3D reasoning: learn slice importance and global volumetric features end-to-end
• Explainability: produce voxel-level attention maps and region-level metrics
• Reproducibility: standardized ADNI→BIDS + Clinica preprocessing pipeline

Key results (paper benchmark)

On ADNI (ADNI1 Complete 1Yr 1.5T), AXIAL reports:

• AD vs CN: ACC 0.856, MCC 0.712
• sMCI vs pMCI (prognosis): ACC 0.725, MCC 0.443

Additionally, the explainability pipeline consistently highlights clinically relevant regions (e.g., hippocampus, amygdala, parahippocampal, inferior lateral ventricles) with voxel-level precision.

Highlights

• 2.5D architecture: 2D backbone + attention aggregation for volumetric modeling
• End-to-end slice importance: learns which slices matter (no heuristics)
• Voxel-level explainability: attention-based 3D maps + GradCAM baseline
• Leakage-safe evaluation: subject-wise splitting for cross-validation
• Reproducible data pipeline: ADNI → BIDS conversion + Clinica preprocessing

Table of Contents

1. Installation
2. Data Preparation
   • ADNI Data Download
   • Clinical Data Download
   • Rename CSV Files
   • Install Clinica Software
   • Convert ADNI to BIDS
3. Preprocessing
   • Run Preprocessing Pipeline
4. Deep Learning approaches
   • Diagnosis Network proposed in the paper (Axial3D)
   • Training and test a model with 5-fold cross-validation
   • Attention-based Explainability
   • GradCAM-based Explainability
5. Acknowledgement
6. License

Abstract

This study presents an innovative method for Alzheimer's disease diagnosis using 3D MRI designed to enhance the explainability of model decisions. Our approach adopts a soft attention mechanism, enabling 2D CNNs to extract volumetric representations. At the same time, the importance of each slice in decision-making is learned, allowing the generation of a voxel-level attention map to produces an explainable MRI. To test our method and ensure the reproducibility of our results, we chose a standardized collection of MRI data from the Alzheimer's Disease Neuroimaging Initiative (ADNI). On this dataset, our method significantly outperforms state-of-the-art methods in (i) distinguishing AD from cognitive normal (CN) with an accuracy of 0.856 and Matthew's correlation coefficient (MCC) of 0.712, representing improvements of 2.4% and 5.3% respectively over the second-best, and (ii) in the prognostic task of discerning stable from progressive mild cognitive impairment (MCI) with an accuracy of 0.725 and MCC of 0.443, showing improvements of 10.2% and 20.5% respectively over the second-best. We achieved this prognostic result by adopting a double transfer learning strategy, which enhanced sensitivity to morphological changes and facilitated early-stage AD detection. With voxel-level precision, our method identified which specific areas are being paid attention to, identifying these predominant brain regions: the hippocampus, the amygdala, the parahippocampal, and the inferior lateral ventricles. All these areas are clinically associated with AD development. Furthermore, our approach consistently found the same AD-related areas across different cross-validation folds, proving its robustness and precision in highlighting areas that align closely with known pathological markers of the disease.

Framework overview

This repository contains code for preprocessing structural Magnetic Resonance Imaging (sMRI) data from the Alzheimer's Disease Neuroimaging Initiative (ADNI) dataset. The code converts the ADNI dataset into the Brain Imaging Data Structure (BIDS) format and applies preprocessing algorithms, including N4 bias field correction, MNI152 registration, and brain extraction. The data preparation is performed using the Clinica software.

Additionally, this repository provides the necessary code to train, validate, and test various deep learning models using PyTorch. Furthermore, the repository includes two explainability approaches one based on GradCAM and the other based on the attention mechanism.

1. 3D Models: These approaches utilizes a 2D backbone to extract feature maps from slices and attention mechanisms to enables capturing slice-level features and their spatial relationships.

2. 2D Models: This approach directly classifies each slice by attaching a classifier to the backbone. The final label for the entire image is determined using a majority voting approach based on the slice predictions.

3. Explainability: The first explainability approach proposed generates attention activation maps at the voxel level, highlighting the brain regions that are most important for the model's decision-making process. The second approach utilizes GradCAM to generate 3D heatmap. The 3D maps are then used to produce XAI metrics that helps to identify the most important brain regions for the model's decision-making process.

The repository aims to provide a comprehensive framework for sMRI preprocessing and deep learning analysis, enabling researchers to efficiently analyze ADNI data and develop advanced models for Alzheimer's disease detection and classification.

Installation

To use the code in this repository, follow these steps:

1. Clone the repository:

git clone https://github.com/GabrieleLozupone/AXIAL.git

2. Install the required dependencies:

pip install -r requirements.txt

Data Preparation

Before performing preprocessing on the ADNI dataset, follow the steps below to prepare the necessary data.

ADNI Data Download

1. Subscribe to the ADNI website at https://ida.loni.usc.edu/login.jsp.

2. Download the desired ADNI image collection. In the case of this work, the image collection name is "ADNI1 Complete 1Yr 1.5T".

Clinical Data Download

1. On the ADNI website, click on "Download" and then select "Study Data".

2. Choose "ALL" to download all available data.

3. In the "Tabular Data (CSV format)" section, select all the files and download them.

Rename CSV Files

Some CSV files in the clinical data may have a date at the end of their name. Remove the date from the file names to ensure compatibility with the preprocessing pipeline.

Install Clinica Software

Install the Clinica software by following the instructions provided at https://aramislab.paris.inria.fr/clinica/docs/public/latest/Converters/ADNI2BIDS/. Clinica is a powerful tool that facilitates the conversion of ADNI data to the BIDS structure.

Convert ADNI to BIDS

To convert the ADNI data to the BIDS structure, use the following command:

clinica convert adni-to-bids -m T1 DATASET_DIRECTORY CLINICAL_DATA_DIRECTORY BIDS_DIRECTORY

Replace DATASET_DIRECTORY with the path to the downloaded ADNI dataset, CLINICAL_DATA_DIRECTORY with the path to the downloaded clinical data, and BIDS_DIRECTORY with the desired output path for the BIDS-formatted dataset. The -m T1 option specifies that only MRI data with T1 weighting should be converted.

Preprocessing

This section describes the preprocessing steps for the sMRI data.

Run Preprocessing Pipeline

To run the preprocessing pipeline on the sMRI data, execute the following command:

python data_preprocessing.py --bids_path /path/to/bids-dataset --n_proc 10 --checkpoint checkpoint.txt

Replace /path/to/bids-dataset with the path to the BIDS-formatted dataset obtained from the data preparation steps. The preprocessing pipeline includes MNI152 registration, brain extraction, and bias field correction with the N4 algorithm. The preprocessed images will be stored in the same path as the original images.

The n_proc argument specifies the number of processes to be used for the preprocessing pipeline.

The checkpointargument specifies the path to the checkpoint file, which is used to keep track of the images that have already been preprocessed. This allows the preprocessing pipeline to be interrupted and resumed later.

Please