# ME-CapsNet: A Multi-Enhanced Capsule Networks with Routing Mechanism

Jerrin Bright2, Suryaprakash Rajkumar2 and Arockia Selvakumar Arockia Doss1
School of Mechanical Engineering, Vellore Institute of Technology, Chennai,
India
<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>2These Authors Contributed Equally To This Work

###### Abstract

Convolutional Neural Networks need the construction of informative features,
which are determined by channel-wise and spatial-wise information at the
network’s layers. In this research, we focus on bringing in a novel solution
that uses sophisticated optimization for enhancing both the spatial and
channel components inside each layer’s receptive field. Capsule Networks were
used to understand the spatial association between features in the feature
map. Standalone capsule networks have shown good results on comparatively
simple datasets than on complex datasets as a result of the inordinate amount
of feature information. Thus, to tackle this issue, we have proposed ME-
CapsNet by introducing deeper convolutional layers to extract important
features before passing through modules of capsule layers strategically to
improve the performance of the network significantly. The deeper convolutional
layer includes blocks of Squeeze-Excitation networks which use a stochastic
sampling approach for progressively reducing the spatial size thereby
dynamically recalibrating the channels by reconstructing their
interdependencies without much loss of important feature information.
Extensive experimentation was done using commonly used datasets demonstrating
the efficiency of the proposed ME-CapsNet, which clearly outperforms various
research works by achieving higher accuracy with minimal model complexity in
complex datasets.

###### Index Terms:

Deep Learning, Capsule Network, Squeeze-Excitation, Stochastic Pooling

##  I Introduction

Convolutional Neural Network (CNN) is a cutting-edge image classification [1,
2, 3] approach that excels in detecting characteristics but struggles to grasp
the spatial connection between them, such as their direction, perspective, and
size. For scene understanding or any vision-based strategy, improving
categorization accuracy is critical. In addition, transmitting information
from one layer to the next (Routing) is less efficient in CNNs. This may be
improved by identifying important characteristics using a variety of
representations and building well-organized routing strategies for
transferring data between layers. Investigation of many elements of network
architecture (spatial, channel) is now a very prominent research subject in
CNN, and it may greatly improve the network’s understanding ability [4].

Figure 1: Comparing various Capsule Network variations based on model
complexity and test accuracy using CIFAR10 dataset.

Capsule Networks (CapsNet) where introduced by Sabour et al. [5] to solve the
common issues with CNNs (rotational invariance and failure to capture spatial
hierarchical information) by fetching most important feature information and
spatial relationship efficiently. Instead of pooling operations, dynamic
routing was adapted to prevent losing important features. Also, capsules where
used instead of neurons thereby passing the spatial relationship from one
layer to next efficiently. But, CapsNet tends to fail as a result of more
feature information in complex datasets. Thus, adding of deeper convolutional
layers to the CapsNet in a strategic way is one approach to solve this issue.

Thus, we are proposing a novel network called Multi-Enhanced Capsule Networks
(ME-CapsNet) which finds robust features, obtaining good channel and spatial
relationships between layers. Squeeze-Excitation network or SENets [6] with
very little computation difference, models the interconnection of channels by
adaptively recalibrating the channel features. Once the important features are
learnt, the likelihood and the properties of each and every feature are
estimated by passing it to the CapsNet module which inturn uses EM Routing to
prevent feature losses. EM Routing [7] is used to group capsules with the
parent capsule using the closeness in the proximity of the corresponding
estimated votes on the pose matrix. Also, reconstruction is included in the
proposed work, unlike the originally proposed CapsNet with EM Routing.

Figure 2: Overall architecture of ME-CapsNet

Some contributions of this research work are mentioned here as follows:

  * •

A novel network which multi enhances the features, that is, both spatial and
channel-wise information at each layer with deep convolutional layers (SENets)
and CapsNet;

  * •

Improvements in the Squeeze phase of the originally proposed SENets by using
Stochastic Spatial Sampling Pooling (S3P) to reduce feature information
losses, computation overload and tranining time when compared to the
originally used average pooling operation.

  * •

Exploring the capabilities of the proposed system by fusing it with several
state-of-the-art base models including ResNet, VGG, Inception and MobileNet
and analyzing the performance of the proposed network by calibrating with
different parameters.

The rest of the research paper is systematized as follows: Section II consists
of the related research works relating to various networks pertaining to
SENets and CapsNets followed by Section III detailing the proposed
architecture along with its modules. Then in Section IV, various
experimentation is done testing by tweaking parameters and in Section V the
proposed architecture is tested with various state-of-the-art datasets and
networks. Also, few predominantly efficient networks were handpicked and
tested alongside our network. Finally, in Section VI the proposed architecture
is summarized with its advantages and is followed by future works premeditated
to increase the robustness of the currently proposed system.

##  II Related Works

Since the introduction of deep convolutional networks [8] significant
improvement in the deep learning regime has been witnessed. Stacking of layers
as a result of deep networks significantly improved the performance but also
introduced vanishing gradient problems [9]. Also, using operating layers like
max pooling introduced loss of spatial information. CapsNet attempts to solve
these issues by nesting layers instead of stacking them. These recent
techniques in deep learning, added fuel to research in this field drastically.

CapsNet are considered a great leap into the neural networks considering its
significance and scope of improvement over the traditional techniques. Some
efficient research works on CapsNet are explained here. Mazzia et al. [10]
proposed a non-iterative routing technique that exploits self-attention
mechanisms. Deliege et al. [11] introduced a counterpart for DigitCaps layer
using a custom Hit-or-Miss layer. Huang et al. [12] proposed addition of two
attention layers to the original CapsNet, to improve the overall performance
of the network.

SENets are predominantly used recently considering the significant improvement
in performance with minimal amount of computation rise. A lot of customization
to SENets have also been done, including Yue Cao et al. [13] performing a
unification of NLnet and SENet resulting in a Global Context (GC) network. A
lot of SENets applications has been established including some works by Qiu et
al. [14], Shunjun Wei et al. [15], and Han et al. [16] where fusion of SENets
were done with various architectures of CNN for fish image classifications,
PRI modulation recognition, sea ice image sensing respectively.

Yang et al. [17] proposes an approach to enhance the feature extraction
capability of CapsNet using a modified version of ResNet called Res2Net and
SENet in the base layers to enhance the performance of the CapsNet. This is
the closest work in relavance to our work. We have strategically employed the
SE blocks in between various layers and used a more powerful variant of
pooling to improve the performance of the squeeze operation.

##  III Methodology

In this section, the different components equipped by our work to complete the
robust architecture is discussed extensively. Our architecture predominantly
has three blocks to multi-enhance the feature relationship of CNN: the SE
Primary Capsule (SE-PC) block, SE Classification Capsule (SE-CC) block and the
Mask by Layer (MbL) block. That is, an input image in form of tensor (_U_) is
fed to our network, which first converts into a tensor upon transformation
($U_{tr}$). Then, $U_{tr}$ is sent through the the SE-PC block wherein Squeeze
and Excitation operation is done which results in $U_{tr}^{o}$ which has
recalibrated weights based on channel dependence. $U_{tr}^{o}$ is then passed
through the primary capsule layer. The Squeeze and Excitation operation is
done three times continuously to get the improved calibrated weights for the
corresponding channels. The Squeeze operation is done using S3P [18] which
will be explained in Section III-A. Then the output of the primary layer of
the SE-PC is fed as input to the SE-CC wherein features are highlighted using
SENets first and then is sent to the classification module. Then, we designed
the MbL block inspired by the works for Huang et al. [19] for reconstructing
the output capsules from the probabilty vectors obtained from the SE-CC block.
Advantages of including the reconstruction block include regularization and
ability to regenerate data.

###  III-A Squeeze-Excitation Network

SENets directly influence the performance of the network and its layers
through adaptive adjusting of the weights obtained from each and every feature
map in the channels of the convolutional block by adding supplementary
parameters. This interfacing with weights radically increases the sensitivity
to beneficial features in the map which will be exploited in succeeding
transformations.

