[Paper Review] AxonNet: A self-supervised Deep Neural Network for Intravoxel Structure Estimation from DW-MRI

Hanna Y. Ehrlich, Mariano Rivera|arXiv (Cornell University)|Mar 19, 2021

Advanced Neuroimaging Techniques and Applications4 citations

TL;DR

AxonNet proposes a self-supervised deep learning framework using a multilayer perceptron (MLP) to estimate intravoxel axonal structure from DW-MRI data without requiring labeled real data. It leverages a generative diffusion multitensor model for synthetic data generation and achieves state-of-the-art accuracy with significantly faster inference times—processing entire volumes in minutes versus hours for SOTA methods—demonstrated on both synthetic and real HARDI datasets including Stanford and local acquisitions.

ABSTRACT

We present a method for estimating intravoxel parameters from a DW-MRI based on deep learning techniques. We show that neural networks (DNNs) have the potential to extract information from diffusion-weighted signals to reconstruct cerebral tracts. We present two DNN models: one that estimates the axonal structure in the form of a voxel and the other to calculate the structure of the central voxel using the voxel neighborhood. Our methods are based on a proposed parameter representation suitable for the problem. Since it is practically impossible to have real tagged data for any acquisition protocol, we used a self-supervised strategy. Experiments with synthetic data and real data show that our approach is competitive, and the computational times show that our approach is faster than the SOTA methods, even if training times are considered. This computational advantage increases if we consider the prediction of multiple images with the same acquisition protocol.

Motivation & Objective

• To develop a self-supervised deep learning method for estimating intravoxel axonal structure from DW-MRI data without relying on real labeled data.
• To address the challenge of limited ground-truth data in DW-MRI analysis by using synthetic signals generated via a formal diffusion multitensor model.
• To improve computational efficiency for high-angular-resolution DW-MRI (HARDI/super-HARDI) by enabling fast inference across entire brain volumes.
• To evaluate performance across different acquisition protocols and spatial support regions (voxelwise vs. patchwise) using both synthetic and real data.
• To demonstrate that self-supervised training with synthetic data can yield competitive results compared to supervised or SOTA methods in fiber orientation estimation.

Proposed method

• Uses a self-supervised training strategy where synthetic DW-MRI signals are generated using a diffusion multitensor model with known parameters (eigenvalues, diffusion tensors, mixing coefficients).
• Employs a multilayer perceptron (MLP) as the neural network architecture to regress from the DW-MRI signal vector to the estimated axonal structure parameters (e.g., fiber orientations and volume fractions).
• Applies two variants: a voxel-wise model (VOX) and a neighborhood-based model (NBH) that uses a 3×3×3 patch of signals to improve estimation robustness.
• Implements normalization of signals (setting S₀ = 1) and enforces physical constraints such as non-negative mixing coefficients (α > 0.1) and ordered α values to improve training stability.
• Uses the Wasserstein distance (Earth Mover’s Distance) as the primary loss function to compare predicted and ground-truth fiber orientation distributions (FODs), enabling better distributional alignment than angular error alone.
• Trains the network on synthetic data generated under specific acquisition protocols (gradient directions and b-values), then deploys it directly on real data without fine-tuning.

Experimental results

Research questions

• RQ1 Can a self-supervised deep learning approach achieve competitive performance in intravoxel axonal structure estimation from DW-MRI without requiring real labeled data?
• RQ2 How does the performance of AxonNet compare to state-of-the-art methods (e.g., NNLS, CSD) in terms of accuracy and computational efficiency on both synthetic and real DW-MRI data?
• RQ3 Does incorporating neighborhood context (via 3×3×3 patches) improve estimation accuracy compared to voxelwise prediction?
• RQ4 To what extent does the model generalize across different acquisition protocols (e.g., Stanford HARDI vs. local 325-signal protocol) without retraining?
• RQ5 Can the model achieve real-time inference on full brain volumes while maintaining high accuracy, especially in comparison to computationally heavy SOTA methods?

Key findings

• AxonNet’s neighborhood model (NBH) achieved a Wasserstein distance of 0.81 on the Stanford HARDI dataset, outperforming NNLS (27.08) and CSD (6.08), indicating superior distributional alignment.
• On a local dataset with higher noise, NBH achieved a Wasserstein distance of 1.90, significantly better than CSD (8.74) and NNLS (114.32), demonstrating robustness to lower SNR.
• Prediction time for the entire Stanford HARDI dataset was 0.8 minutes (48 seconds) using the VOX model, compared to hours for CSD and NNLS, highlighting a major computational advantage.
• The VOX model trained in 1.24 seconds on Stanford data and 1.24 seconds on the local dataset, while the NBH model required 1.91 and 3.34 seconds respectively, showing efficient training even on complex data.
• Visual inspection confirmed that AxonNet predictions closely matched ground truth in both high- and low-error cases, with accurate peak detection and reduced artifacts compared to CSD and NNLS.
• The method demonstrated strong generalization across different acquisition protocols, including varying b-values, gradient counts (150 vs. 64), and signal repetition, without retraining.