[Paper Review] NEURO-DRAM: a 3D recurrent visual attention model for interpretable neuroimaging classification

David Wood, James H. Cole+1 more|arXiv (Cornell University)|Oct 10, 2019

Machine Learning in Healthcare46 references4 citations

TL;DR

NEURO-DRAM is a 3D recurrent visual attention model that uses reinforcement learning to navigate volumetric MRI scans, selectively attending to the most diagnostically relevant brain regions for neuroimaging classification. It achieves state-of-the-art accuracy in Alzheimer’s disease prediction and mild cognitive impairment prognosis, while automatically identifying known radiological biomarkers without explicit supervision.

ABSTRACT

Deep learning is attracting significant interest in the neuroimaging community as a means to diagnose psychiatric and neurological disorders from structural magnetic resonance images. However, there is a tendency amongst researchers to adopt architectures optimized for traditional computer vision tasks, rather than design networks customized for neuroimaging data. We address this by introducing NEURO-DRAM, a 3D recurrent visual attention model tailored for neuroimaging classification. The model comprises an agent which, trained by reinforcement learning, learns to navigate through volumetric images, selectively attending to the most informative regions for a given task. When applied to Alzheimer's disease prediction, NEURODRAM achieves state-of-the-art classification accuracy on an out-of-sample dataset, significantly outperforming a baseline convolutional neural network. When further applied to the task of predicting which patients with mild cognitive impairment will be diagnosed with Alzheimer's disease within two years, the model achieves state-of-the-art accuracy with no additional training. Encouragingly, the agent learns, without explicit instruction, a search policy in agreement with standardized radiological hallmarks of Alzheimer's disease, suggesting a route to automated biomarker discovery for more poorly understood disorders.

Motivation & Objective

• Address the limitations of standard convolutional neural networks in neuroimaging, including overfitting, poor generalization, and lack of interpretability.
• Develop a model that can integrate non-imaging clinical and demographic data naturally to guide attention without relying on them for final classification.
• Improve model interpretability by visualizing attention trajectories to reveal which brain regions are most influential in classification decisions.
• Enable zero-shot transfer to out-of-sample data and new clinical tasks without retraining, enhancing generalization and clinical utility.
• Explore automated biomarker discovery for poorly understood neurological disorders by learning attention policies aligned with known pathological regions.

Proposed method

• Employ a reinforcement learning agent that navigates a 3D MRI volume in discrete steps, attending to small 3D patches at a time.
• Use a 3D convolutional neural network to extract features from each attended patch, feeding them into an LSTM-based recurrent controller.
• Train the agent using policy gradient reinforcement learning to maximize classification accuracy, with sparse rewards based on prediction performance.
• Incorporate non-imaging meta-data (e.g., age, sex, clinical scores) via a context network that modulates the agent’s attention policy.
• Visualize attention trajectories to identify the most informative anatomical regions, providing interpretability and potential biomarker discovery.
• Use a hierarchical architecture where the LSTM hidden state maintains context across steps, enabling long-range reasoning over the brain volume.

Experimental results

Research questions

• RQ1 Can a reinforcement learning-based 3D visual attention model achieve superior classification performance on neuroimaging data compared to standard 3D CNNs?
• RQ2 Can the model generalize to out-of-sample MRI data from different scanners without fine-tuning?
• RQ3 Does the agent learn to attend to clinically relevant brain regions associated with Alzheimer’s disease, such as the hippocampus and temporoparietal cortex?
• RQ4 Can the model leverage non-imaging clinical and demographic data to improve attention policy without directly using them in the final decision?
• RQ5 Can the attention trajectory serve as a proxy for automated biomarker discovery in neurodegenerative disorders?

Key findings

• NEURO-DRAM achieves state-of-the-art classification accuracy in predicting Alzheimer’s disease from structural MRI, significantly outperforming a baseline 3D convolutional neural network.
• The model generalizes effectively to out-of-sample data, showing no performance degradation when tested on an external dataset from a different scanner.
• When applied to predicting progression from mild cognitive impairment to Alzheimer’s disease within two years, NEURO-DRAM achieves state-of-the-art accuracy with no additional training.
• The agent learns a search policy that systematically visits known radiological hallmarks of Alzheimer’s disease, including the hippocampus, parahippocampal gyrus, lateral ventricles, and parietal cortex.
• The attention trajectories are interpretable and align with clinical knowledge, suggesting the model can guide automated biomarker discovery for less understood disorders.
• The integration of non-imaging meta-data improves attention guidance without compromising the model’s reliance on imaging features for final classification.