Figure 3: Schema of the input (SE-ResNet) architecture for Capsule Network

Generally, in the squeeze phase, Global Average Pooling (GAP) or Maximum
Pooling (GMP) techniques are equipped, each has its limitations and
advantages. GAP (represented as $z_{avg}$ in equation 1 results in smooth
average pixels and doesn’t preserve pixel information. GAP takes into
consideration all the features, thus it can result in a biased feature map.

$z_{avg}=\frac{1}{H\cdot W}\sum_{L=1}^{H}{\sum_{j=1}^{W}{u_{c}(i,j)}}$ (1)

GMP (represented as $z_{max}$ in equation 2) extracts low-level features
(edges, textures) better than GAP but captures a lot of noise and is
independent of neighboring pixels.

$z_{max}=max_{i,j}\cdot{u_{c}(i,j)}$ (2)

In this work, we will be using a pooling approach proposed by Zhai et al.
[18]. A two step process including GMP ($z_{max}$) and stochastic down
sampling ($\mathcal{D}_{g}^{s}(.)$) is done on the feature map to learn better
abstract representations in a larger receptive field. Considering it is
stochastic, different feature map is obtained for the same training sample.
This means it performs a virtual data augmentation [20]. Equation 3 defines
the S3P pooling, where _s_ represents strides and _g_ represents the grid
size.

$z_{s3p}=\mathcal{D}_{g}^{s}\cdot(z_{max})$ (3)

Now, the output of the squeeze phase ($z_{m}ix$) is passed to the excitation
phase for adaptive scaling of weights in the respective channel. Excitation is
done using the Multi-Layer Perceptron (MLP) approach. That is, two fully
connected layers are used with weights $W_{1}$ and $W_{2}$. The first fully
connected layer is connected to ReLU activation [21] followed by another fully
connected layer, which in turn is connected to the Sigmoid activation
function. This can be represented as shown in equation 4. While training,
$W_{1}$ and $W_{2}$ are learned.

$s=\sigma\cdot(W_{2}\cdot\delta(W_{1}\cdot z_{mix}))$ (4)

In other words, channel-wise dependence is found inside the excitation phase
by learning the weights, $W_{1}$ and $W_{2}$ using a compression ratio to
enable efficient learning.

Figure 3 shows the schema of the SENet with the Residual Network (ResNet)
backbone leveraged in this work. In other words, the output of figure 3 is the
input of the primary capsule layer. A very similar architecture without
ResNet, leveraging the output of capsule layers are used for the SENets
equipped before the classification capsule layer.

###  III-B Capsule Network

Appropriate activation can be given if the likeliness of the feature and the
feature information is passed via each and every neuron. Thus, these neurons
are termed capsules and they output a vector (called an activity vector)
instead of single feature information. Knowing the likeliness of a feature, we
can easily extrapolate different feature variants in a capsule, thereby
reducing the training data. The product of the pose matrix ($M_{i}$)
represented by a gaussian distribution with the transformation ($W_{ij}$) is
computed to calculate the vote ($v_{ij}$).

$v_{ij}=M_{i}\cdot W_{ij}$ (5)

In equation 5, _i_ corresponds to the parent capsule and _j_ to the current
capsule. The transformation is learned through back-propagation and a cost
function.

###  III-C Routing Mechanism

Transferring information from one layer to another is called Routing. ReLU is
used as the routing mechanism in fully convolutional neural networks.

$y_{i}=ReLU(\sum_{j}W_{ji}\cdot x_{j}+b_{i})$ (6)

Instead of ReLU, a squashing function is used to convert small vector
magnitudes to zero and big vector magnitudes to unit vectors. Grouping of
capsules using the EM clustering technique is called EM Routing. Parent
capsules are predicted in the pose matrix in terms of ‘votes’. It is based on
proximity, so even when the viewpoints change, the transformation matrix
remains constant. Assignment probabilities ($r_{ij}$) estimate the runtime
connection with respect to the parent and other capsules. The cost of all low-
level capsules is shown in equation 7.

$J_{j}^{h}=(ln(\sigma_{j}^{h})+k)\sum_{i}r_{ij}$ (7)

Activation ($A_{j}$) in a particular capsule can be estimated using the
inverse temperature parameter ($\lambda$) and parameters iteratively estimated
using EM Routing by equation 8.

$A_{j}=sigmoid(\lambda(b_{j}-\sum_{h}J_{j}^{h}))$ (8)

The data points fit into a gaussian model and E and M steps are triggered
alternatively iterating for ‘ _t_ ’ steps. The Gaussian model is updated using
the M step based on the $r_{ij}$ initialized (distributed uniformly) which
will then be reshaped to form the pose matrix of the parent capsule. The
$r_{ij}$ is recomputed using the E-step for every data point.

Upon the completion of EM Routing, the capsules with best activation is
selected to make the appropriate predictions using the MLP approach used in
the excitation phase of the SE Blocks of ME-CapsNet.

###  III-D Reconstruction

The output of the classification capsule is used for reconstruction of the
input image by passing through a decoder. The output is a vector which
consists the probability vector squashed to get values between 0 to 1. The
idea behind the used non-linear squashing function is to compress small
vectors to zero and long vectors to one. The decoder network consists of two
fully connected layers with ReLU activation and a sigmoid function. Mean
Square Error (MSE) is used for estimating the reconstruction loss between true
($y$) and predicted ($y^{0}$) labels as shown in equation 9.

$L(y,y^{0})=\frac{1}{N}\sum^{N}_{i=0}(y-y^{0})$ (9)

|  |  |
---|---|---|---
(a) MNIST Training Results | (b) MNIST Loss Results | (c) FMNIST Training Results | (d) FMNIST Loss Results
|  |  |
(e) CIFAR10 Training Results | (f) CIFAR10 Loss Results | (g) KMNIST Training Results | (h) KMNIST Loss Results
Figure 4: Training and Validation results of accuracy and overall loss using
various state-of-the-art datasets (MNIST, FashionMNIST- represented as FMNIST,
CIFAR10, KMNIST)

##  IV Experimentation

In this section, experiments are steered from the architecture proposed in the
previous section using MNIST, FashionMNIST and CIFAR-10 datasets, and the
performance is evaluated with respect to other networks, previous effective
research works and parameter tuning.

Implementation of ME-CapsNet is done using several tuned parameters for
efficient training and some important parameters are mentioned below. We have
built ME-CapsNet using Keras and tensorflow framework. Tensorboard
visualization toolkit was used for analyzing the performance of the proposed
network. The network was compiled using the NVIDIA GTX 1650 GPU. The dataset
was publicly available for the training of the model. For the training
procedure, 20 epochs were done uniformly for all datasets, a batch size of
128, Adam optimizer [22], squashing constant of 0.5 and decay factor of 0.9.

The section has been split into three based on the target calibration. In
section IV-B, the parameters of SENet is tweaked and the resulting top-1% and
top-10% error rates are obtained. This is followed by section IV-A, where the
capsule parameters are calibrated using various dimensions and quantity.

###  IV-A Capsule parameters

The number of capsules affects the overall accuracy and performance of the
proposed ME-CapsNet. In table I, different capsule dimensions are selected and
the influence of capsule dimension on the overall accuracy of the network was
tested. We have got the best accuracy when testing ME-CapsNet with dimension
of capsule set to be 8. This testing has been done using CIFAR10 dataset.

TABLE I: Capsule dimension

Dimension | Accuracy
---|---
8 | 89.84%
12 | 89.12%
16 | 88.57%

In table II, we tested our architecture with three variations in the total
number of capsules and the respective accuracy has been logged corresponding
to it.Best results where obtained when ME-CapsNet was tested with 5 capsules.
This testing has been done using CIFAR10 dataset and the dimension of capsule
set to be 8.

TABLE II: No. of capsules

Number | Accuracy
---|---
5 | 88.79%
8 | 89.84%
10 | 86.44%

###  IV-B SE Network parameters

As discussed in previous sections, SENets are split into Squeeze and
Excitation operation. The squeeze operation consists of a pooling technique
which includes max pooling, average pooling and the proposed S3P. All the
pooling techniques are extensively studied in section III-A. Similarly,
excitation operation is predominantly carried by the activation function like
Sigmoid, ReLU, LeakyReLU and TanH.

TABLE III: Comparison with SE parameters.

| top-1% | top -5%
---|---|---
Squeeze - Max | 22.57 | 6.09
Squeeze - Avg | 22.28 | 6.03
Squeeze - S3P (Ours) | 21.78 | 5.83
Excitation - Sigmoid | 22.28 | 6.03
Excitation - ReLU | 23.47 | 6.98
Excitation - LeakyReLU | 23.22 | 6.91
Excitation - TanH | 23.00 | 6.38

Sigmoid is an activation function that adds non-linearity and ranges between 0
to 1. Similarly, TanH is an activation function very similar to Sigmoid’s
properties, but ranges from -1 to 1. ReLU is a linear function which outputs
the input if positive and 0 if negative. Meanwhile, LeakyReLU works similar to
ReLU for positive inputs, but outputs a very small positive value when the
input tends to be negative.

##  V Results

In this section, ME-CapsNet equipped with the best parameters obtained from
the experimentation done in Section IV is compared alongside different
backbone state-of-the-art models (Section V-B) and novel research works
pertaining to CapsNet (Section V-C). Also details on the used datasets used
for testing the propsed architecture is explained in V-A.

###  V-A Datasets

Some commonly used datasets used to experiment the performance of ME-CapsNet
include MNIST [23], FashionMNIST [24], CIFAR10 [25] and KMNIST [26]. The
performance of ME-CapsNet when tested alongside these datasets can be
visualized in 4

MNIST and FashionMNIST both has 10 classes with 70,000 images of size
$28\times 28\times 1$.CIFAR10 has 10 classes with 60,000 images of size
$32\times 32\times 3$. KMNIST or Kuzushiji-MNIST dataset has 10 classes is a
drop-in replacement of MNIST and FashionMNIST, consisting of 70,000 images of
size $28\times 28\times 1$ just like those datasets. For MNIST, KMNIST and
FashionMNIST datasets 50,000 images where used for training and 10,000 images
for testing. For CIFAR10 dataset, 60,000 images where used for training and
10,000 for testing.

###  V-B Comparing various backbones

In this subsection, various backbone architectures including residual network
with various levels of layers, inception network and VGG network is fused as
input layer for ME-CapsNet and the respective test accuracy is logged and
compared systematically for four different datasets.

TABLE IV: Comparison of various backbone architectures

| MNIST | FashionMNIST | CIFAR10 | KMNIST
---|---|---|---|---
ResNet-50 | 99.57 | 94.85 | 88.87 | 96.75
ResNet-101 | 99.83 | 95.63 | 89.37 | 97.84
ResNet-152 | 99.78 | 95.91 | 89.84 | 97.11
VGG-16 [27] | 98.54 | 95.11 | 87.72 | 96.49
Inception [28] | 98.63 | 95.38 | 88.98 | 97.34

###  V-C Comparing architectures

In this subsection, some state-of-the-art research works aligning to CapsNet
is studied in terms of accuracy, number of parameters and the ability to
reconstruct. Table V clearly shows our network, ME-CapsNet outperforms the
other research works significantly with comparatively decent number of
parameters using CIFAR10 dataset.

TABLE V: Comparison with various proposed network.

| Accuracy | # of Params | Recon.
---|---|---|---
HitNet [11] | 73.30% | 8.89M | Yes
MS-CapsNet [29] | 75.70% | 11.20M | Yes
CapsNet Baseline | 79.24% | 11.98M | No
DeeperCaps [30] | 81.29% | 5.81M | Yes
DCNet [31] | 82.63% | 11.88M | Yes
DA-CapsNet [12] | 85.47% | 7M | Yes
Cv-CapsNet++ [32] | 86.70% | 2.69M | No
AR-CapsNet [33] | 88.94% | 9.60M | Yes
Sabour et. al. [5] | 89.40% | 14.36M | Yes
DCNet++ [31] | 89.71% | 13.4M | Yes
ME-CapsNet (Ours) | 89.84% | 7.24M | Yes

ME-CapsNet easily outperforms the research works on CapsNet in CIFAR10
datasets, thus demonstrating its ability to capture important features and
route without feature loss. Also, 1.23% better performance in terms of
accuracy was estimated using our model with FashionMNIST dataset. Though the
difference in accuracy between the proposed approach and the previous best
approach is only 0.13%, computationally our approach tends to be very
computationally cheaper. This is because of the less computational increase
when using deep convolutional layers via SENets in between the CapsNet layers
strategically with well-defined parameters.

##  VI Conclusion and Future Work

In this research, we have proposed a novel deep learning architecture that
simultaneously enhances the spatial and channel-wise relationship of features
with minimum computation. This multi-enhancement based network has been coined
ME-CapsNet which uses a stochastic pooling approach to strategically improve
the performance of the SENets used in CapsNet. The originally proposed CapsNet
was used as the base network and SENets have been fused strategically before
and in between capsule layers. The squeeze operation proposed in the original
SENets where altered in this work using S3P, a stochastic sampling-based
pooling approach which in turn significantly improved the performance of the
SENets used in the proposed ME-CapsNet. The architecture was evaluated with
state-of-the-art datasets and comparatively better results in terms of
accuracy and computational cost were obtained upon testing with the proposed
model.

Future study will focus on evaluating ME-CapsNet on broad, complex datasets.
Though the aim of the originally proposed CapsNet was to prevent the use of
pooling operations to ward off the loss of information, ME-CapsNet like many
other networks uses the advantages of pooling operators to enhance the
performance of CapsNets especially for complex datasets. One right research
direction to proceed from here would be to look for alternatives of pooling
operators for increasing the performance especially for complex datasets.
Also, adding attention to routing mechanism can be added to our network to
improve its performance. Computational time can be further reduced by
replacing the fully connected layers with capsule layers.

## References

  * [1] A. Krizhevsky, I. Sutskever, and G. E. Hinton, “Imagenet classification with deep convolutional neural networks,” Commun. ACM, vol. 60, p. 84–90, may 2017.
  * [2] M. D. Zeiler and R. Fergus, “Visualizing and understanding convolutional networks,” in Computer Vision – ECCV 2014 (D. Fleet, T. Pajdla, B. Schiele, and T. Tuytelaars, eds.), (Cham), pp. 818–833, Springer International Publishing, 2014.
  * [3] R. Mukhometzianov and J. Carrillo, “Capsnet comparative performance evaluation for image classification,” CoRR, vol. abs/1805.11195, 2018.
  * [4] X. Jiang, Y. Wang, W. Liu, S. Li, and J. Liu, “Capsnet, cnn, fcn: Comparative performance evaluation for image classification,” International Journal of Machine Learning and Computing, vol. 9, pp. 840–848, 12 2019.
  * [5] S. Sabour, N. Frosst, and G. E. Hinton, “Dynamic routing between capsules,” CoRR, vol. abs/1710.09829, 2017.
  * [6] J. Hu, L. Shen, and G. Sun, “Squeeze-and-excitation networks,” in 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 7132–7141, 2018.
  * [7] G. E. Hinton, S. Sabour, and N. Frosst, “Matrix capsules with EM routing,” in International Conference on Learning Representations, 2018.
  * [8] A. Krizhevsky, I. Sutskever, and G. E. Hinton, “Imagenet classification with deep convolutional neural networks,” Commun. ACM, vol. 60, p. 84–90, may 2017.
  * [9] G. Ananthi and A. Arockia Selvakumar, Review on Deep Learning Algorithms in Medical Devices, pp. 1–27.  Cham: Springer International Publishing, 2020.
  * [10] V. Mazzia, F. Salvetti, and M. Chiaberge, “Efficient-capsnet: Capsule network with self-attention routing,” CoRR, vol. abs/2101.12491, 2021.
  * [11] A. Deliège, A. Cioppa, and M. V. Droogenbroeck, “Hitnet: a neural network with capsules embedded in a hit-or-miss layer, extended with hybrid data augmentation and ghost capsules,” CoRR, vol. abs/1806.06519, 2018\.
  * [12] W. Huang and F. Zhou, “Da-capsnet: dual attention mechanism capsule network,” Sci Rep, 2020.
  * [13] Y. Cao, J. Xu, S. Lin, F. Wei, and H. Hu, “Gcnet: Non-local networks meet squeeze-excitation networks and beyond,” in 2019 IEEE/CVF International Conference on Computer Vision Workshop (ICCVW), pp. 1971–1980, 2019.
  * [14] C. Qiu, S. Zhang, C. Wang, Z. Yu, H. Zheng, and B. Zheng, “Improving transfer learning and squeeze- and-excitation networks for small-scale fine-grained fish image classification,” IEEE Access, vol. 6, pp. 78503–78512, 2018\.
  * [15] S. Wei, Q. Qu, Y. Wu, M. Wang, and J. Shi, “Pri modulation recognition based on squeeze-and-excitation networks,” IEEE Communications Letters, vol. 24, no. 5, pp. 1047–1051, 2020.
  * [16] Y. Han, C. Wei, R. Zhou, Z. Hong, Y. Zhang, and S. Yang, “Combining 3d-cnn and squeeze-and-excitation networks for remote sensing sea ice image classification,” Mathematical Problems in Engineering, vol. 2020, pp. 364–375, 4 2020.
  * [17] S. Yang, F. Lee, R. Miao, J. Cai, L. Chen, W. Yao, K. Kotani, and Q. Chen, “Rs-capsnet: An advanced capsule network,” IEEE Access, vol. 8, pp. 85007–85018, 2020.
  * [18] S. Zhai, H. Wu, A. Kumar, Y. Cheng, Y. Lu, Z. Zhang, and R. Feris, “S3pool: Pooling with stochastic spatial sampling,” in Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), July 2017.
  * [19] G. Huang, Z. Liu, and K. Q. Weinberger, “Densely connected convolutional networks,” CoRR, vol. abs/1608.06993, 2016.
  * [20] P. Simard, D. Steinkraus, and J. Platt, “Best practices for convolutional neural networks applied to visual document analysis,” in Seventh International Conference on Document Analysis and Recognition, 2003. Proceedings., pp. 958–963, 2003.
  * [21] V. Nair and G. E. Hinton, “Rectified linear units improve restricted boltzmann machines,” in Proceedings of the 27th International Conference on International Conference on Machine Learning, ICML’10, (Madison, WI, USA), p. 807–814, Omnipress, 2010.
  * [22] D. P. Kingma and J. Ba, “Adam: A method for stochastic optimization,” International Conference for Learning Representations, 2014.
  * [23] L. Deng, “The mnist database of handwritten digit images for machine learning research,” IEEE Signal Processing Magazine, vol. 29, no. 6, pp. 141–142, 2012.
  * [24] H. Xiao, K. Rasul, and R. Vollgraf, “Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms,” CoRR, vol. abs/1708.07747, 2017\.
  * [25] A. Krizhevsky, “Learning multiple layers of features from tiny images,” Citeseer, Tech. Rep, pp. 32–33, 2009.
  * [26] T. Clanuwat, M. Bober-Irizar, A. Kitamoto, A. Lamb, K. Yamamoto, and D. Ha, “Deep learning for classical japanese literature,” CoRR, vol. abs/1812.01718, 2018.
  * [27] K. Simonyan and A. Zisserman, “Very deep convolutional networks for large-scale image recognition,” arXiv 1409.1556, 09 2014.
  * [28] C. Szegedy, W. Liu, Y. Jia, P. Sermanet, S. E. Reed, D. Anguelov, D. Erhan, V. Vanhoucke, and A. Rabinovich, “Going deeper with convolutions,” CoRR, vol. abs/1409.4842, 2014.
  * [29] C. Xiang, L. Zhang, Y. Tang, W. Zou, and C. Xu, “Ms-capsnet: A novel multi-scale capsule network,” IEEE Signal Processing Letters, vol. 25, no. 12, pp. 1850–1854, 2018.
  * [30] J. Rajasegaran, V. Jayasundara, S. Jayasekara, H. Jayasekara, S. Seneviratne, and R. Rodrigo, “Deepcaps: Going deeper with capsule networks,” CoRR, 2019\.
  * [31] S. S. R. Phaye, A. Sikka, A. Dhall, and D. R. Bathula, “Dense and diverse capsule networks: Making the capsules learn better,” CoRR, vol. abs/1805.04001, 2018.
  * [32] X. Cheng, J. He, J. He, and H. Xu, “Cv-capsnet: Complex-valued capsule network,” IEEE Access, vol. 7, pp. 85492–85499, 2019.
  * [33] J. Choi, H. Seo, S. Im, and M. Kang, “Attention routing between capsules,” CoRR, vol. abs/1907.01750, 2019.


$\displaystyle=\int
x^{2}e^{-2\pi\frac{tl^{2}x^{2}}{t^{2}+q^{2}}}(2t)^{1/2}|l|(t^{2}+q^{2})^{-1/2}dx$
$\displaystyle=\int y^{2}(2t)^{-1}|l|^{-2}(t^{2}+q^{2})e^{-\pi y^{2}}dy,$
$\displaystyle\text{and \quad}\mathcal{V}\bigl{(}v_{t}\bigr{)}$
$\displaystyle=\int x^{2}e^{-2\pi tx^{2}}(2t)^{1/2}dx$ $\displaystyle=\int
y^{2}(2t)^{-1}e^{-\pi y^{2}}dy.$

We get then

$\displaystyle\mathcal{V}\bigl{(}\mathcal{M}_{p,l,q}v_{t}\bigr{)}\mathcal{V}\bigl{(}v_{t}\bigr{)}$
$\displaystyle=(2t)^{-2}|l|^{-2}(t^{2}+q^{2})\left(\int y^{2}e^{-\pi
y^{2}}dy\right)^{2}$
$\displaystyle=(2t)^{-2}|l|^{-2}(t^{2}+q^{2})\Bigl{(}\frac{1}{\pi^{3/2}}\Gamma(3/2)\Bigr{)}^{2}$
$\displaystyle=|l|^{-2}(1+q^{2}t^{-2})\frac{\pi}{4^{2}\pi^{3}}=|l|^{-2}(1+q^{2}t^{-2})\frac{1}{(4\pi)^{2}},$

and thus

(4.2.23)
$\sqrt{\mathcal{V}\bigl{(}\mathcal{M}_{p,l,q}v_{t}\bigr{)}\mathcal{V}\bigl{(}v_{t}\bigr{)}}={{\frac{1}{4\pi|l|}\sqrt{1+q^{2}t^{-2}}}\begin{matrix}\\\
\longrightarrow\\\ \hfill{\scriptstyle
t\rightarrow+{\infty}}\end{matrix}\frac{1}{4\pi|l|}}.$

####  4.2.3. Warm-up in one dimension

As a warm-up, let us check thoroughly the one-dimensional case. We start with
a symplectic matrix in $\textsl{Sp}(2,\mathbb{R})=\textsl{Sl}(2,\mathbb{R})$,
say

$\Xi=\begin{pmatrix}\alpha&\beta\\\
\gamma&\delta\end{pmatrix},\quad\alpha\delta-\beta\gamma=1.$

$\bullet$ We assume first that $\beta\not=0$ and we find that, using the
Notation (2.2.40)

$\Xi=\Lambda_{p,l,q},\quad\alpha=l^{-1}q,\beta=l^{-1},\gamma=pl^{-1}q-l,\delta=pl^{-1}.$

Looking for a metaplectic transformation $M$ in the fiber of $\Xi$, we get

(4.2.24) $\mathcal{M}_{p,l,q}\text{ with
kernel\quad}e^{-\frac{i\pi}{4}}e^{\frac{i\pi
m}{2}}{|l|}^{1/2}e^{i\pi(px^{2}-2lxy+qy^{2})}.$

We consider now, with $a,b$ linear form on $\mathbb{R}^{2}$,

$2\operatorname{Re}\langle{\text{{O}p}_{\text{{w}}}(a)}\mathcal{M}_{p,l,q}v,i\mathcal{M}_{p,l,q}{\text{{O}p}_{\text{{w}}}(b)}v\rangle=\frac{1}{2\pi}\left\\{a\circ\Xi,b\right\\}\|v\|^{2}$

We choose $b(x,\xi)=x$ and we look for $a$ satisfying

$(a\circ\Xi)(x,\xi)=\xi,\quad\text{i.e.\quad}a(\alpha x+\beta\xi,\gamma
x+\delta\xi)=\xi+\rho x.$

We choose then

$a(y,\eta)=\beta^{-1}y\quad\text{so that \quad}a(\alpha x+\beta\xi,\gamma
x+\delta\xi)=\beta^{-1}(\alpha x+\beta\xi)=\xi+\rho x.$

We obtain thus that

$2{|\beta|}^{-1}\mathcal{V}\bigl{(}\mathcal{M}v\bigr{)}^{1/2}\mathcal{V}(v)^{1/2}\geq\frac{1}{2\pi}\|v\|^{2}$

that is

$\mu(\mathcal{M})\geq\frac{|\beta|}{4\pi}.$

Let us now show the equality $\mu(\mathcal{M})=\frac{|\beta|}{4\pi}.$ With
$\mathcal{M}$ given by (4.2.24), we follow the calculations of Section 4.2.2
in our Appendix and we get $\mu(M)\leq\frac{|\beta|}{4\pi}$, proving the
sought result.

$\bullet$ We assume now that $\beta=0$, so that

$\Xi=\begin{pmatrix}\alpha&0\\\
\gamma&\alpha^{-1}\end{pmatrix},\quad\alpha\not=0.$

We find then that

$\Xi=\Xi_{a,b,0},\quad b=\alpha^{-1},a=b\gamma.$

The operator

$(M_{a,b,0}v)(x)=e^{\frac{i\pi m}{2}}|b|^{1/2}\int
e^{i\pi(ax^{2}+2bx\eta)}\hat{v}(\eta)d\eta$

Since $\mathcal{V}(e^{i\pi ax^{2}}w)=\mathcal{V}(w)$, we may assume that $a=0$
and

$(Mv)(x)=b^{1/2}v(bx),$

yielding $\mu(M)=0.$ We have thus proven the following result.

######  Theorem 4.11.

Let $M\in\textsl{Mp}(1)$. Then we have
$\Xi=\Psi(M)\in\textsl{Sp}(2,\mathbb{R})=\textsl{Sl}(2,\mathbb{R})$ and

$\Xi=\begin{pmatrix}\alpha&\beta\\\ \gamma&\delta\end{pmatrix}.$

Then we have $\mu(M)=\frac{|\beta|}{4\pi}.$

######  Lemma 4.12.

Let $\omega,c$ be real numbers and let
$\mathcal{H}_{c,\omega}={\text{{O}p}_{\text{{w}}}((\omega
x)^{2}+(c\xi+x)^{2})}.$ Then we have

(4.2.25)
$\inf\bigl{[}\operatorname{Spectrum}\mathcal{H}_{c,\omega}\bigr{]}=\frac{|\omega
c|}{2\pi}.$

###### Proof.

We may assume that $c\not=0$, since the result is obvious when $c=0$. (see
e.g. Remark 1.3.) We define then $\alpha=\sqrt{|\omega/c|},$ and we obtain
that, thanks to the symplectic covariance of the Weyl calculus,
$\mathcal{H}_{c,\omega}$ is unitarily equivalent to the operator with Weyl
symbol

$\omega^{2}\alpha^{-2}x^{2}+\bigl{(}c\alpha\xi+\alpha^{-1}x\bigr{)}^{2}=|c\omega|\bigl{(}x^{2}+(\xi+c^{-1}\alpha^{-2}x)^{2}\bigr{)},$

which is unitarily equivalent to the harmonic oscillator
$|c\omega|{\text{{O}p}_{\text{{w}}}(\xi^{2}+x^{2})}$ (also from the symplectic
covariance), whose spectrum is

$\frac{|c\omega|}{\pi}\Bigl{(}\frac{1}{2}+\mathbb{N}\Bigr{)},$

concluding the proof of the lemma. ∎

## References

  * [1] W. O. Amrein and A. M. Berthier, _On support properties of $L^{p}$-functions and their Fourier transforms_, J. Functional Analysis 24 (1977), no. 3, 258–267. MR 461025
  * [2] Michael Benedicks, _On Fourier transforms of functions supported on sets of finite Lebesgue measure_ , J. Math. Anal. Appl. 106 (1985), no. 1, 180–183. MR 780328
  * [3] Aline Bonami and Bruno Demange, _A survey on uncertainty principles related to quadratic forms_ , Collect. Math. (2006), 1–36. MR 2264204
  * [4] Aline Bonami, Bruno Demange, and Philippe Jaming, _Hermite functions and uncertainty principles for the Fourier and the windowed Fourier transforms_ , Rev. Mat. Iberoamericana 19 (2003), no. 1, 23–55. MR 1993414
  * [5] Jean-Michel Bony, _Evolution equations and generalized Fourier integral operators_ , Advances in phase space analysis of partial differential equations, Progr. Nonlinear Differential Equations Appl., vol. 78, Birkhäuser Boston, Boston, MA, 2009, pp. 59–72. MR 2664605
  * [6] John B. Conway, _A course in operator theory_ , Graduate Studies in Mathematics, vol. 21, American Mathematical Society, Providence, RI, 2000. MR 1721402
  * [7] Nuno C. Dias, Franz Luef, and João N. Prata, _Uncertainty principle via variational calculus on modulation spaces_ , J. Funct. Anal. 283 (2022), no. 8, Paper No. 109605, 30. MR 4453622
  * [8] I. Ekeland and H. Hofer, _Symplectic topology and Hamiltonian dynamics_ , Math. Z. 200 (1989), no. 3, 355–378. MR 978597
  * [9] Ivar Ekeland and Helmut Hofer, _Symplectic topology and Hamiltonian dynamics. II_ , Math. Z. 203 (1990), no. 4, 553–567. MR 1044064
  * [10] C. Fefferman and D. H. Phong, _The uncertainty principle and sharp Gȧrding inequalities_ , Comm. Pure Appl. Math. 34 (1981), no. 3, 285–331. MR 611747
  * [11] Charles L. Fefferman, _The uncertainty principle_ , Bull. Amer. Math. Soc. (N.S.) 9 (1983), no. 2, 129–206. MR 707957
  * [12] Gerald B. Folland, _Harmonic analysis in phase space_ , Annals of Mathematics Studies, vol. 122, Princeton University Press, Princeton, NJ, 1989\. MR 983366
  * [13] Saifallah Ghobber and Philippe Jaming, _Uncertainty principles for integral operators_ , Studia Math. 220 (2014), no. 3, 197–220. MR 3173045
  * [14] Karlheinz Gröchenig, _Foundations of time-frequency analysis_ , Applied and Numerical Harmonic Analysis, Birkhäuser Boston, Inc., Boston, MA, 2001\. MR 1843717
  * [15] M. Gromov, _Pseudo holomorphic curves in symplectic manifolds_ , Invent. Math. 82 (1985), no. 2, 307–347. MR 809718
  * [16] Lars Hörmander, _A uniqueness theorem of Beurling for Fourier transform pairs_ , Ark. Mat. 29 (1991), no. 2, 237–240. MR 1150375
  * [17] Jean Leray, _Analyse lagrangienne et mécanique quantique_ , Séminaire sur les Équations aux Dérivées Partielles (1976–1977), I, Collège de France, Paris, 1977, pp. Exp. No. 1, 303. MR 0501198
  * [18] by same author, _Lagrangian analysis and quantum mechanics_ , MIT Press, Cambridge, Mass., 1981, A mathematical structure related to asymptotic expansions and the Maslov index, Translated from the French by Carolyn Schroeder. MR MR644633 (83k:58081a)
  * [19] Nicolas Lerner, _Metrics on the phase space and non-selfadjoint pseudo-differential operators_ , Pseudo-Differential Operators. Theory and Applications, vol. 3, Birkhäuser Verlag, Basel, 2010. MR 2599384
  * [20] by same author, _Integrating the Wigner distribution on subsets of the phase space, a survey_ , Memoirs of the European Mathematical Society, vol. 12, EMS Press, Berlin, 2024. MR 4726947
  * [21] Elliott H. Lieb, _The stability of matter: from atoms to stars_ , fourth ed., Springer, Berlin, 2005, Selecta of Elliott H. Lieb, Edited by W. Thirring, and with a preface by F. Dyson. MR 2766495
  * [22] Elliott H. Lieb and Robert Seiringer, _The stability of matter in quantum mechanics_ , Cambridge University Press, Cambridge, 2010. MR 2583992
  * [23] Michael Reed and Barry Simon, _Methods of modern mathematical physics. I_ , second ed., Academic Press, Inc. [Harcourt Brace Jovanovich, Publishers], New York, 1980, Functional analysis. MR 751959
  * [24] Yiyu Tang, _A simple observation on Heisenberg-like uncertainty principles_ , Zap. Nauchn. Sem. S.-Peterburg. Otdel. Mat. Inst. Steklov. (POMI) 503 (2021), 113–120. MR 4337266
  * [25] André Unterberger, _Oscillateur harmonique et opérateurs pseudo-différentiels_ , Ann. Inst. Fourier (Grenoble) 29 (1979), no. 3, xi, 201–221. MR 552965
  * [26] Hermann Weyl, _Gruppen theorie und Quanten mechanik_ , second ed., Wissenschaftliche Buchgesellschaft, Darmstadt, 1977. MR 0450450
  * [27] Avi Wigderson and Yuval Wigderson, _The uncertainty principle: variations on a theme_ , Bull. Amer. Math. Soc. (N.S.) 58 (2021), no. 2, 225–261. MR 4229152
  * [28] Eugene Paul Wigner, _The collected works of Eugene Paul Wigner. Part A. The scientific papers. Vol. IV_ , Springer-Verlag, Berlin, 1997, Part I. Physical chemistry, Part II. Solid state physics, With a preface by Jagdish Mehra and Arthur S. Wightman, With annotations by Nandor Balazs and Walter Kohn, Edited by Wightman. MR 1635991


# 3D GRID-ATTENTION NETWORKS FOR INTERPRETABLE AGE AND ALZHEIMER’S DISEASE
PREDICTION FROM STRUCTURAL MRI

###### Abstract

We propose an interpretable 3D Grid-Attention deep neural network that can
accurately predict a person’s age and whether they have Alzheimer’s disease
(AD) from a structural brain MRI scan. Building on a 3D convolutional neural
network, we added two attention modules at different layers of abstraction, so
that features learned are spatially related to the global features for the
task. The attention layers allow the network to focus on brain regions
relevant to the task, while masking out irrelevant or noisy regions. In
evaluations based on 4,561 3-Tesla T1-weighted MRI scans from 4 phases of the
Alzheimer’s Disease Neuroimaging Initiative (ADNI), salience maps for age and
AD prediction partially overlapped, but lower-level features overlapped more
than higher-level features. The brain age prediction network also
distinguished AD and healthy control groups better than another state-of-the-
art method. The resulting visual analyses can distinguish interpretable
feature patterns that are important for predicting clinical diagnosis. Future
work is needed to test performance across scanners and populations.

Index Terms—  Deep learning, computer-aided diagnosis, Alzheimer’s disease,
convolutional neural network, biomarker, aging

##  1 Introduction

BrainAge [1] estimation refers to the task of estimating a person’s age from
their brain imaging data; this task is often performed by computing features
from healthy individuals’ structural MRI data (which may be supplemented by
diffusion or functional MRI [2]). A statistical or machine learning model is
then fitted to predict age, and tested on previously unseen subjects’ data.
The ‘gap’ or ‘delta’ between a person’s estimated BrainAge and their true
chronological age has been used as a marker of accelerated or delayed brain
aging. BrainAge gap has shown to be an accurate predictor of conversion from
mild cognitive impairment (MCI) to Alzheimer’s disease [3].

Successful BrainAge prediction algorithms have been developed based on linear
regression, using SVD to reduce the input feature dimension [2], and classical
machine learning methods such as support vector regression [4]. More recently,
deep learning methods such as convolutional neural networks (CNN) have been
applied to preprocessed and raw T1-weighted MRI data [5], and adapted to
accommodate multiple datasets using transfer learning [6].

However, these networks can overfit to the task of age prediction and be
invariant to disease [7], limiting the utility of the derived brain age as a
biomarker. Additionally, [8] brings up concerns regarding correcting for true
age when predicting brain age delta and stresses that more interpretable
methods need to be developed.

We developed a model that can learn BrainAge from distributed patterns in the
image at different levels, maintaining sensitivity to disease without limiting
training capacity, as in [7]. Our model also uses attention to produce
saliency maps for features at different levels of complexity.

Attention layers allow the network to focus on regions relevant to the task
while masking out noise or irrelevant regions. They have been used in medical
imaging analysis and have improved accuracy in classification [9] and
segmentation [10] tasks. Our model uses a 3D version of grid-attention
originally proposed for fetal ultrasound screening by [9]. Our contributions
are as follows:

  * •

We propose a new architecture for age prediction and Alzheimer’s disease
classification based on structural MRI.

  * •

Our method shows group differences between brain age delta for Alzheimer’s
disease and healthy controls without limiting the accuracy of age prediction.

  * •

Our method can visualize salient regions for classification at different
points in the model with different complexities of features.

  * •

We compare the saliency maps for age with those important for Alzheimer’s
classification and find that lower-level features have significantly more
overlap than higher level features. To our knowledge we are the first to
introduce this visual analysis to distinguish interpretable feature patterns
used to predict age versus Alzheimer’s disease classification.

Fig. 1: Model Architecture

##  2 METHOD

###  2.1 3D Grid-Attention Main Network

The backbone of our model is a simple 3D convolutional neural network (Figure
1). Esmaeilzadeh et al.[11] show that simple convolution models tend to
outperform more complex models for classification of Alzheimer’s disease by
reducing overfitting.

The network consists of 4 convolution blocks followed by a final convolutional
layer. Each convolution block consists of a single 3D convolution layer
(kernel size=3) followed by a batch norm and ReLU activation and an average-
pooling layer (kernel size=2, stride=2). We find that average-pooling performs
slightly better than max pooling for disease classification and our results
are consistent with those of Jin et al. [12]. Global-average-pooling (GAP) is
performed on the output of the final convolutional layer and the resulting
feature vector is input to a dense layer with varying number of nodes
depending on the classification/regression task. The numbers of filters for
each convolution are as follows: [32, 64, 64, 128, 128]. The number of output
nodes is 23 for the age prediction task and 2 for Alzheimer’s disease
classification.

Fig. 2: Grid-Attention Mechanism

###  2.2 Grid-Attention Mechanism

We use two attention modules based on [9] to better integrate low-level
features into the age prediction and AD classification. The attention module
is designed so that features learnt are spatially related to the global
features for the task. The architecture of our attention module is shown in
Figure 2. The ”attention score” is represented by the following equation:

$\alpha_{\mathit{l}}=\boldsymbol{\psi}\sigma\left(\mathbf{W}_{\mathit{f}}\mathbf{f}_{\mathit{l}}+\mathbf{W}_{\mathit{g}}\mathbf{g}\right)+\mathbf{b}_{\psi}$
(1)

where $\mathbf{f}_{\mathit{l}}\in\mathbb{R}^{H_{\mathit{l}}\times
W_{\mathit{l}}\times D_{\mathit{l}}}$ represents the output activation maps
for block $\mathit{l}$ before the pooling operation and
$\mathbf{g}\in\mathbb{R}^{H_{\mathit{g}}\times W_{\mathit{g}}\times
D_{\mathit{g}}}$ represents the output activation maps of the final
convolution layer before global average pooling.

$\mathbf{W}_{\mathit{f}}\in\mathbb{R}^{\mathit{C_{int}\times}C_{f}}$,
$\mathbf{W}_{\mathit{g}}\in\mathbb{R}^{\mathit{C_{int}\times}C_{g}}$, and
$\mathbf{W}_{\mathit{\psi}}\in\mathbb{R}^{\mathit{C_{int}\times}1}$ are linear
transformations, $\sigma$ is a non-linear transformation, and
$\mathbf{b}_{\mathit{\psi}}\in\mathbb{R}$ is a bias term. We choose ReLU for
the non-linearity and find that it results in more accurate predictions than
Softmax.

The resulting attention score $\alpha_{\mathit{l}}$ has a dimensionality of
1xHxWxD. This score is normalized per module by subtracting
$min(\alpha_{\mathit{l,i,j,k}})$ and then dividing by
$\sum{\alpha_{\mathit{l,i,j,k}}}$. This normalized attention score is
multiplied by each input activation map to produce the output of the module.
Finally, a sum is taken over feature maps and is connected to a dense layer
with the same number of nodes as the backbone network. Thus, regions with a
higher normalized attention score are more relevant to the task and are
highlighted, while noise and irrelevant features are masked out.

Our implementation connects attention modules after convolutions 3 and 4, as
shown in Figure 1. We implement $\mathbf{W}_{\mathit{f}}$ and
$\mathbf{W}_{\mathit{g}}$ as 1x1x1 convolution operations and use n=64, 128
intermediate filter dimensions respectively. Global feature maps are upsampled
using trilinear upsampling to match the dimensionality of earlier activation
maps. We aggregate the results by calculating the average of the dense output
layers for the two attention modules as well as the backbone network.

###  2.3 Kullback-Leibler Divergence (KLD) Loss

We formulate age prediction tasks similarly to [13] and [14]. Age is
represented by a discrete Gaussian distribution with mean equal to the true
age of a subject and with a standard-deviation of 2. We train our age
prediction model to minimize the Kullback-Leibler Divergence (KLD) between
this distribution and the soft-max network output. There are 23 bins in our
distribution: each bin represents an interval of 2 years and corresponds to 23
nodes in the age prediction network.

###  2.4 Data and Preprocessing

Our experiments analyze brain structural (3T) T1-weighted MRI scans from the
Alzheimer’s Disease Neuroimaging Initiative (ADNI) dataset. We use scans from
ADNI1, ADNI2, ADNI3, and ADNI-GO studies; our dataset consists of 4,561 scans
in total (2,004 non-accelerated, 2,557 accelerated) and contains 329 unique
subjects with AD (191 male, 138 female, mean age=75.7 years, standard
deviation=8.18) and 687 unique healthy controls (288 male, 399 female, mean
age=73.0, std=7.0). Scans are processed through a standard processing pipeline
consisting of reorientation using [15], skull stripping using HDBet [16],
nonparametric intensity normalization for bias field correction using
N4BiasCorrection [17], and 6-dof linear registration to MNI152 space using FSL
[15]. Scans are subsequently down-sampled to 2-mm resolution using ANTS [18].
The final dimensions of each scan are 91x109x109.

###  2.5 Experiment Set-up and Training

We split our data into 3 non-overlapping sets: training, validation (for
hyperparameter tuning and model stopping), and test sets. We create an age-
balanced validation set of n=109 (mean age=74.76 years, std=8.2, 53 AD, 56 CN)
and use the remaining scans to create 10 pairs of training/testing sets
(10-fold split). As our dataset may contain multiple scans for a single
subject, we split our data by unique subject ID to ensure that no subjects are
in common between training, validation, and testing sets within a single fold.
Follow-up and duplicate scans are used for data augmentation for our training
set but are filtered out from our validation and testing sets. For our age-
prediction models, we further filter out subjects with Alzheimer’s disease
from our training and validation sets. For testing, we use a single scan per
subject chosen at random; for standardization, we use a non-accelerated scan
if it exists but include accelerated scans for subjects without any non-
accelerated scans. [19] shows comparable group discrimination in accelerated
vs. non-accelerated scans. For training our age models, we sample from a
unique set of subject IDs and randomly select at most one of the multiple
scans for each subject. For training the Alzheimer’s classification model, we
split our subjects into two sets based on diagnostic class and sample equally
from both sets.

Scan intensities are standardized to mean zero and unit variance for non-zero
(non-background) pixels. We augment our training data using shift (uniform
(-0.1, 0.1)*std) and scale (uniform (-0.1, 0.1)) intensity augmentation on
non-zero pixels and random (p=0.5) hemispherical mirroring.

We initialize our model weights using Kaiming weight initialization [20] and
train our model using the ADAM [21] optimizer, with batch size of 16 and
learning rate of 1e-5. We use a weight decay of 1e-2 to prevent overfitting.
We minimize KLD Loss (detailed in section 2.3) to train our age model and
2-class cross entropy for the Alzheimer’s disease classification model and
train our models until 10 epochs without performance improvement on the
validation set. Test data is evaluated on the state with the lowest validation
loss.

##  3 RESULTS

###  3.1 Models Results

Our final age prediction model has a mean absolute error (MAE) of 3.96 years
on healthy controls (RMSE=5.17, Corr=0.686) for n=991 subjects from the ADNI
dataset. For comparison, we replicated our training/evaluation procedure with
[13] which received state-of-the-art performance on the UK Biobank dataset and
received an MAE of 4.36 on healthy controls (RMSE of 5.56, Corr=0.66).
Although state-of-the art shows much lower MAE for brain age prediction (2.86
MAE [14]), we attribute this to having far more training samples (7,312
scans). Our final AD prediction model has an accuracy of 0.868 (balanced
accuracy=0.84, precision=0.81, recall=0.76) over n=991 subjects.

Fig. 3: The histograms show the brain age delta for our attention model (top)
and for an alternative state-of-the-art model [13] (bottom). Our method shows
a greater sensitivity to AD: the average brain age delta is 3.09 (SD=6.76) for
subjects with AD (Orange) and -0.89 (SD=5.10) for healthy controls(Blue). The
compared model has an average brain age delta of 1.00 (SD=6.99) for subjects
with AD and -0.58 (SD=5.62) for healthy controls.

###  3.2 Sensitivity to Disease Effects

The distributions of brain age delta (the difference between predicted and
true age) for AD and healthy controls are shown in Figure 2 (top). The MAE for
AD subjects is 5.7 years (RMSE: 7.4 years). For comparison, we also included
the distribution for the model based on [13] (bottom); the MAE for AD subjects
is 5.56 years (RMSE: 7.01 years). As can be observed, our model detects a
greater group effect on BrainAge for AD. This observation corresponds to a
higher Cohen’s d effect size (95 percent confidence interval) for our model
(0.66 $\pm$ 0.13; P-value $<$ 0.00001) as compared to [13] (0.25 $\pm$ 0.13;
P-value = 0.0002). We attribute this to our model using fewer convolutional
blocks for the backbone architecture and the use of attention modules to
enforce the use of lower level features for age prediction.

Fig. 4: Saliency visualization maps produced by the model for two different
subjects. Three left panels show saliency maps for the first attention module
and the rightmost three panels correspond to the second module. Red colors
denote voxels relevant for AD classification, blue voxels for age prediction,
and purple colors show the region of overlap of voxels relevant for each task.
Top row uses a threshold of 2 standard deviations above mean attention, while
the bottom row uses a threshold of 1 std above the mean.

###  3.3 Visualization and Comparison

|  A1 (CN/AD) |  A2 (CN/AD)
---|---|---
$\mu+2\sigma$ |  0.39/0.39 |  0.29/0.30
$\mu+1\sigma$ |  0.43/0.41 |  0.20/0.26
$\mu+0.5\sigma$ |  0.52/0.52 |  0.10/0.20
Table 1: Average overlap scores between saliency maps for the age prediction
and AD prediction tasks, separate for AD subjects (AD) versus healthy controls
(CN) for varying intensity levels of thresholded attention value. Column A1 is
the average of scores for the first attention module while column A2 is the
average of scores for the second attention module.

We visualize region saliency using the attention masks (normalized attention
scores) produced by the attention modules. We upsample these masks to the
original size of our scan using nearest neighbor interpolation and mask out
artifacts by multiplying by the binary brain mask. For analysis, we apply
binary thresholding based on the mean attention value. Figure 4 shows examples
of saliency maps produced by our model at different thresholds. We note that
saliency maps produced by the network differ by subject and are based on
relevant areas used to classify each scan.

In addition, we analyze the difference in region overlap between salient
features for age prediction and AD classification. For each subject, we
calculate the Jaccard score between the thresholded saliency maps for AD
classification and age prediction:

$\mathbf{J}=\frac{|AD\cap AGE|}{|AD\cup AGE|}$ (2)

where AD and AGE are binary (thresholded) importance regions. Results are
presented in Table 1. As the $\sigma$ threshold decreases, there is a greater
difference in overlap between attention modules. For $\sigma=1$ there is a
0.52 overlap for the first attention module compared to only 0.1 for the
subsequent module. This implies that feature complexity plays a crucial role
in common region saliency between tasks. Visually, this can also be seen in
Figure 4. Age prediction focuses on a much larger region of features compared
to AD classification which focuses mainly on the ventricles. This may explain
why more complex age prediction models tend to more accurately predict age in
diseased subjects. We also observed little difference between average scores
for subjects with AD versus healthy controls.

##  4 CONCLUSION

In this work, we propose an interpretable 3D Grid-Attention network that can
accurately predict age and Alzheimer’s disease from structural MRI. We compare
saliency maps produced by this model between age prediction and Alzheimer’s
disease and find that the features used for each task do overlap, but are
partially distinct: age prediction tends to use more widespread features,
while AD prediction makes use of ventricular expansion, which is mild in
normal aging but can be severe in AD. When brain aging and AD classification
models are forced to use the same small set of features, it is almost
inevitable that people with AD appear to have accelerated aging. The current
brain aging model also focuses on brain regions that are less affected by
pathology. This may also explain why complex deep learning models trained for
brain age prediction may be somewhat invariant to detecting shifts in brain
age, in disease [7]. The overlap of the salience maps for each task offers a
more intuitive understanding of what image regions are used for each task, and
their differences encourages further study of localized brain aging for
biomarker development. Future work will involve extending our analysis to
include subjects with MCI, a precursor to Alzheimer’s disease and extending
our architecture to predict both age and Alzheimer’s disease status
simultaneously using multi-task learning. Further work should also test the
approach on multiple datasets to fully assess generalization performance for
other scanners and populations. Ongoing work in domain adaptation and
adversarial networks will also be needed to train models that generalize well
and offer invariance to scanning protocols [22] [23].

##  5 Compliance with Ethical Standards

Data used in this article were obtained from the public Alzheimer’s Disease
Neuroimaging Initiative (ADNI) database (adni.loni.usc.edu). Further ethical
approval was not required, as we studied previously collected, anonymized, de-
identified data.

##  6 Acknowledgments

This research was supported in part by the National Institutes of Health (NIH)
under grants U01AG068057, R01AG059874, and RF1AG051710, and by a research
grant from Biogen, Inc.

## References

  * [1] Katja Franke and Christian Gaser,  “Ten years of BrainAGE as a Neuroimaging Biomarker of Brain Aging: What Insights Have We Gained?,”  Frontiers in Neurology, vol. 10, pp. 789, October 2019.
  * [2] Stephen M. Smith et al.,  “Estimation of brain age delta from brain imaging,”  Neuroimage, vol. 200, pp. 528–539, October 2019.
  * [3] Christian Gaser, Katja Franke, Stefan Klöppel, Nikolaos Koutsouleris, and Heinrich Sauer,  “BrainAGE in Mild Cognitive Impaired Patients: Predicting the Conversion to Alzheimer’s Disease,”  PLOS ONE, vol. 8, pp. 1–15, June 2013.
  * [4] Laura K.M. Han et al.,  “Brain aging in major depressive disorder: Results from the ENIGMA Major Depressive Disorder Working Group,”  Molecular Psychiatry, May 2020.
  * [5] James H. Cole et al.,  “Predicting brain age with deep learning from raw imaging data results in a reliable and heritable biomarker,”  Neuroimage, vol. 163, pp. 115–124, December 2017.
  * [6] B.A. Jonsson et al.,  “Deep learning based brain age prediction uncovers associated sequence variants,”  Nature Communications, November 2019.
  * [7] Vishnu M. Bashyam et al.,  “MRI signatures of brain age and disease over the lifespan based on a deep brain network and 14,468 individuals worldwide,”  BRAIN, vol. 143, pp. 2312–2324, July 2020.
  * [8] Ellyn R. Butler et al.,  “Statistical pitfalls in brain age analyses,”  bioRxiv, 2020.
  * [9] Jo Schlemper et al.,  “Attention-Gated Networks for Improving Ultrasound Scan Plane Detection,”  arXiv, 2018.
  * [10] Ashish Sinah and Jose Dolz,  “Multi-scale self-guided attention for medical image segmentation,”  arXiv, 2019.
  * [11] Soheil Esmaeilzadeh, Dimitrios I. Belivanis, Kilian M. Pohl, and Ehsan Adeli,  “End-To-End Alzheimer’s Disease Diagnosis and Biomarker Identification,”  arXiv, 2018.
  * [12] Dan Jin et al.,  “Attention-based 3D Convolutional Network for Alzheimer’s Disease Diagnosis and Biomarkers Exploration,”  in 2019 IEEE International Symposium on Biomedical Imaging (ISBI) 2019 April 8-11, 2019, 2019, pp. 1481–1486.
  * [13] Han Peng et al.,  “Accurate brain age prediction with lightweight deep neural networks,”  bioRxiv, 2019.
  * [14] Pradeep K. Lam et al.,  “Accurate brain age prediction using recurrent slice-based networks,”  bioRxiv, 2020.
  * [15] Mark Jenkinson, Peter Bannister, Michael Brady, and Stephen M. Smith,  “FSL,”  NeuroImage, vol. 62, pp. 782–790, August 2012.
  * [16] Fabian Isensee et al.,  “Automated brain extraction of multisequence MRI using artificial neural networks,”  Human Brain Mapping, vol. 40, no. 17, pp. 4952–4964, 2019.
  * [17] Nicholas J. Tustison et al.,  “N4ITK: Improved N3 Bias Correction,”  IEEE Transactions on Medical Imaging, vol. 29, pp. 1310–1320, June 2010.
  * [18] Mark Jenkinson, Peter Bannister, Michael Brady, and Stephen M. Smith,  “Large-scale evaluation of ANTs and FreeSurfer cortical thickness measurements,”  NeuroImage, vol. 99, pp. 166–179, October 2014.
  * [19] Prashanthi Vemuri et al.,  “Accelerated vs. unaccelerated serial MRI based TBM-SyN measurements for clinical trials in Alzheimer’s disease,”  Neuroimage, vol. 113, pp. 61–69, March 2015.
  * [20] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun,  “Delving Deep into Rectifiers: Surpassing Human-Level Performance on Imagenet Classification,”  in 2015 IEEE International Conference on Computer Vision (ICCV), 2015, pp. 1026–1034.
  * [21] Diederik Kingma and Jimmy Ba,  “Adam: A Method for Stochastic Optimization,”  in 2015 International Conference on Learning Representations (ICLR), 2015.
  * [22] Hao Guan et al.,  “Attention-guided Deep Domain Adaptation for Brain Dementia Identification with Multi-site Neuroimaging Data,”  in Domain Adaptation and Representation Transfer, and Distributed and Collaborative Learning. 2020, pp. 31–40, Springer International Publishing.
  * [23] Nicola K. Dinsdale et al.,  “Unlearning Scanner Bias for MRI Harmonisation,”  in Medical Image Computing and Computer Assisted Intervention – MICCAI 2020. 2020, pp. 369–378, Springer International Publishing.