# Multimodal Brain Disease Classification with Functional Interaction Learning
from Single fMRI Volume

Wei Dai, Ziyao Zhang, Lixia Tian, Shengyuan Yu, Shuhui Wang, Zhao Dong, and
Hairong Zheng  Wei Dai and Ziyao Zhang contributed equally;
Corresponding authors: Shuhui Wang; Zhao Dong.
W. Dai, S. Yu, and Z. Dong are The Department of Neurology, Chinese PLA
General Hospital, Beijing 100853, China. Email<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
Z. Zhang and H. Zheng are with the Paul C. Lauterbur Research Center for
Biomedical Imaging, Shenzhen Institutes of Advanced Technology, Chinese
Academy of Sciences, Shenzhen, Guangdong 518000, China and Peng Cheng
Laboratory, Shenzhen, Guangdong 518000, China. Email<EMAIL_ADDRESS><EMAIL_ADDRESS>
L. Tian is with the School of Computer and Information Technology, Beijing
Jiaotong University, Beijing 100044, China. Email<EMAIL_ADDRESS>
S. Wang is with the Institute of Computing Technology, Chinese Academy of
Sciences, Beijing 100190, China. Email<EMAIL_ADDRESS>

###### Abstract

In neuroimaging analysis, functional magnetic resonance imaging (fMRI) can
well assess the function changes for brain diseases with no obvious structural
lesions. To date, most deep-learning-based fMRI studies have employed
functional connectivity (FC) as the basic feature for disease classification.
However, FC is calculated on time series of predefined regions of interest and
neglects detailed information contained in each voxel. Another drawback of
using FC is the limited sample size for the training of deep models. The low
representation ability of FC leads to poor performance in clinical practice,
especially when dealing with multimodal medical data involving multiple types
of visual signals and textual records for brain diseases. To overcome this
bottleneck problem in the fMRI feature modality, we propose BrainFormer, an
end-to-end functional interaction learning method for brain disease
classification with single fMRI volume. Unlike traditional deep learning
methods that construct convolution and transformers on FC, BrainFormer learns
the functional interaction from fMRI signals, by modeling the local cues
within each voxel with 3D convolutions and capturing the global correlations
among distant regions with specially designed global attention mechanisms from
shallow layers to deep layers. Meanwhile, BrainFormer can deal with multimodal
medical data including fMRI volume, structural MRI, FC features and phenotypic
data to achieve more comprehensive brain disease diagnosis. We evaluate
BrainFormer on five independent multi-site datasets on autism, Alzheimer’s
disease, depression, attention deficit hyperactivity disorder and headache
disorders. The results demonstrate its effectiveness and generalizability for
multiple brain diseases diagnosis with multimodal features. BrainFormer may
promote precision of neuroimaging-based diagnosis in clinical practice and
motivate future studies on fMRI analysis. Code is available at:
https://github.com/ZiyaoZhangforPCL/BrainFormer.

###### Index Terms:

Functional MRI, brain disease classification, multimodal, Transformer, 3D CNN.

##  1 Introduction

Recent years have witnessed the significant development in deep learning.
Extensive progresses have been made in areas such as vision [1, 2, 3] and
language [4]. Deep learning is also proliferating in medical domains,
including digital health, patient monitoring, disease diagnostics and so on.
These application scenarios involve data in multiple forms or modalities such
as various types of medical images, clinical records and tabular data. There
are urgent demands for multimodal medical analysis due its wide application
prospect [5]. Methods for disease classification, lesion detection, structural
segmentation, image registration, multimodal image synthesis and fusion, using
cutting-edge deep learning models, _e.g._ , Convolutional Neural Networks
(CNN) [6, 7, 8], Graph Convolution Network(GCN) [9, 10, 11], and Transformers
[12, 13, 14, 15, 5, 16, 17, 18, 19, 20], are developing quickly.

Figure 1: (a) sMRI and fMRI data of the same person; (b) 4D brain fMRI data:
3D fMRI volumes of the brain along temporal dimension, and sections of
sagittal, coronal and axial views are shown for better visualization; (c) the
functional connectivity between two predefined brain regions is calculated by
the correlation coefficients between the time series of the two regions.

Among diversified types of medical imaging modalities, magnetic resonance
imaging (MRI) is a powerful non-invasive tool for disease diagnosis. For
example, structural MRI (sMRI) can record the detailed structure and reflect
the signals of lesion areas. sMRI is particularly useful for evaluating
complex organs such as the brain, and promising progresses have been achieved
recently using CNN [21] and Transformers [12, 13, 14, 15] as the backbone.
However, sMRI appears to be less effective for evaluating diseases displaying
no obvious structural changes, such as psychological disorders, autism,
headache disorders and so on. These diseases are traditionally considered as
brain functional disorders rather than structural abnormalities. Diagnoses of
these diseases are based on clinical interviews and behavioral assessments, so
the accuracy is highly dependent on the skills and professionalism of doctors.
Thus, sMRI cannot facilitate diagnosis or treatment well on these diseases in
medical practice.

Unlike sMRI, functional MRI (fMRI) uses the blood oxygen level dependent
(BOLD) signals to reflect neural activity along the temporal dimension with a
4D tensor, see Fig. 1 (a). As shown in Fig. 1 (b), the 4D tensor consists of a
sequence of scanned 3D brain volumes arranged along the temporal dimension,
where each voxel in the 3D volume recording the BOLD signals of a specific
brain region in a given time interval. Since fMRI provides an indirect measure
of neural activity, the major challenge is to model the actual functional
interactions among brain areas from the fluctuating fMRI data. There have been
a bulk of techniques for extracting features from fMRI data in a hand-crafted
manner to model the connectivity from spatial-temporal correlation, frequency,
region or modularity aspects [22]. In particular, the functional connectivity
(FC) is the most widely used technique, which records the pair-wise
correlation coefficients between sets of predefined regions of interest (ROI).
Two ROIs are assumed to be functionally connected if they display synchronous
functional activity fluctuations, see Fig. 1 (c). FC has been playing an
essential role in fMRI-based brain disease classification, in companion with
off-the-shelf models, e.g., Support Vector Machine (SVM) [23], Deep Neural
Network (DNN) [24, 25, 26], GCN [9, 11] and Transformers [5].

However, FC-based models have often been reported to be of relatively low
accuracies, due to the following reasons. The FC feature is built on the whole
4D tensor of one specific person with predefined brain ROIs, and calculated as
the overall correlation between brain region pairs using all 3D volumes. This
leads to extremely low dimensions of FC, and accordingly insufficient
representation ability and training data size for consequent classifications.
Given the coarse granularity of FC, the disease biomarker may be hard to
detect even with strong DNN. Of even greater concern, due to the physical
differences among MRI scanners used in different study sites, FC features
calculated using a single fMRI dataset can hardly be calibrated and aligned to
those of other datasets, and thus possesses poor generalizability across
different datasets and brain diseases.

Given the limitations of traditional hand-crafted features and the
corresponding classification paradigms, we employ deep feature learning
directly on fMRI volume instead of the 4D tensor in existing paradigms [22,
23, 24, 25, 26, 27, 10, 11], aiming to model the functional interaction among
different brain areas in a fine-grained manner. We assume reasonably that each
3D volume contains all the information in need, so working on the 3D volume
input enhances the available training data size by hundreds of times through
using volume-wise feature instead of the subject-wise FC feature, meanwhile
avoiding drastic information losses brought by FC feature, thus the feature
representation ability has the potential to be stronger.

On the 3D volume input, the key to learn functional interaction is to capture
the correlation among components of different levels flexibly. Towards this
goal, we apply 3D CNN [28, 29, 30] together with a data normalization layer to
extract the informative features from the 3D volume. The 3D convolution is
more accurate than any hand-crafted techniques in modeling the correlation
within a certain spatial-temporal receptive field, while the data
normalization is useful in alleviating the distribution divergence among
different fMRI volumes collected from multiple sites. On top of 3D CNN, we
design global attention mechanism inspired by Transformer [4] to model the
long-range function connectivity among different brain areas. Directly
applying the original self-attention [4] to fMRI data leads to prohibitive
computation complexity, especially at shallow layers with high feature
resolution. We thus design two attention mechanisms towards better
computational efficiency, i.e., Shallow Global Attention (SGA) for shallow
layers to exchange information globally through fully connected layer, and
Deep Global Attention (DGA) for deep layers to fuse global information through
global attention mask. Our hierarchical feature extraction framework is
constructed with collaboration of the global attention and 3D CNN capturing
local feature correlations (Fig. 2), and trained end-to-end without manual
brain region labeling.

Furthermore, it is necessary to fuse the complementary modalities for disease
classification. Different neuroimaging methods provide different views for the
same brain. Take Alzheimer’s disease (AD) for example, AD in early stage
usually shows no structural changes but functional alterations of the brain,
thus fMRI is capable to evaluate the dynamic function. With the progression of
AD, sMRI can detect signs of brain atrophy due to cell death and tissue loss.
Other neuroimaging methods can provide information including metabolic
activity, biochemical changes, and so on. Phenotypic data provide clinical
information regarding patients’ disease symptoms, as well as relevant
demographic data, such as age, ethnicity and sex. Therefore, combining
multiple modalities can provide more comprehensive information for disease
evaluation and prediction, and probe indepth into the underlying
pathophysiology.

Accordingly, we propose BrainFormer, an ROI label-free multimodal brain
disease classification framework. As for the fMRI modality, it can effectively
learn the functional interaction features in various brain disease
classification tasks. Beyond that, it can also fully leverage other
complementary modalities such as sMRI, FC and phenotypic data via a simple
multimodal feature fusion structure to further boost the brain disease
classification accuracy.

We evaluate BrainFormer on five independently acquired brain disease datasets
that contain multi-site data: ABIDE, ADNI, MPILMBB, ADHD-200 and ECHO. The
diseases evaluated include major depressive disorder (MDD), AD and mild
cognitive impairment (MCI), autism spectrum disorder (ASD) and Asperger’s
syndrome (AS), attention deficit hyperactivity disorder (ADHD), migraine
disease (MD) and medication-overuse headache (MOH). BrainFormer achieves
state-of-the-art performance on all the eight brain diseases. In clinical
practice, we have also applied our method on different types of headache
detection.

The main contributions of this study are summarized as follows:

  * •

We propose BrainFormer, the first brain disease classification model, which
learns functional interaction among brain regions directly on single fMRI
volume. BrainFormer can fuse multimodal medical data to achieve more
comprehensive brain disease diagnosis. It demonstrates strong generalizability
on eight types of brain diseases.

  * •

BrainFormer well captures both local and distant feature correlations by
specifically designed SGA and DGA. The disease-related biomarkers can be
precisely located within fMRI volumes by gradient-based localization,
facilitating brain disease diagnosis in clinical practice.

  * •

Experiments on five independent datasets with eight types of brain diseases
validate the promising ability of our multimodal brain disease classification
framework, and the generalizability on brain disease diagnosis without obvious
structure changes.

##  2 Related Work

fMRI-based classification: Extensive studies on brain disease classification
using fMRI are based on FC, which characterizes the pairwise correlation
between two predefined brain regions [31, 32, 33]. Wang _et al._ [31] perform
sparse multiview task-centralized ensemble learning for ASD binary
classification with age and sex information on FC features. Wang _et al._ [33]
conduct 3-class classification on ASD by introducing label distribution
learning to FC matrix (the labels describing either shared or specific disease
features of the subjects). GCN is also widely used for fMRI data analysis [9,
10, 11]. Parisot _et al._ [9] design a spectral graph convolution that predict
brain disease based on FC features. Huang _et al._ [10] propose edge-
variational graph convolutional networks to predict ASD, AD and ocular
disease. Huang _et al._ [34] construct multiple FC networks for each
participant, and create a multi-task learning model to learn the most
discriminative features in each template for ASD classification. A recent work
by Zhang _et al._ [5] introduces diffusion kernel attention module in
Transformer to construct functional brain networks and conduct classification
on ADHD and AD. However, FC only encodes the correlation between different
brain regions, and loses detailed cues within each individual region;
therefore, it is difficult to locate disease-related biomarkers that aid
diagnosis. Very few studies employ single fMRI volumes for disease
classification [35]. Vu _et al._ [35] build a vanilla 3D CNN for single fMRI
volume classification, which consists of 3 convolution layers, however, the
correlation cues between distant regions are still lost.

Multimodal classification: Deep models have been built to process and
correlate information from multiple modalities for joint feature
representations (including but not limited to sMRI, fMRI, positron emission
tomography, electroencephalogram and phenotypic data) [36]. Different
modalities can provide information in various aspects and compensate the
shortcomings of each other to boost the performance for classification. Liu
_et al._ [36] combine sMRI and fMRI with a multiheaded gating fusion model and
achieve remarkable accuracies for ASD, ADHD and schizophrenia classification.
Huang _et al._ [37] use GCN to incorporate multimodal data including sMRI,
fMRI, fundus photos and phenotypic data for AD, ASD and ocular disease
classification. Ji _et al._ [38] perform deep graph hashing learning to
restore semantic space as well as topological features in fMRI for brain-
network-based classification for ADHD and ASD, using age, sex, handedness and
clinical features. Supekar _et al._ [39] develop a spatiotemporal deep neural
network model that can identify interpretable dynamic markers, which combines
fMRI time series, age, sex and study site to predict ASD from neurotypical
individuals. Sun _et al._ [40] take FC, two BOLD-derived features (ReHo and
ALFF), and gray matter volume from sMRI as input, and select features using
two-nested leave-one-out cross validation method for MDD prediction.

Figure 2: The illustration of proposed (a) BrainFormer and (b) multimodal
framework for brain disease classification. BrainFormer is used to process 3D
volumes include fMRI and sMRI, the functional connectivity (FC) feature and
phenotypic data are processed by Multi-Layer Perceptron (MLP).

3D CNN: The 3D CNN encodes the spatial-temporal cues or 3D structure through
sliding 3D convolution kernel along consecutive video frames [41, 42]. Tran
_et al._ [28] design a convolution 3D network (C3D) through directly stacking
several 3D convolution layers. However, the 3D CNN lacks pretrained parameters
on large-scale datasets, which increases the model training difficulty.
Carreira _et al._ [41] build a deep Inflated 3D (I3D) network through
inflating the 2D convolution kernels in 2D CNN to the corresponding 3D
kernels. Qiu _et al._ [42] factorize the 3D convolution kernels into 2D
spatial kernel and 1D temporal kernel, which significantly reduce the
parameters without performance drop. Li _et al._ [29] build a compact Multi-
scale 3D (M3D) convolution to learn multi-scale temporal cues with dilated
temporal convolution. 3D CNN is also suitable for fMRI data. However, 3D CNN
is unable to capture long-range relations among distant brain regions.
Moreover, previous 3D CNN methods are designed for single-site data, which
suffer serious performance drop when directly applied to multi-site fMRI data.

Transformer: Transformer [4] is proposed to model long-range dependencies
among tokens. Inspired by the successes in language, recent studies are trying
to introduce Transformer into computer vision tasks [3, 2, 43, 20].
Transformer is also used in medical image classification [5], segmentation
[13, 14, 12, 17, 15] , registration [18], multimodal image synthesis and
fusion [16, 19]. Specifically, Gao _et al._ [13] propose a hybrid transformer
architecture for sMRI image segmentation. Chen _et al._ [18] create a hybrid
CNN-Transformer model for volumetric medical image registration in MRI and CT.
In multimodal medical image synthesis, Dalmaz _et al._ [16] has improved the
capture ability of contextual relations while maintaining the power of
localization by aggregating convolutional operators and transformers.
Multiscale adaptive Transformer is proposed by Tang _et al._ [19] to achieve
multimodal medical image fusion, facilitating clinical diagnosis and surgical
navigation. Transformer is capable of modeling long-range dependencies between
distant brain regions in fMRI volume. Nevertheless, computing fully connected
attention mask for high-dimensional fMRI data may suffer unaffordable
computation.

Our BrainFormer is an end-to-end architecture for brain disease classification
based on single fMRI volumes. Compared with FC-based method, our method
jointly captures local details within each brain region and global
relationships between the regions directly on single fMRI volumes, and
exhibits stronger capability in brain disease classification. Moreover, our
method directly make prediction on fMRI volume, hence is more suitable to
locate disease-related biomarkers, which may further facilitate neuroimaging-
based disease diagnosis in clinical practice.

##  3 Proposed Method

###  3.1 Multimodal brain disease classification framework

In this study, we consider multimodal data for brain disease classification,
including fMRI volume, sMRI volume, FC feature and phenotypic data. As shown
in Fig. 2(b), the brain fMRI of one person can be viewed as a sequence of 3D
tensor
$\mathbf{V^{f}}=\\{\mathcal{V}_{1}^{f},\mathcal{V}_{2}^{f},...,\mathcal{V}_{T}^{f}\\},\
\mathcal{V}_{i}^{f}\in\mathbb{R}^{D\times H\times W}$, where $T$ denotes the
temporal length, and $D$, $H$ and $W$ represent the depth, height and width of
a 3D volume, respectively. We only use single 3D fMRI volume in our method.
The sMRI can also be represented as 3D tensor
$\mathbb{V}^{s}\in\mathbb{R}^{D\times H\times W}$. The FC feature is a fully
connected similarity matrix $\mathbf{f}^{fc}\in\mathbb{R}^{p\times p}$, where
$p$ is the partition number brain ROIs. And the phenotypic data $\mathbf{s}$
is a series of scalar values.

As shown in Fig. 2, in our multimodal framework, the fMRI and sMRI volumes are
processed by BrainFormer. BrainFormer can be used for single fMRI volume
classification, and can also be used to process sMRI in multimodal framework.
The FC feature and phenotypic data are processed by a 3-layer Multi-Layer
Perceptron (MLP). The features from data modalities are combined for
classification. The following sections will introduce BrainFormer in detail.

###  3.2 Formulation

As discussed above, our multimodal framework consists of four modal input
data, including fMRI volume, sMRI volume, FC and scalar phenotypic data. Based
on these multimodal data, the brain disease classification can be performed
by,

$p_{i}=f(\mathcal{V}^{f},\mathcal{V}^{s},\mathbf{f}^{fc},\mathbf{s}),$ (1)

where $p_{i}\in\mathbb{R}^{n}$ denotes the classification score for the $i$-th
volume, and $n$ is the predefined class number, _e.g._ , 2 classes for
diseased and healthy individual. We use $f$ to denote our classification
model. Our model takes multimodal diagnosis data as input and outputs the
predicted classification score. Fig. 2 illustrates the overall framework. As
shown in Fig. 2(b), our multimodal framework use BrainFormer to process fMRI
and sMRI volumes, and use MLP process FC feature and clinical data. The
features of multimodal data are unified in a self-attention model for
classification. The following section will introduce the BrainFormer model in
detail, consists of a data normalization layer, 3D CNN backbone and global
attention blocks.

###  3.3 BrainFormer

####  3.3.1 Data Normalization Layer

The fMRI datasets are usually collected from multiple sites, among which data
distributions differ. Therefore, we apply a data normalization layer to
normalize all fMRI volumes to identical distribution. We apply instance
normalization on 3D volume $\mathcal{V}_{i}$ to normalize it to zero-mean and
unit-variance distribution,

$\bar{\mathcal{V}}_{i}=\frac{\mathcal{V}_{i}-\operatorname{mean}(\mathcal{V}_{i})}{\operatorname{std}(\mathcal{V}_{i})},$
(2)

where $\bar{\mathcal{V}}_{i}$ is the normalized volume.

####  3.3.2 3D CNN Backbone

We employ 3D CNN as backbone feature extractor. The 3D CNN accepts 3D fMRI
volume $\mathcal{V}\in\mathbb{R}^{D\times H\times W}$ as input and outputs the
corresponding 3D feature maps $\mathbf{f}\in\mathbb{R}^{d\times h\times w}$
(we omit the channel dimension for simplicity). Our 3D CNN consists of several
convolution blocks, which are made up of 3D convolution layers, batch
normalization layers and nonlinear activation layers. The stride of last
convolution layer in each block is set to 2.

3D convolution layer: The 3D convolution layer encodes the 3D structural cues
through sliding 3D kernel along each dimension of 3D fMRI volume. A 3D
convolution kernel can be formulated as a 3D tensor
$\mathbf{W}\in\mathbb{R}^{d\times h\times w}$, where $d$, $h$ and $w$ are the
depth, height and width of kernel (the input and output channels are omitted).
Computation of a $3\times 3\times 3$ sized 3D convolution can be formulated
as,

$\mathbf{f}^{out}_{d,h,w}=\sum_{i=0}^{2}\sum_{j=0}^{2}\sum_{k=0}^{2}\mathbf{f}^{in}_{d-1+i,h-1+j,w-1+k}\times\mathcal{W}_{i,j,k},$
(3)

where $\mathcal{W}$ is the weight of 3D convolution kernel, and
$\mathbf{f}^{in}$ and $\mathbf{f}^{out}$ are input and output features,
respectively.

Figure 3: The detailed structure of (a) Shallow Global Attention (SGA) block
which consists of spatial-mixing block and channel-mixing block, and (b) Deep
Global Attention (DGA) block which consists of Multi-head Self-Attention and
Feed-Forward layer. $C$ is the number of channels. $D$, $H$, $W$ are depth,
height and width of kernel. $N=D\times H\times W$ is the total number of
features.

Batch normalization layer: A batch Normalization (BN) layer is inserted to
mitigate this internal covariate shift [44] and speed up training. The BN
layer is achieved through a normalization step and a transformation step,

$\hat{x}=\gamma\frac{x-\mu}{\sqrt{\sigma^{2}+\epsilon}}+\beta,$ (4)

where $\mu$ and $\sigma$ are mean and variance for normalization, and $\gamma$
and $\beta$ are learnable parameters for scaling and shifting.

Nonlinear activation layer: We employ ReLU layer to provide nonlinear ability
in BrainFormer,

$\phi(x)=\max(0,x)=\left\\{\begin{matrix}x,\ \ \operatorname{if}\ x\geq 0,\\\
0,\ \ \operatorname{if}\ x<0.\\\ \end{matrix}\right.$ (5)

The positive values are kept and negative values are set to $0$.

####  3.3.3 Global Attention

3D CNNs only encode information within each local region, but are unable to
capture the correlation between distant brain regions. The self-attention is
widely used to formulate global relationship in computer vision and natural
language processing tasks [4, 2, 43]. Traditional self-attention learns global
relationship through computing fully connected attention mask for 1D/2D input
signal. However, the attention computation for 3D fMRI volume will suffer
unaffordable computation complexity, especially on the high resolution feature
in shallow layer. Towards more effective attention modeling for fMRI data, we
design different attention blocks for shallow and deep layers, respectively.

Shallow Global Attention: Fig. 3(a) depicts the detailed structure of SGA. It
accepts a sequence of volume features, shaped as “#seq$\times$channels” as
input, and uses two MLP blocks for spatial-mixing and channel-mixing,
respectively. The spatial-mixing block allows communication between different
spatial locations, by operating on each channel independently. The channel-
mixing block allows information communication between different channels.
These two types of blocks are interleaved to enable interaction among input
dimensions.

The input 3D feature is first flattened to $x\in\mathbb{R}^{N\times C}$, where
$C$ is the number of channels and $N=D\times H\times W$ is the resulting
number of features (the input sequence length). The first MLP block acts on
$x$-columns to mix global spatial information, and the second MLP block acts
on the $x$-rows of to mix channel information. Each MLP block contains two
fully connected layers with nonlinear mapping,

$\displaystyle x^{\prime}_{*,i}$
$\displaystyle=x_{*,i}+\mathbf{W}_{2}^{S}\phi(\mathbf{W}_{1}^{S}x_{*,i}),\ \
\operatorname{for}i=1...C,$ (6) $\displaystyle x^{\prime\prime}_{j,*}$
$\displaystyle=x^{\prime}_{j,*}+\mathbf{W}_{4}^{S}\phi(\mathbf{W}_{3}^{S}x^{\prime}_{j,*}),\
\ \operatorname{for}j=1...N,$

where $\phi$ is ReLU activation layer and $\mathbf{W}_{1}^{S}$,
$\mathbf{W}_{2}^{S}$, $\mathbf{W}_{3}^{S}$ and $\mathbf{W}_{4}^{S}$ are the
weights of the fully connected layers. Both MLP blocks are inserted with
residual connection. $\mathbf{W}_{4}^{S}$ is initialized to zero to retain the
original initialization. For an input $x\in\mathbb{R}^{N\times C}$, the SGA
reduces the computational complexity of self-attention from
$\mathcal{O}(N^{2})$ to $\mathcal{O}(N)$, thus relieving the memory and
computation costs in shallow layers where the feature resolution is high.

Deep Global Attention: DGA is designed to fuse the global information in deep
layer. Fig. 3(b) depicts the detailed DGA structure. A DGA block consists of a
Multi-head Self-Attention (MSA) and a Feed Forward (FF) layer. Considering
that the computation complexity of MSA is $\mathcal{O}(N^{2})$, we only insert
it into deep layers. Similarly, the DGA also flattens the input to
$x\in\mathbb{R}^{N\times C}$. Extra position embedding
$\mathbf{E}_{pos}\in\mathbb{R}^{N\times C}$ is added to volume features to
retain the positional information. The resulting feature sequence serves as
the input to the DGA block. The computation of the DGA block can be formulated
as,

$\displaystyle
z_{0}=x+\mathbf{E}_{pos},~{}z_{1}=\operatorname{MSA}(z_{0})+z_{0},~{}z_{2}=\operatorname{FF}(z_{1})+z_{1}.$
(7)

An MSA layer consists of multiple Self-Attentions (SA). A standard SA can be
formulated as mapping a query $\mathbf{q}$, key $\mathbf{k}$ and value
$\mathbf{v}$ to an output, where queries, keys, values, and outputs are all
volume features. The output is computed as a weighted sum of values, where the
weight for each value is computed by a similarity function between query and
key.

Figure 4: Visualization of (a) Functional Connectivity (FC) features with
different partitions, and (b-f) attention mask of SGA and DGA in different
layers..

For input features $z\in\mathbb{R}^{N\times C}$, SA first computes
$\mathbf{q}$, $\mathbf{k}$, and $\mathbf{v}$ through fully connected layer.
The attention mask $\mathbf{M}$ is then computed through pairwise similarity
between $\mathbf{q}$ and $\mathbf{k}$,

$\displaystyle=\mathbf{W}_{1}^{D}z,~{}\mathbf{M}=\operatorname{softmax}(\mathbf{q}\mathbf{k}^{T}/\sqrt{C_{h}}),~{}SA(\mathbf{v})$
$\displaystyle=\mathbf{M}\mathbf{v},$ (8)

where $\mathbf{W}_{1}^{D}\in\mathbb{R}^{C\times 3C_{h}}$ is the learnable
weight matrix in fully connected layer, $C_{h}$ is the channel number of
$\mathbf{q},\mathbf{k}$ and $\mathbf{v}$, $\sqrt{C_{h}}$ is scale factor for
normalization, and $\mathbf{M}\in\mathbb{R}^{N\times N}$ is the attention
mask. The Softmax function is used to normalize the attention mask to
probabilities. The MSA is an extension of the SA in which $k$ SAs are computed
in parallel,

$\operatorname{MSA}(z_{0})=[\operatorname{SA}_{1}(z_{0}),...,\operatorname{SA}_{k}(z_{0})]\mathbf{W}_{2}^{D},$
(9)

where $\mathbf{W}_{2}^{D}\in\mathbb{R}^{kC_{h}\times C}$ is the weight matrix
of the linear layer. The $\mathbf{W}_{2}^{D}$ is initialized to zero to retain
the original initialization. Finally, the weighted sum of value is added to
original input,

$z_{1}=z_{0}+\operatorname{MSA}(z_{0}),$ (10)

where $z_{1}\in\mathbb{R}^{N\times C}$ is the MSA output. MSA can effectively
capture the global correlation between distant brain regions.

Besides MSA, each DGA block contains an extra FF layer. An FF layer consists
of two fully connected layers with ReLU activation, and the output of FF layer
is added to input feature with residual manner,

$\displaystyle FF(z_{1})$
$\displaystyle=max(0,\mathbf{W}_{3}^{D}z_{1}+b_{1})\mathbf{W}_{4}^{D}+b_{2},$
(11) $\displaystyle z_{2}$ $\displaystyle=z_{1}+FF(z_{1}),$

where $\mathbf{W}_{3}^{D}$ and $\mathbf{W}_{4}^{D}$ are weights of fully
connected layers in FF, and $z_{2}$ is the final output of DGA block. We
visualize the FC feature and the attention mask in SGA and DGA in Fig. 4. In
comparison with FC computed by Pearson correlation between the mean time
series of predefined ROIs, the attention map in SGA and DGA automatically
learns correlation between finer brain voxel through end-to-end training,
hence could learn more reliable correlations.

TABLE I: Evaluation of individual components in proposed method Name | output size | layers
---|---|---
input | $64\times 72\times 64$ | -
stem conv | $32\times 36\times 32$ |  $7\times 7\times 7$, 64, stride 2
conv block1 | $32\times 36\times 32$ | $\begin{bmatrix}3\times 3\times 3,64\\\ 3\times 3\times 3,64\end{bmatrix}\times 2$
SGA block
conv block2 | $16\times 18\times 16$ | $\begin{bmatrix}3\times 3\times 3,128\\\ 3\times 3\times 3,128\end{bmatrix}\times 2$
SGA block
conv block3 | $8\times 9\times 8$ | $\begin{bmatrix}3\times 3\times 3,256\\\ 3\times 3\times 3,256\end{bmatrix}\times 2$
DGA block
conv block4 | $8\times 9\times 8$ | $\begin{bmatrix}3\times 3\times 3,512\\\ 3\times 3\times 3,512\end{bmatrix}\times 2$
DGA block
classifier | $1\times 1\times 1$ |  average pool, $512\to c$ fc, softmax
SGA:Shallow Global Attention, DGA:Deep Global Attention

####  3.3.4 Classifier

The output feature map is pooled with global average pooling and then
classified with classifier. The classifier consists of a fully connected layer
to map the features to class index, and a softmax function that normalizes the
predicted probability to $[0,1]$. We choose cross-entropy loss to optimize our
model,

$\mathcal{L}_{ce}=-\sum_{i=1}^{n}l_{i}\log(p_{i}),\ \
p_{i}=\frac{e^{\textbf{W}x_{i}}}{\sum_{j=1}^{n}e^{\textbf{W}x_{j}}},$ (12)

where $n$ is the number of classes, $l_{i}$ is one-hot label for
classification, W is the weight of fully connected layer, and $p_{i}$ is the
normalized classification score for the $i$-th class.

####  3.3.5 Multimodal fusion

With the designed data normalization layer, 3D CNN and global attention
blocks, we build BrainFormer based on 3D ResNet18 [6]. The overall BrainFormer
model is illustrated in Fig. 2(a) and the detailed structure is summarized in
Table I. Based on BrainFormer, we design the multimodal framework in Fig.
2(b). The features from four data modalities are concatenated for final
classification.

TABLE II: Summary of five brain disease datasets Dataset | #individual | #healthy control | #diseased | #class | disease category | #site | volume number
---|---|---|---|---|---|---|---
MPILMBB | 426 | 198 | 228 | 2 | MDD | 1 | 657
ADNI | 448 | 167 | 127AD, 154MCI | 3 | AD, MCI | 55 | 145
ABIDE | 988 | 556 | 339ASD, 93AS | 3 | ASD, AS | 17 | 82-320 (avg:194)
ADHD-200 | 1,151 | 842 | 309 | 2 | ADHD | 8 | 76-261 (avg:162)
ECHO | 187 | 52 | 107MD, 38MOH | 3 | MD, MOH | 2 | 240
MDD: major depressive disorder, AD: Alzheimer’s disease, MCI: mild cognitive
impairment
ASD: autism spectrum disorder, AS: Asperger’s syndrome, ADHD: attention
deficit hyperactivity disorder
MD: migraine disease, MOH: medication-overuse headache

Figure 5: (a) Multimodal comparison on ABIDE and ADNI datasets. (b) Evaluation
of individual components, and (c) Performance on each disease class on ECHO
dataset

##  4 Experiments

###  4.1 Dataset

We evaluate our method on five independently acquired neuroimaging datasets in
Table II.

MPILMBB dataset [45] contains 426 fMRI sequences from 318 individuals,
including 228 with MDD and 198 healthy controls (HC). Individuals who
underwent more than one scan are labeled as either depression or HC. The
dataset can be obtained from the MPI-Leipzig Mind-Brain-Body project111MPILMBB
dataset is available on https://www.neuroconnlab.org/data/.

ADNI dataset [46] includes subjects with AD, MCI, and HC from 55 study sites.
We filter out samples with irregular shapes, and finally get 448 fMRI
sequences. Most prior studies conducted binary classifications (AD _vs._ HC,
or MCI _vs._ HC). Here, we report both binary and 3-class classification
performances. The ADNI dataset can be accessed from the Alzheimer’s Disease
Neuroimaging Initiative project222ADNI dataset is available on
https://adni.loni.usc.edu.

ABIDE dataset [47] contains fMRI data of 988 subjects acquired in 17
neuroimaging centers worldwide with different imaging protocols. The dataset
includes 339 individuals with ASD, 93 with AS, and 546 HC. Most previous
studies conducted binary classification (ASD _vs._ HC). Here, we report both
binary and 3-class performance. The ABIDE dataset can be accessed from Autism
Brain Imaging Data Exchange project333ABIDE dataset is available on
http://preprocessed-connectomes-project.org/abide/index.html.

ADHD-200 dataset444ADHD-200 dataset is available on
http://fcon_1000.projects.nitrc.org/indi/adhd200/ [48] contains fMRI data of
1,151 subjects, including 309 with ADHD and 842 HC. The dataset is acquired
from eight different neuroimaging centers worldwide.

ECHO We have an ongoing project for constructing a new fMRI dataset Evaluation
and Classification of Headache disOrders using MRI, (ECHO) from 12 study
sites. Two types of headache from two sites in ECHO dataset, including MD,
MOH, as well as HC are included in the analysis, which consists 107, 38 and 52
individuals, respectively. The study was approved by the Ethics Committee of
Chinese PLA General Hospital and registered at Chinese Clinical Trial Registry
(ChiCTR2200056871).

###  4.2 Implementation Details

####  4.2.1 Data Preprocessing

Before training, we preprocess all fMRI data (except ABIDE and ADHD-200) using
DPARSF toolset555http://rfmri.org/DPARSF based on MATLAB2013b (8.2.0.701). The
steps are as follows: (1) remove the first five volumes; (2) slice timing
correction; (3) realignment; (4) normalize to Montreal Neurological Institute
(MNI) standard space through DARTEL and resample to $3\times 3\times 3$ mm3
resolution; (5) spatial smoothing with a 6-mm full width at half maximum
(FWHM) Gaussian kernel. Above process results $T\times 61\times 73\times 61$
tensor for each fMRI sequence, where $T$ is the length of sequence. For the
ABIDE dataset, the Configurable Pipeline for the Analysis of Connectomes
(CPAC) minimally preprocessed data [47] is used. For the ADHD-200 dataset, the
Athena preprocessed data [48] is used. All sMRI data are in nifty format and
trained as they are.

####  4.2.2 Training

We apply 3D ResNet18 [41] as backbone feature extractor and initialize it with
weights pretrained on ImageNet [1]. All our models are trained with Adam
optimizer for 10 epochs, and each batch contains 16 3D volumes. The initial
learning rate is set to 0.0001, and is reduced by factor of 10 at the 8-th
epoch. The 3D fMRI and sMRI volumes are padded to $64\times 72\times 64$
before input to model. We flatten the CC200 [47] as FC feature, and the scalar
phenotypic data is concatenated as one vector before input into model. All
models are trained with only cross-entropy loss.

####  4.2.3 Testing

We use 5-fold cross validation to verify our method on MPILMBB, ADNI, ABIDE
and ECHO datasets, and report the average performances. For participants who
have multiple scans in one dataset, the splitting ensures the data of the same
individual enter either training or testing data. As the ADHD-200 dataset
provides split of training and testing data, we follow the recommended
protocol when verifying our method on this dataset.

TABLE III: The classification accuracy with multimodal data on five brain disease datasets  Modal | MPILMBB | ADNI | ADHD | ABIDE | ECHO
---|---|---|---|---|---
fMRI | 97.2 | 95.3 | 72.0 | 72.5 | 81.2
fMRI+FC | 97.2 | 95.2 | 72.2 | 72.8 | 81.4
fMRI+PD | 97.3 | 95.4 | 72.4 | 72.9 | 81.4
fMRI+sMRI | 97.2 | 95.8 | 72.2 | 73.3 | 81.5
fMRI+sMRI+FC+PD | 97.3 | 95.9 | 72.3 | 73.3 | 81.5
FC: functional connectivity, PD: phenotypic data |
TABLE IV: Evaluation of individual components in BrainFormer Method | MPILMBB | ADNI | ABIDE | ADHD | ECHO
---|---|---|---|---|---
BrainFormer | 97.2 | 95.3 | 72.0 | 72.5 | 81.2
w/o data norm | 94.2 | 93.8 | 58.7 | 56.7 | 79.3
w/o SGA | 95.3 | 93.4 | 68.7 | 69.2 | 75.0
w/o DGA | 94.7 | 92.2 | 67.9 | 68.5 | 75.9
w/o SGA, DGA | 94.1 | 90.6 | 65.5 | 67.1 | 73.9

###  4.3 Multimodal experiments

In this section, we first verify the effectiveness of each modality for brain
disease classification on five dataset. We compare our model trained with
different data modalities, including fMRI volumes, sMRI volume, FC and scalar
phenotypic information. The experimental results are summarized in Table III.
It can be observed from table that, fusing multiple modalities can improve the
classification accuracy. For example, introducing sMRI leads to the improved
classification accuracy on ADNI dataset by 0.5%. The reason may be that AD
patients exhibit brain structural changes that can also be detected by sMRI.
Nevertheless, introducing FC feature does not significantly improve
classification accuracy on the other five datasets. This result again supports
our main hypothesis that FC neglects the detailed information in each voxel
that is important to identify the disease, and cannot compensate our model by
using single fMRI volume. When combining all the four modalities, our method
achieves the best performance on all the five dataset. This proves the
effectiveness of multimodal fusion for brain disease classification.

We further compare the proposed methods with recent multimodal methods on
ABIDE and ADNI dataset. The result are summarized in Fig. 5(a). The compared
methods include Ridge Classifier [37], MLP [37], MMGL [11], BNC-DGHL [38],
Spectral-GCN [27] and DNN-JFC [37]. It can be observed from the figure that,
our method achieves the best performance among all the compared methods.

###  4.4 Ablation Study

####  4.4.1 Individual Component

We first investigate the effectiveness of each component in BrainFormer,
including the data norm layer, SGA and DGA. The experimental results on five
datasets are summarized in Table. IV and Fig. 5(b). In the table,
“BrainFormer” denotes the complete BrainFormer model, and “w/o” refers to
discard the corresponding component.

It can be observed from figure that, without data normalization, BrainFormer
suffers marked performance drop on all five datasets. This shows our data
normalization layer can effectively unify the distribution of multi-site data.
The largest performance drops can be observed on ABIDE and ADHD-200 datasets
(13.3% and 15.8% respectively). Presumably the reason is that ABIDE and
ADHD-200 datasets are collected from more study sites, thus benefit more from
data normalization. Table. IV also shows that both SGA and DGA significantly
boost the performance. By combining SGA and DGA, BrainFormer achieves the best
performance, _e.g._ , Acc=81.2% on ECHO dataset. We conclude that both SGA and
DGA are effective to model the correlation between distant brain areas,
thereby enhancing performance, and are optimized when used in combination.
Global correlation learning is complementary between shallow and deep layers .

####  4.4.2 Class-wise Analysis on ECHO dataset

The ECHO dataset contains 3 classes, _e.g._ , MD, MOH, HC, and the overall
metric cannot measure the performance for each class. Therefore, we evaluate
class-wise performance on ECHO dataset in Fig. 5(c). It can be observed that
our method achieves superiority recall and fine accuracy on MD class. This
shows that our method can effectively detect migraine while prevent missed
diagnosis in clinical practice. Fig. 5(c) also shows that our method achieves
93.1% accuracy on HC class, indicating its ability to obtain normal range.

TABLE V: Evaluation of attention blocks. Attn block | MPI | ADNI | ABIDE | ADHD | ECHO | Mem (GB) | Gflops
---|---|---|---|---|---|---|---
S-S-S-S | 93.1 | 88.9 | 64.3 | 66.0 | 71.5 | 0.2 | 34.2
S-S-S-D | 94.1 | 90.6 | 65.5 | 67.1 | 73.9 | 0.3 | 34.5
S-S-D-D | 95.3 | 93.4 | 68.7 | 69.2 | 75.0 | 0.3 | 34.9
S-D-D-D | 92.7 | 91.2 | 66.9 | 66.5 | 73.9 | 1.9 | 77.1
D-D-D-D | 90.2 | 89.8 | 61.7 | 62.5 | 71.9 | 2.3 | 119.9

####  4.4.3 Evaluation of attention blocks

In this part, we evaluate the two kinds of attention blocks at different
layers. The evaluation results are summarised in Table V. In the table, “S”
denotes we add SGA block in corresponding layer, and “D” denotes we insert DGA
block. It can be observed from the table that, using DGA block in shallow
layer will significantly increase the memory cost and computation. which will
also increase the optimization difficulty. While inserting SGA block in deep
layer doesn’t significantly influence the memory cost and computation, but
reduces the performance of BrainFormer. Based on this observation, we use
“S-S-D-D” in BrainFormer.

TABLE VI: The classification accuracy of different methods on five brain disease datasets. Methods | MPILMBB | ADNI | ABIDE | ADHD | ECHO
---|---|---|---|---|---
CNN-based | 2D CNN | 51.2$\pm$2.3 | 53.4$\pm$2.7 | 47.9$\pm$3.4 | 49.7$\pm$4.2 | 31.2$\pm$3.6
3D CNN | 92.1$\pm$3.2 | 91.5$\pm$4.8 | 65.3$\pm$3.3 | 61.2$\pm$5.6 | 73.9$\pm$2.7
3D CNNA | 92.3$\pm$3.3 | 92.4$\pm$3.9 | 65.9$\pm$3.5 | 61.2$\pm$5.6 | 73.9$\pm$2.7
SVIG [7] | 92.2$\pm$1.7 | 91.1$\pm$2.4 | 67.7$\pm$4.4 | 66.0$\pm$2.0 | 74.1$\pm$3.0
GCN-based | PopGCNA [9] | 83.2$\pm$2.7 | 81.3$\pm$3.2 | 57.6$\pm$4.1 | 60.2$\pm$4.6 | 65.6$\pm$5.3
EVGCNA [10] | 90.7$\pm$1.9 | 88.1$\pm$4.3 | 63.7$\pm$4.6 | 64.1$\pm$3.9 | 70.8$\pm$3.1
MMGLA [11] | 92.3$\pm$0.9 | 91.1$\pm$1.7 | 68.4$\pm$5.3 | 66.6$\pm$4.6 | 74.1$\pm$4.3
Transformer-based | 3D ViTA [20] | 89.2$\pm$2.9 | 87.1$\pm$3.8 | 60.4$\pm$1.9 | 55.6$\pm$4.3 | 68.0$\pm$5.6
MT-UNet* [12] | 85.4$\pm$3.5 | 79.8$\pm$3.4 | 52.5$\pm$4.6 | 55.0$\pm$4.6 | 67.6$\pm$6.0
UNETR* [13] | 87.1$\pm$2.1 | 82.4$\pm$4.1 | 57.1$\pm$2.5 | 60.5$\pm$3.8 | 68.2$\pm$4.1
SwinUNETR* [14] | 89.3$\pm$2.5 | 84.0$\pm$1.5 | 58.2$\pm$3.8 | 64.0$\pm$4.2 | 68.6$\pm$3.9
BrainAgingNet [17] | 91.8$\pm$1.9 | 90.6$\pm$3.1 | 65.1$\pm$4.9 | 67.1$\pm$4.4 | 73.5$\pm$6.0
TransBTS* [15] | 94.1$\pm$2.1 | 92.5$\pm$1.3 | 70.1$\pm$5.7 | 68.4$\pm$3.7 | 77.0$\pm$4.9
| KD-TransformerA [5] | 95.1$\pm$1.2 | 78.1$\pm$1.9 | 60.6$\pm$2.9 | 70.7$\pm$1.3 | 78.4$\pm$2.7
Ours | BrainFormer | 97.2$\pm$0.4 | 95.3$\pm$0.6 | 72.5$\pm$1.4 | 72.0$\pm$2.3 | 81.2$\pm$2.7
BrainFormerA | 97.3$\pm$0.3 | 95.9$\pm$0.4 | 73.3$\pm$1.1 | 72.3$\pm$1.9 | 81.5$\pm$2.2
The CNN-based methods directly take 2D/3D fMRI data as input for
classification.
The GCN-based methods take FC feature as input and learn correlation between
different brain ROIs with GCNs.
The Transformer-based methods introduce transformer to learn the correlation
between different brain areas.
* denotes the method is originally designed for medical image segmentation.
A denotes the methods use multimodal data for training and testing.

Figure 6: Comparison with recent methods on MPI, ADNI, ABIDE and ADHD200
datasets.

###  4.5 Comparison

In this section, we compare BrainFormer with representative methods on all
five brain disease datasets.

####  4.5.1 Comparison with typical methods on five datasets

We first compare BrainFormer with some typical methods used for medical image
processing in Table 6, including CNN-based methods, GCN-based methods and
Transformer-based methods. We repeat the methods with the codes provided by
the authors in Table 6 with our train/test protocol mentioned in Sec. 4.2.
SVIG [7] transforms 3D fMRI volumes into 2D images and uses 2D CNN for
prediction. The GCN-based methods use FC features as input instead of the
original 3D fMRI volumes. Considering there are very few Transformer-based
methods proposed for fMRI data classification, we modify medical image
segmentation methods including MT-UNet [12], UNETR [13], SwinUNETR [14] and
TransBTS [15] for comparison. Since these methods are originally designed for
sMRI data segmentation, they may show unsatisfactory performance for fMRI
classification. The methods use multimodal data are annotated.

Table 6 shows that 2D CNN exhibits inferior performance, since that 2D CNN
cannot capture the 3D structure inside 3D volume. The performance of GCN-based
methods is lower than transformer-based methods even though the former
introduce extra multimodal information on age, sex, site-id and so on, since
FC neglects the detailed information contained in each voxel of 3D volumes.
The Transformer-based methods achieve reasonable performance on all five
datasets. BrainFormer outperforms all compared methods. For example,
BrainFormer outperforms TransBTS [15] by 4.2%. Our method is able to capture
both local details and long-range relationships from fMRI and sMRI volumes,
and also effectively fuse the multimodal data, it hence achieves promising
performance.

####  4.5.2 MPILMBB

Fig. 6(a) compares BrainFormer with recent studies on the MPILMBB dataset.
BrainFormer outperforms MCD [49] by a large margin.

####  4.5.3 ADNI

As in Fig. 6(b), BrainFormer is compared with CNN-based methods such as
DeepFusion [50], DenseCNN [51], CMTL [52] and Cascaded CNN [53]. Our method
achieves the best performance on ADNI dataset. It achieves 95.3% accuracy for
binary classification (AD _vs._ CN), outperforming state-of-the-art Cascaded
CNN [53] by 2.1%. Cascaded CNN [53] also uses 3D CNN for fMRI classification.
Our method is equipped with two global attention blocks, and achieves better
results with global correlation learning ability.

####  4.5.4 ABIDE

As in Fig. 6(c), we compare our method with SVM [26], MLP-based method
including Autoencoder+MLP [26], FCNN [26] and ASD-DiagNet [24]. BrainFormer
achieves 72.5% of accuracy, outperforming current state-of-the-art ASD-DiagNet
[24]. Note that our method does not apply extra data augmentation or auto-
encoder for pretraining.

####  4.5.5 ADHD-200

In Fig. 6(d), the compared methods include CNN-based method, _e.g._ , 3D CNN
[54], SC-CNN [55], and GCN-based methods MMTGCN [56] and MTGCN [56].
BrainFormer outperforms MTGCN [56], and MMTGCN [56] on both accuracy and
precision, which further demonstrates its transferability for multi-site fMRI
analysis.

####  4.5.6 3-class classification

We implement 3-class classification on ADNI, ABIDE and ECHO dataset in Fig.
6(e). It can be observed from figure that BrainFormer also achieves promising
performance for 3-class classification, which further demonstrates its
superiority.

####  4.5.7 Training curve

We visualize the training loss and accuracy curves in Fig. 7. It can be
observed from the figure that, our method converges quickly during training.

Figure 7: Training curve of BrainFormer on ECHO dataset.
Figure 8: Visualization results of disease-related biomarkers achieved by GL.
The activation map highlights brain regions that are associated with disease
classification. (a) from MPILMBB, (b) from ADNI, (c) from ABIDE, (d) from
ADHD-200, (e) and (f) from ECHO.

###  4.6 Visualization

Fig. 8 uses Gradient-based Localization (GL) [57] to visualize the disease-
related biomarkers obtained from the five datasets. The greatest advantage of
BrainFormer is that it directly locates the possible disease biomarkers on 3D
brain volumes and provides probability distribution. As for traditional FC,
brain regions are predefined and the mean BOLD values are computed, making it
a blurring method to locate biomarkers. All regions illustrated in Fig. 8 have
frequently been associated with certain diseases. In terms of MDD, the
critical region for classification locates mainly in the prefrontal cortex
[58]. Regions that may identify AD are located in the frontal, parieto-
temporal, and occipital cortices [59]. In terms of ASD, the prefrontal,
temporal, and insular lobes and part of the cerebellum show high response
[60]. Regions that may identify MOH locate in the dorsal anterior cingulate
cortex, midbrain and cerebellum [61]. For ADHD, the critical regions for
classification locate in the occipital lobe and cerebellum [62]. MD-related
regions are brainstem and cerebellum [63].

##  5 Discussion and Conclusion

We propose BrainFormer, a hybrid Transformer architecture working on the
single fMRI volume for universal multimodal brain disease classification.
BrainFormer shows solid generalizability for multiple brain diseases
classification with multimodal data, and good transferability for multi-site
data using different scanning parameters. The visualization results
demonstrate the superiority of BrainFormer to directly locate possible
disease-related biomarker on fMRI image. The results are achieved by first
modeling the local cues within each brain region through 3D CNN, and capturing
the global correlations among distant regions through two global attention
blocks, and then applying a data normalization layer to handle the multi-site
data. For multimodal data classification, we use BrainFormer to process MRI
images, and the FC feature and phenotypic data are processed by a 3-layer MLP,
and combine all the data modalities for classification. Besides, we further
introduce gradient-based localization to locate disease-related biomarker. The
results demonstrate that BrainFormer, working on single fMRI volume, performs
quite well on five diseases and combining multiple feature modalities to
achieve even better classification performance.

Towards better clinical practice, BrainFormer provides an ROI-labeling-free
choice to avoid hand-crafted bias in fMRI brain disease classification. With
the input of 3D single fMRI images, our method expands the sample size by
hundreds of times for deep learning application. This can largely reduce the
afford in collecting neuroimaging data, which has been rather time-consuming
and labour-expensive. BrainFormer can precisely locate the disease-related
biomarkers on fMRI volume, which also motivates future studies on fMRI data
analysis.

We discuss limitations as follows. First, we apply a simple normalization
layer to handle multi-site data, which could be further optimized by better
design. Second, although BrainFormer can classify multiple diseases, it can
only handle one disease from a spectrum of them per training. In the future,
we aim to study a more unified framework to classify multiple diseases by
multi-task learning.

Acknowledgments This work is supported xxx.

## References

  * [1] A. Krizhevsky, I. Sutskever, and G. E. Hinton, “Imagenet classification with deep convolutional neural networks,” _Advances in neural information processing systems_ , vol. 25, 2012.
  * [2] N. Carion, F. Massa, G. Synnaeve, N. Usunier, A. Kirillov, and S. Zagoruyko, “End-to-end object detection with transformers,” in _European conference on computer vision_. Springer, 2020.
  * [3] X. Wang, R. Girshick, A. Gupta, and K. He, “Non-local neural networks,” in _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2018, pp. 7794–7803.
  * [4] A. Vaswani _et al._ , “Attention is all you need,” _Advances in neural information processing systems_ , vol. 30, 2017.
  * [5] J. Zhang, L. Zhou, L. Wang, M. Liu, and D. Shen, “Diffusion kernel attention network for brain disorder classification,” _IEEE Transactions on Medical Imaging_ , 2022.
  * [6] K. He, X. Zhang, S. Ren, and J. Sun, “Deep residual learning for image recognition,” in _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2016, pp. 770–778.
  * [7] M. R. Ahmed, Y. Zhang, Y. Liu, and H. Liao, “Single volume image generator and deep learning-based asd classification,” _IEEE Journal of Biomedical and Health Informatics_ , vol. 24, no. 11, pp. 3044–3054, 2020.
  * [8] O. Ronneberger, P. Fischer, and T. Brox, “U-net: Convolutional networks for biomedical image segmentation,” in _International Conference on Medical image computing and computer-assisted intervention_. Springer, 2015, pp. 234–241.
  * [9] S. Parisot, S. I. Ktena, E. Ferrante, M. Lee, R. G. Moreno, B. Glocker, and D. Rueckert, “Spectral graph convolutions for population-based disease prediction,” in _International conference on medical image computing and computer-assisted intervention_. Springer, 2017, pp. 177–185.
  * [10] Y. Huang and A. Chung, “Edge-variational graph convolutional networks for uncertainty-aware disease prediction,” in _International Conference on Medical Image Computing and Computer-Assisted Intervention_. Springer, 2020, pp. 562–572.
  * [11] S. Zheng, Z. Zhu, Z. Liu, Z. Guo, Y. Liu, Y. Yang, and Y. Zhao, “Multi-modal graph learning for disease prediction,” _IEEE Transactions on Medical Imaging_ , 2022.
  * [12] H. Wang, S. Xie, L. Lin, Y. Iwamoto, X.-H. Han, Y.-W. Chen, and R. Tong, “Mixed transformer u-net for medical image segmentation,” in _ICASSP 2022-2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_. IEEE, 2022, pp. 2390–2394.
  * [13] Y. Gao, M. Zhou, and D. N. Metaxas, “Utnet: a hybrid transformer architecture for medical image segmentation,” in _International Conference on Medical Image Computing and Computer-Assisted Intervention_. Springer, 2021, pp. 61–71.
  * [14] A. Hatamizadeh, V. Nath, Y. Tang, D. Yang, H. R. Roth, and D. Xu, “Swin unetr: Swin transformers for semantic segmentation of brain tumors in mri images,” in _International MICCAI Brainlesion Workshop_. Springer, 2022, pp. 272–284.
  * [15] W. Wang, C. Chen, M. Ding, H. Yu, S. Zha, and J. Li, “Transbts: Multimodal brain tumor segmentation using transformer,” in _International Conference on Medical Image Computing and Computer-Assisted Intervention_. Springer, 2021, pp. 109–119.
  * [16] O. Dalmaz, M. Yurt, and T. Çukur, “Resvit: residual vision transformers for multimodal medical image synthesis,” _IEEE Transactions on Medical Imaging_ , vol. 41, no. 10, pp. 2598–2614, 2022.
  * [17] S. He, P. E. Grant, and Y. Ou, “Global-local transformer for brain age estimation,” _IEEE Transactions on Medical Imaging_ , vol. 41, no. 1, pp. 213–224, 2021.
  * [18] J. Chen, E. C. Frey, Y. He, W. P. Segars, Y. Li, and Y. Du, “Transmorph: Transformer for unsupervised medical image registration,” _Medical Image Analysis_ , vol. 82, p. 102615, 2022.
  * [19] W. Tang, F. He, Y. Liu, and Y. Duan, “Matr: Multimodal medical image fusion via multiscale adaptive transformer,” _IEEE Transactions on Image Processing_ , vol. 31, pp. 5134–5149, 2022.
  * [20] A. Dosovitskiy _et al._ , “An image is worth 16x16 words: Transformers for image recognition at scale,” _arXiv preprint arXiv:2010.11929_ , 2020.
  * [21] C. D. Billones, O. J. L. D. Demetria, D. E. D. Hostallero, and P. C. Naval, “Demnet: a convolutional neural network for the detection of alzheimer’s disease and mild cognitive impairment,” in _2016 IEEE region 10 conference (TENCON)_. IEEE, 2016, pp. 3724–3727.
  * [22] M. Khosla, K. Jamison, A. Kuceyeski, and M. R. Sabuncu, “Ensemble learning with 3d convolutional neural networks for functional connectome-based prediction,” _NeuroImage_ , vol. 199, pp. 651–662, 2019.
  * [23] Y. Tang, C. Wang, Y. Chen, N. Sun, A. Jiang, and Z. Wang, “Identifying adhd individuals from resting-state functional connectivity using subspace clustering and binary hypothesis testing,” _Journal of attention disorders_ , vol. 25, no. 5, pp. 736–748, 2021.
  * [24] T. Eslami, V. Mirjalili, A. Fong, A. R. Laird, and F. Saeed, “Asd-diagnet: a hybrid learning approach for detection of autism spectrum disorder using fmri data,” _Frontiers in neuroinformatics_ , vol. 13, p. 70, 2019.
  * [25] M. Chen, H. Li, J. Wang, J. R. Dillman, N. A. Parikh, and L. He, “A multichannel deep neural network model analyzing multiscale functional brain connectome data for attention deficit hyperactivity disorder detection,” _Radiology: Artificial Intelligence_ , vol. 2, p. e190012, 2019.
  * [26] J. Hu, L. Cao, T. Li, B. Liao, S. Dong, and P. Li, “Interpretable learning approaches in resting-state functional connectivity analysis: the case of autism spectrum disorder,” _Computational and Mathematical Methods in Medicine_ , vol. 2020, 2020.
  * [27] S. Parisot, S. I. Ktena, E. Ferrante, M. Lee, R. Guerrero, B. Glocker, and D. Rueckert, “Disease prediction using graph convolutional networks: application to autism spectrum disorder and alzheimer’s disease,” _Medical image analysis_ , vol. 48, pp. 117–130, 2018.
  * [28] D. Tran, L. Bourdev, R. Fergus, L. Torresani, and M. Paluri, “Learning spatiotemporal features with 3d convolutional networks,” in _Proceedings of the IEEE international conference on computer vision_ , 2015\.
  * [29] J. Li, S. Zhang, and T. Huang, “Multi-scale 3d convolution network for video based person re-identification,” in _Proceedings of the AAAI Conference on Artificial Intelligence_ , vol. 33, 2019, pp. 8618–8625.
  * [30] ——, “Multi-scale temporal cues learning for video person re-identification,” _IEEE Transactions on Image Processing_ , vol. 29, pp. 4461–4473, 2020.
  * [31] J. Wang, Q. Wang, H. Zhang, J. Chen, S. Wang, and D. Shen, “Sparse multiview task-centralized ensemble learning for asd diagnosis based on age-and sex-related functional connectivity patterns,” _IEEE transactions on cybernetics_ , vol. 49, no. 8, pp. 3141–3154, 2018.
  * [32] J. Wang _et al._ , “Multi-class asd classification based on functional connectivity and functional correlation tensor via multi-source domain adaptation and multi-view sparse representation,” _IEEE transactions on medical imaging_ , vol. 39, no. 10, pp. 3137–3147, 2020.
  * [33] ——, “Multi-class asd classification via label distribution learning with class-shared and class-specific decomposition,” _Medical Image Analysis_ , vol. 75, p. 102294, 2022.
  * [34] F. Huang, E.-L. Tan, P. Yang, S. Huang, L. Ou-Yang, J. Cao, T. Wang, and B. Lei, “Self-weighted adaptive structure learning for asd diagnosis via multi-template multi-center representation,” _Medical Image Analysis_ , vol. 63, p. 101662, 2020.
  * [35] H. Vu, H.-C. Kim, M. Jung, and J.-H. Lee, “fmri volume classification using a 3d convolutional neural network robust to shifted and scaled neuronal activations,” _NeuroImage_ , vol. 223, p. 117328, 2020.
  * [36] R. Liu, Z.-A. Huang, Y. Hu, Z. Zhu, K.-C. Wong, and K. C. Tan, “Attention-like multimodality fusion with data augmentation for diagnosis of mental disorders using mri,” _IEEE Transactions on Neural Networks and Learning Systems_ , 2022.
  * [37] Y. Huang and A. C. Chung, “Disease prediction with edge-variational graph convolutional networks,” _Medical Image Analysis_ , vol. 77, p. 102375, 2022\.
  * [38] J. Ji and Y. Zhang, “Functional brain network classification based on deep graph hashing learning,” _IEEE Transactions on Medical Imaging_ , vol. 41, no. 10, pp. 2891–2902, 2022.
  * [39] K. Supekar, S. Ryali, R. Yuan, D. Kumar, C. de Los Angeles, and V. Menon, “Robust, generalizable, and interpretable artificial intelligence–derived brain fingerprints of autism and social communication symptom severity,” _Biological Psychiatry_ , vol. 92, no. 8, pp. 643–653, 2022.
  * [40] K. Sun, Z. Liu, G. Chen, Z. Zhou, S. Zhong, Z. Tang, S. Wang, G. Zhou, X. Zhou, L. Shao _et al._ , “A two-center radiomic analysis for differentiating major depressive disorder using multi-modality mri data under different parcellation methods,” _Journal of Affective Disorders_ , vol. 300, pp. 1–9, 2022.
  * [41] J. Carreira and A. Zisserman, “Quo vadis, action recognition? a new model and the kinetics dataset,” in _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2017.
  * [42] Z. Qiu, T. Yao, and T. Mei, “Learning spatio-temporal representation with pseudo-3d residual networks,” in _Proceedings of the IEEE international conference on computer vision_ , 2017.
  * [43] Z. Liu _et al._ , “Swin transformer: Hierarchical vision transformer using shifted windows,” in _Proceedings of the International Conference on Computer Vision_ , 2021, pp. 10 012–10 022.
  * [44] S. Ioffe and C. Szegedy, “Batch normalization: Accelerating deep network training by reducing internal covariate shift,” in _International conference on machine learning_. PMLR, 2015, pp. 448–456.
  * [45] N. Mendes _et al._ , “A functional connectome phenotyping dataset including cognitive state and personality measures,” _Scientific data_ , vol. 6, pp. 1–19, 2019.
  * [46] R. C. Petersen _et al._ , “Alzheimer’s disease neuroimaging initiative (adni): clinical characterization,” _Neurology_ , vol. 74, no. 3, pp. 201–209, 2010.
  * [47] C. Craddock _et al._ , “The neuro bureau preprocessing initiative: open sharing of preprocessed neuroimaging data and derivatives,” _Frontiers in Neuroinformatics_ , vol. 7, p. 27, 2013.
  * [48] P. Bellec, C. Chu, F. Chouinard-Decorte, Y. Benhajali, D. S. Margulies, and R. C. Craddock, “The neuro bureau adhd-200 preprocessed repository,” _Neuroimage_ , vol. 144, pp. 275–286, 2017.
  * [49] M. Mousavian, J. Chen, Z. Traylor, and S. Greening, “Depression detection from smri and rs-fmri images using machine learning,” _Journal of Intelligent Information Systems_ , vol. 57, no. 2, pp. 395–418, 2021.
  * [50] U. Senanayake, A. Sowmya, and L. Dawes, “Deep fusion pipeline for mild cognitive impairment diagnosis,” in _2018 IEEE 15th international symposium on biomedical imaging_. IEEE, 2018, pp. 1394–1997.
  * [51] F. Li _et al._ , “Alzheimer’s disease diagnosis based on multiple cluster dense convolutional networks,” _Computerized Medical Imaging and Graphics_ , vol. 70, pp. 101–110, 2018.
  * [52] K. Aderghal, A. Khvostikov, A. Krylov, J. Benois-Pineau, K. Afdel, and G. Catheline, “Classification of alzheimer disease on imaging modalities with deep cnns using cross-modal transfer learning,” in _2018 IEEE 31st international symposium on computer-based medical systems (CBMS)_. IEEE, 2018, pp. 345–350.
  * [53] M. Liu, D. Cheng, K. Wang, and Y. Wang, “Multi-modality cascaded convolutional neural networks for alzheimer’s disease diagnosis,” _Neuroinformatics_ , vol. 16, no. 3, pp. 295–308, 2018.
  * [54] L. Zou, J. Zheng, C. Miao, M. J. Mckeown, and Z. J. Wang, “3d cnn based automatic diagnosis of attention deficit hyperactivity disorder using functional and structural mri,” _IEEE Access_ , vol. 5, 2017.
  * [55] T. Zhang _et al._ , “Separated channel attention convolutional neural network (sc-cnn-attention) to identify adhd in multi-site rs-fmri dataset,” _Entropy_ , vol. 22, p. 893, 2020.
  * [56] D. Yao _et al._ , “A mutual multi-scale triplet graph convolutional network for classification of brain disorders using functional or structural connectivity,” _IEEE Transactions on Medical Imaging_ , vol. 40, no. 4, pp. 1279–1289, 2021.
  * [57] R. R. Selvaraju, M. Cogswell, A. Das, R. Vedantam, D. Parikh, and D. Batra, “Grad-cam: Visual explanations from deep networks via gradient-based localization,” in _Proceedings of the IEEE international conference on computer vision_ , 2017, pp. 618–626.
  * [58] D. A. Pizzagalli and A. C. Roberts, “Prefrontal cortex and depression,” _Neuropsychopharmacology_ , vol. 47, no. 1, pp. 225–246, 2022.
  * [59] I. M. McDonough, S. B. Festini, and M. M. Wood, “Risk for alzheimer’s disease: A review of long-term episodic memory encoding and retrieval fmri studies,” _Ageing Research Reviews_ , vol. 62, p. 101133, 2020.
  * [60] C. M. Murphy _et al._ , “Abnormal functional activation and maturation of ventromedial prefrontal cortex and cerebellum during temporal discounting in autism spectrum disorder,” _Human Brain Mapping_ , vol. 38, no. 11, pp. 5343–5355, 2017.
  * [61] W. Dai _et al._ , “Enhanced functional connectivity between habenula and salience network in medication-overuse headache complicating chronic migraine positions it within the addiction disorders: an ica-based resting-state fmri study,” _The Journal of Headache and Pain_ , vol. 22, no. 1, pp. 1–9, 2021\.
  * [62] A. Sathyanesan, J. Zhou, J. Scafidi, D. H. Heck, R. V. Sillitoe, and V. Gallo, “Emerging connections between cerebellar development, behaviour and complex brain disorders,” _Nature Reviews Neuroscience_ , vol. 20, no. 5, pp. 298–313, 2019.
  * [63] W. Dai _et al._ , “Cortical mechanisms in migraine,” _Molecular Pain_ , vol. 17, p. 17448069211050246, 2021.


# Thermalization in a Spin-Orbit coupled Bose gas by enhanced spin Coulomb
drag

D. J. Brown<EMAIL_ADDRESS>Dodd-Walls Centre for Photonic and Quantum
Technologies, Department of Physics, University of Auckland, Private Bag
92019, Auckland, New Zealand  Present address: Light-Matter Interactions for
Quantum Technologies Unit, Okinawa Institute of Science and Technology
Graduate University, Onna, Okinawa 904-0495, Japan  M. D. Hoogerland  Dodd-
Walls Centre for Photonic and Quantum Technologies, Department of Physics,
University of Auckland, Private Bag 92019, Auckland, New Zealand

###### Abstract

An important component of the structure of the atom, the effects of spin-orbit
coupling are present in many sub-fields of physics. Most of these effects are
present continuously. We present a detailed study of the dynamics of changing
the spin-orbit coupling in an ultra-cold Bose gas, coupling the motion of the
atoms to their spin. We find that the spin-orbit coupling greatly increases
the damping towards equilibrium. We interpret this damping as spin drag, which
is enhanced by spin-orbit coupling rate, scaled by a remarkable factor of
$8.9(6)$ s. We also find that spin-orbit coupling lowers the final temperature
of the Bose gas after thermalization.

##  I Introduction

The understanding of the transport, diffusion and damping of spin, in contrast
to those of charge, is important to the field of spintronics Wolf  _et al._
(2001), where the spin of particles, rather than the quantity of particles
(charge) carries information. Spin currents, in contrast to charge currents,
are damped due to collisions between particles of opposite spin, as their
relative momentum is not conserved. This damping is known as Spin Drag D’Amico
and Vignale (2000); Weber  _et al._ (2005).

Analogous to the spin drag in bilayer electron systems, systems of ultracold
bosons can also demonstrate spin drag, with the drag enhanced by the familiar
bosonic enhancement Duine and Stoof (2009) prominent in ultracold boson
systems. There is a detailed collection of work over the years into the
presence of spin Coulomb drag in ultracold atomic systems Sommer  _et al._
(2011); Goulko  _et al._ (2013); Fava  _et al._ (2018) with a small selection
of the work featuring the inclusion of spin-orbit coupling Li  _et al._
(2019).

The other important effect in two dimensional electron systems is spin-orbit
coupling (SOC) Heida  _et al._ (1998); Iñarrea (2017); Zhu  _et al._ (2016).
However, solid state materials used to investigate the effects of spin-orbit
coupling are often challenging due to the limited control of individual
experimental parameters. The body of work surrounding surrounding the topic of
spin-drag with added spin-orbit coupling is limited with theoretical
investigations looking at the impact of weak coupling on the drag in a 2D
electron system Tse and Das Sarma (2007), or the behaviour of impurities in a
spinor condensate system Liao  _et al._ (2016).

In depth understanding of this combination could lead to better understanding
of systems such as the topological insulators Kane and Mele (2005); Bernevig
_et al._ (2006) with their famed protected edge states are dependent on the
spin-momentum locking caused by the SOC within the material and have been
investigated as a potential platform for fault tolerant quantum computation
Nayak  _et al._ (2008); Raussendorf  _et al._ (2007).

Ultracold atoms provide an ideal environment for testing the effects of SOC on
the spin coulomb drag in quantum systems due to the ability to control many of
the crucial parameters accurately.

Previous experimental and theoretical work by Li et. al. Li  _et al._ (2019)
demonstrated the generation of spin currents using the same technique of a
quench of a spin-orbit coupled Bose-Einstein condensate (BEC) and investigated
the increased damping of the out of equilibrium system. GPE simulations showed
good qualitative agreement with the experimental results, and gathering
insight into the BEC shape oscillations and the miscible-imiscible phase
transition. As stated by Li et. al. the simulations underestimate the damping
of the BEC oscillations, potentially due to the lack of thermal atoms in the
simulations. In particular, it was shown in Duine and Stoof (2009) that the
spin-Coulomb drag between thermal atoms and the condensate dominates over the
mean field effects, rendering the mean field GPE only partially effective.

In this article we present our experiments on investigating the thermodynamic
behaviour of spin-orbit coupled systems within a conservative potential, and
attempt to explain the enhanced damping of the atomic oscillations in the
presence of SOC as Coulomb spin drag Duine and Stoof (2009) by comparing the
results to theoretical calculations.

We create synthetic SOC using the ground state manifold of a Rubidium-87
(87Rb) BEC, following the Raman laser scheme first demonstrated in the
experiments of Spielman et. al Lin  _et al._ (2011).

A bias magnetic field induces a Zeeman shift, breaking the degeneracy of the
F=1 ground state of the atom separating them in energy. A quadratic Zeeman
shift $\epsilon$ shifts the $m_{F}=+1$ state further than the $m_{F}=-1$,
allowing us to effectively decouple the latter from the system. The atoms in
the different Zeeman sublevels also differ in momentum by $\delta
k_{y}=2k_{R}$, where $k_{R}=2\pi/\lambda$ is the recoil momentum gained by the
atom due to absorption of a photon with wavelength $\lambda$. The Hamiltonian
of the coupled $F=1$ state as a function of the atomic quasimomentum
$\tilde{k}_{y}$ , is as follows,

$\hat{H}_{y}(\tilde{k}_{y})=\begin{pmatrix}\frac{\hbar^{2}(\tilde{k}_{y}+2k_{R})^{2}}{2m}-\hbar\delta&\frac{\hbar\Omega_{R}}{2}&0\\\
\frac{\hbar\Omega_{R}}{2}&\frac{\hbar^{2}\tilde{k}_{y}^{2}}{2m}-\hbar\epsilon&\frac{\hbar\Omega_{R}}{2}\\\
0&\frac{\hbar\Omega_{R}}{2}&\frac{\hbar^{2}(\tilde{k}_{y}-2k_{R})^{2}}{2m}+\hbar\delta\end{pmatrix}$
(1)

Diagonalizing the Hamiltonian gives the energies of the spin-orbit coupled
dressed states which for Raman coupling strengths $\hbar\Omega_{R}<4E_{R}$
features a double minimum. Here, $E_{R}=\hbar^{2}k_{R}^{2}/2m$, $\delta$ is
the two photon detuning between the bare states, and $\epsilon$ indicates the
quadratic shift. Correctly choosing the detuning for a given coupling results
in the ground state being an equal superposition of the two pseudospin states
$|\uparrow\rangle$, $|\downarrow\rangle$, corresponding to a spin-orbit
coupled state.

However, when $\hbar\Omega_{R}>4E_{R}$, there is only a single minimum at
quasimomentum $\tilde{k}_{y}=k_{R}$. In the experiments reported here, we
initially prepare a Bose-Einstein Condensate, trapped in a harmonic trap, in
the latter state with $\hbar\Omega_{R}=5.5E_{R}$. The system is then quenched
to a lower $\hbar\Omega_{R}<2E_{R}$, which takes the system out of
equilibrium, and allowed to thermalise. We find the time constant for this
thermalisation, and find that the rate scales with the coupling strength
$\Omega_{R}$.

##  II Experimental apparatus

Figure 1: Schematic of the experimental setup. (a) Geometry of the laser
orientation and polarization, along with the bias magnetic field. (b) The
level scheme of the $F=1$ ground state and the Raman transitions induced by
the coupling lasers. The Rabi frequencies of the individual transitions are
$\Omega_{1}$ and $\Omega_{2}$.

Our experiments begin with an all-optical BEC composed of approximately 104
87Rb atoms, optically pumped into the $|F=1;m_{F}=0\rangle$ before
evaporation, as described in our previous work Brown  _et al._ (2018). The BEC
is held in a harmonic trap, with aspect ratios
$\omega_{y}:\omega_{x}:\omega_{z}=1:1.2:2$, formed by a crossed-beam optical
dipole trap. We use two values of $\omega_{y}$ for our experiments. The lowest
trapping frequencies correspond to the the experiments performed with the trap
held at the final power 66 mW achieved after evaporation, with
$\omega_{y}=2\pi\times 85$ s-1. For the larger trapping frequencies, we
adiabatically increase the power of the dipole trapping laser to 90 mW,
corresponding to $\omega_{y}=2\pi\times 112$ s-1. Note that the larger trap
frequency corresponds to a larger trap depth, and thus increases the number of
atoms retained in the trap during the thermalization process and increases the
rate of collisions between the atoms.

During the evaporation to BEC a magnetic bias field By is ramped up in the
last 2 seconds to 8.35 G providing the a measured $\omega_{B}/2\pi=5.845$ MHz
Zeeman shift, and a measured quadratic Zeeman shift of $\epsilon/2\pi=5$ kHz.
A schematic of the coupling scheme and geometry is shown in Fig. 1. The two-
photon coupling strength $\Omega_{R}$ is experimentally determined by
observing the Rabi oscillations between the populations of the states
$|-1\rangle$, $|0\rangle$ and $|+1\rangle$ zero detuning, and fitting the time
evolution with the three state optical Bloch equations.

To induce spin-orbit coupling we use two orthogonally polarized laser beams
with wavelength $\lambda=790.2~{}$nm counter-propagating along $y$ that are
focused to an 150 $\mu$m diameter beam onto the center of the dipole trap.
This wavelength of $790.2$ nm was chosen to minimise the scalar AC Stark shift
in the atoms, which would have led to undesirable extraneous forces induced by
these beams. The two beams are derived from the same laser, but differ in
frequency by $\Delta\omega\approx\omega_{B}+\epsilon$ and couple two of the
internal $m_{F}$ levels of the BEC atoms. For sufficient quadratic Zeeman
shift, that is $h\epsilon>E_{R}$ the $m_{F}=+1$ internal state is tuned out of
resonance for the two photon Raman coupling. The coupled system becomes an
effective two level system of spin-momentum states which we label
$|m_{F}=-1,\tilde{k}_{y}+2k_{R}\rangle=|\uparrow^{\prime}\rangle$ and
$|m_{F}=0,\tilde{k}_{y}\rangle=|\downarrow^{\prime}\rangle$.

##  III Experimental Procedure

The condensate is prepared in the lowest energy dressed band of the Raman
coupled system by adiabatically increasing the Raman coupling to 5.5 $E_{R}$
over 50 ms, where there is a single minimum in the dispersion curve as
illustrated in Fig.2(a). The adiabatic increase of the coupling prevents
unwanted heating and oscillations of the condensate in the trap caused by
synthetic electric fields Lin  _et al._ (2011). We hold the Raman coupling on
for a further 30 ms at a constant value in order to ensure the system is in
the lowest energy dressed band. We confirm that the ramp is adiabatic by
measuring the total momentum of the atoms, obtained from the weighted sum of
the quasimomentum of all momentum components, during this 30 ms period and
confirming that it is zero at each point in time. If the total momentum is
non-zero during this phase, the ramp speed must be adjusted to ensure the
atoms remain in the lowest energy dressed state.

To take the system out of equilibrium, a synthetic electric force is imparted
on the dressed BEC by abruptly reducing the Raman coupling strength from the
initial $\Omega_{i}=5.5E_{R}/\hbar$ to a final value $\Omega_{f}$ in 1 ms. The
rapid decrease of Raman coupling constitutes a quench of the system. The
condensate separates into the two pseudospin states
$|\uparrow^{\prime}\rangle$ and $|\downarrow^{\prime}\rangle$ through the
synthetic electric force, each accelerating towards one of the new minima (see
Fig. 2(b)) of the dispersion relation, where they then oscillate in the
harmonic trap with maximal momentum $|k_{\uparrow,\downarrow}|=\pm 1\hbar
k_{R}$.

To compensate for both the impact of the $m_{F}=+1$ state and the changing AC
stark shift as the laser intensity changes, we adjust the laser frequency
difference to maintain equal populations of the two states, by an amount up to
$\hbar\delta_{\mathrm{AC}}=1E_{R}$. This shift in the two photon resonance
condition is extremely sensitive to small changes in experimental parameters,
such as the magnetic fields. Although care is taken to maintain equal
populations of the spin components, the final spin-orbit coupled state after
the quench will occasionally have non-equal populations in each component. To
group the data we calculate the population imbalance

Figure 2: Dispersion relations for the coupled BEC for two coupling strengths.
(a) $\hbar\Omega_{R}=5.5E_{R}$, features a single minima where both spins have
equal populations and the same quasimomentum. (b) $\hbar\Omega_{R}=1E_{R}$
features two minima of the dispersion. Each pseudospin occupies one of the
minima with the corresponding quasimomentum.

$F_{\uparrow\downarrow}=\frac{N_{\uparrow}-N_{\downarrow}}{N_{\uparrow}+N_{\downarrow}},$
(2)

where $N_{\uparrow}$ is the population of the $m_{F}=-1$ state and
$N_{\downarrow}$ is the population of the $m_{F}=0$ state. In this paper we
focus on the case with balanced populations, where
$|F_{\uparrow\downarrow}|<0.1$, by post-selecting the data.

We let the two pseudospin states oscillate in the dipole trap for time $t$ up
to 20 ms before switching the trap and Raman coupling off simultaneously,
projecting the atoms onto their bare spin-momentum states. The bare states
expand for 15 ms in a Stern-Gerlach gradient separating the spin components in
the $x$ dimension before being imaged with a resonant absorption method.

We measure the rate of thermalization of the system by evaluating the momentum
distribution of both spin states as a function of time. We numerically
determine the mean momentum of each of the spin ensembles as they oscillate in
the trap with a decaying amplitude. We fit the decay of the oscillation of the
momentum difference $k_{t}=|k_{\uparrow}-k_{\downarrow}|$ and measure the
final temperature of the thermalized ensembles.

##  IV Thermalization of a Spin-Orbit coupled BEC

Figure 3:  Plots of the momentum difference between the two components of the
system for two coupling strengths (a) $\hbar\Omega_{f}=1.5E_{R}$ and (b)
$\hbar\Omega_{f}=2E_{R}$ fit with a decaying cosine function. The insets show
the spin-momentum distribution after the system has returned to equilibrium,
and clearly demonstrate the return of the system to a spin orbit coupled state
with the final momentum distribution reflecting the non-zero quasimomentum
before release from the trap.

It is important to note that the momentum imparted on the pseudospins during
the quench depends of the final coupling strength, which arises from the
quasimomentum minima shifting as a function of the coupling strength. The
shift in this work on the order of $\pm 0.05\hbar k$ for each spin, accounting
for a 5% difference in the total momentum of the atoms of $\hbar\Omega_{f}=0$
and $\hbar\Omega_{f}=2E_{R}$

Once the system has thermalized and the oscillations have completely damped, a
small fraction of the condensate remains, with the atoms occupying the minima
of the new spin-orbit coupled band. For Raman coupling above
$\hbar\Omega_{R}=1E_{R}$ the time-of-flight images show clearly the system has
returned to equilibrium in a spin-orbit coupled state with the pseudospin
momentum clearly being non-zero. We confirm the non-zero momentum of the atoms
comes from the quasimomentum of the spin-orbit coupled state, rather than
residual oscillation energy, by noting the momentum remains unchanged over 5
ms of evolution.

Fig. 3 demonstrates two cases where the final pseudospins are separated from
zero momentum when reaching equilibrium. A clear example is shown in the inset
of Fig. 3(b) shows the pseudospins are positioned $k_{\uparrow,\downarrow}=\pm
0.25\hbar k_{R}$ , corresponding to quasimomentum before release
$\tilde{k}_{y}=\mp 0.75\hbar k_{R}$, the locations of the dispersion minima
obtained from exact diagonalization of the Hamiltonian. Even though the trap
frequencies are the same for the two situations in the figure, the dispersion
relation is different for different coupling, giving rise to the observed
difference in oscillation frequency. It is also clear that the higher coupling
strength (b) gives rise to a stronger damping of the oscillation.

As mentioned, some data was also obtained for imbalanced populations. In this
case, qualitatively we observe that the smaller population oscillation damps
more rapidly while the larger population continues to oscillate. Accurately
controlling the imbalanced populations proves to be difficult and therefore we
do not include these results in this paper.

##  V Spin Coulomb drag

For a situation with no spin-orbit coupling, the damping coefficient can be
determined theoretically for an ultracold Bose gas. The spin drag between two
components can be calculated from two expressions for the non-condensed and
the condensed atoms respectively Duine and Stoof (2009),

$\displaystyle\gamma_{22}=\frac{n\Lambda
a^{2}_{\uparrow\downarrow}}{\hbar\beta}\frac{1}{6\pi^{2}}\frac{1}{\left(n\Lambda^{3}\right)^{2}}\int_{0}^{\infty}\frac{\,dq\,d\omega\,q^{2}}{\sinh^{2}\left(\omega/2\right)}$
$\displaystyle\times\ln\left(\frac{\exp\left[q^{2}/16\pi+\beta
gn_{0}-\omega/2+\pi\omega^{2}/q^{2}\right]-\exp\left[-\omega\right]}{\exp\left[q^{2}/16\pi+\beta
gn_{0}-\omega/2+\pi\omega^{2}/q^{2}\right]-1}\right),$ (3)

and

$\displaystyle\gamma_{12}=\frac{n\Lambda
a^{2}_{\uparrow\downarrow}}{\hbar\beta}\frac{64n_{0}a}{3(2\pi)^{3}n\Lambda}\int_{0}^{\infty}\,dp_{1}\,dp_{3}\,p_{1}p_{3}^{3}$
$\displaystyle\times\left[1+\frac{1}{\exp\left[\left(p_{1}^{2}+p_{3}^{2}\right)/4\pi+2\beta
gn_{0}\right]-1}\right]$
$\displaystyle\times\frac{1}{\exp\left[p_{1}^{2}/4\pi+\beta
gn_{0}\right]-1}\times\frac{1}{\exp\left[p_{3}^{2}/4\pi+\beta
gn_{0}\right]-1}$
$\displaystyle\times\Theta\left(\frac{p_{1}p_{3}}{2\pi}-\beta gn_{0}\right).$
(4)

Here, $\Theta$ is the Heaviside function,
$\Lambda=\sqrt{2\pi\hbar^{2}/mk_{B}T}$ is the thermal de Broglie wavelength,
$\beta=1/k_{B}T$ the inverse thermal energy and $g=4\pi\hbar^{2}a/2m$ the
interparticle interaction strength. Calculations were performed using a script
provided by Jogundas Armaitis Duine and Stoof (2009) with our experimental
parameters, returning the total spin drag relaxation rate for a given density
of atoms. Due to the fact the atoms are oscillating in the trap and only
overlap and only interact periodically, we multiply by a scaling factor
calculated based on the interaction time of the two spin clouds overlapping in
the trap.

Figure 4: The damping rate of the system as a function of the Raman coupling
strength, for low (green circles) and high (red crosses) trapping frequencies.
The damping increases linearly over the range of coupling, with the gradient
being a combination of both collisional damping and spin drag. The red and
green dashed lines indicates the theoretical spin drag damping rates.

At the time of writing we are unable to obtain theoretical calculations for
the effects of the spin-orbit coupling on the spin drag, so we compare the
experiments for the uncoupled case with the theory. For the $y$ trapping
frequencies $\omega_{y}=2\pi\times 85$ s-1 and $\omega_{y}=2\pi\times 112$
s-1, we calculate a spin drag damping rate of
$\gamma_{s}=\gamma_{12}+\gamma_{22}=$13.6 s-1. Comparing this to our
experimentally observed damping rate of $\gamma=72(9)$ s-1, we observe that
part of $\gamma_{e}$ is caused by collisions of the condensate with thermal
atoms, and is dependent on the atomic density and is also present regardless
of spin drag.

Figure 5: Fractional temperature T/Tc of the system as a function of the final
coupling strength. The blue line shows the temperature before the quench
averaged over many shots, with the shaded region indicating the error. The
temperature of the equilibrium state is shown with red circles, demonstrating
a clear decrease in the heating for larger coupling strengths. Inset: The
integrated 1D profile showing the bimodal fit to one of the spin components

Taking into account the elastic collision rate, it is clear from Fig. 4 that
the increase in the spin-orbit coupling corresponds to a significantly
increased damping rate, with a linear dependence over the range measured. We
find that the damping rate also scales with the calculated spin-drag damping
rate and e summarize our results for the damping coefficient $\gamma$ by the
expression

$\gamma=\gamma_{e}+\xi\Omega_{R}\gamma_{s}$ (5)

where $\xi$ is a constant. The fit parameters in Fig. 4, as well as the spin
drag constant without spin-orbit coupling are summarised in table 1.

$\omega_{y}/2\pi$ (s-1) |  $\gamma_{e}$ (s-1) | $\xi\gamma_{s}$ |  $\gamma_{s}$ (s-1)
---|---|---|---
85 | 18(6) | 63(4) | 6.4
112 | 67(14) | 107(13) | 13.6
Table 1: Scaling of the fit parameters in figure 4

At the time of writing, we have not found a way to find $\xi$ from theoretical
considerations for our experimental configuration, but from table 1 we find
that $\xi=8.9(6)$ s, which is remarkably large as spin-orbit coupling affects
the condensate fraction much more strongly than the thermal fraction. We
envision that finding an accurate theoretical value may take a truncated
Wigner type simulation Blakie  _et al._ (2008); Brown  _et al._ (2018) to
include both spin components, the spin-orbit coupling along with their
interactions with both the opposing spin condensate atoms, but also the atoms
belonging to the thermal cloud.

Finally, we measure the final temperature of the system once it has reached
equilibrium. We integrate the time-of-flight region for each spin to obtain a
1D density profile and fit them with a sum of a Bose enhanced Gaussian and a
Thomas-Fermi profile. Integrating the fits we obtain the atom number for the
BEC and thermal component, which we use to obtain the fractional temperature
T/Tc. We plot the measured temperatures in Fig. 5 along with the initial
temperature and uncertainty. We fit a straight line to the temperature,
obtaining $T/T_{c}=1.08(0.01)-0.11(0.02)\Omega_{R}$.

It is interesting to note that the condensate is off center with respect to
the thermal cloud, consistent with the quasimomentum of the final spin orbit
coupled state mentioned in an earlier section. The results surprisingly show
that for increasing Raman coupling, the damping results in a reduced final
temperature, possibly indicating the spin-orbit coupling plays a significant
role in the relaxation process. The means by which the temperature decreases
is not so obvious, however the condensate fraction remaining at the end of the
experiment is increased for increasing coupling.

##  VI Conclusion

We have presented experiments performed to investigate the impact of spin-
orbit coupling on the thermalization processes present in an out of
equilibrium system of ultracold bosons. We measure the spin drag damping rate
of the atoms and compare the uncoupled case to theoretical calculations. We
show that introducing the spin-orbit coupling into the system strongly
increases the rate at which the system returns to equilibrium, while also
reducing the temperature increase caused by the excitation. Finally, we have
shown that the equilibrium state of the system after rethermalization is a
spin-orbit coupled BEC, with the quasimomentum measurements after reaching
equilibrium corresponding to the dispersion relation calculated through exact
diagonalization of the Hamiltonian. We anticipate that this work will lead to
new understanding of thermalization in the presence of spin-orbit coupling.

## References

  * Wolf  _et al._ (2001) S. A. Wolf, D. D. Awschalom, R. A. Buhrman, J. M. Daughton, S. von Molnár, M. L. Roukes, A. Y. Chtchelkanova, and D. M. Treger, Science 294, 1488 (2001).
  * D’Amico and Vignale (2000) I. D’Amico and G. Vignale, Phys. Rev. B 62, 4853 (2000).
  * Weber  _et al._ (2005) C. P. Weber, N. Gedik, J. E. Moore, J. Orenstein, J. Stephens, and D. D. Awschalom, Nature 437, 1330 (2005).
  * Duine and Stoof (2009) R. A. Duine and H. T. C. Stoof, Phys. Rev. Lett. 103, 170401 (2009).
  * Sommer  _et al._ (2011) A. Sommer, M. Ku, G. Roati, and M. W. Zwierlein, Nature 472, 201 (2011).
  * Goulko  _et al._ (2013) O. Goulko, F. Chevy, and C. Lobo, Phys. Rev. Lett. 111, 190402 (2013).
  * Fava  _et al._ (2018) E. Fava, T. Bienaimé, C. Mordini, G. Colzi, C. Qu, S. Stringari, G. Lamporesi, and G. Ferrari, Phys. Rev. Lett. 120, 170401 (2018).
  * Li  _et al._ (2019) C.-H. Li, C. Qu, R. J. Niffenegger, S.-J. Wang, M. He, D. B. Blasing, A. J. Olson, C. H. Greene, Y. Lyanda-Geller, Q. Zhou, C. Zhang, and Y. P. Chen, Nature Communications 10, 375 (2019).
  * Heida  _et al._ (1998) J. P. Heida, B. J. van Wees, J. J. Kuipers, T. M. Klapwijk, and G. Borghs, Phys. Rev. B 57, 11911 (1998).
  * Iñarrea (2017) J. Iñarrea, Scientific Reports 7, 13573 (2017).
  * Zhu  _et al._ (2016) C. Zhu, L. Dong, and H. Pu,  49, 145301 (2016).
  * Tse and Das Sarma (2007) W.-K. Tse and S. Das Sarma, Phys. Rev. B 75, 045333 (2007).
  * Liao  _et al._ (2016) R. Liao, O. Fialko, J. Brand, and U. Zülicke, Phys. Rev. A 93, 023625 (2016).
  * Kane and Mele (2005) C. L. Kane and E. J. Mele, Phys. Rev. Lett. 95, 146802 (2005).
  * Bernevig  _et al._ (2006) B. A. Bernevig, T. L. Hughes, and S.-C. Zhang, Science 314, 1757 (2006).
  * Nayak  _et al._ (2008) C. Nayak, S. H. Simon, A. Stern, M. Freedman, and S. Das Sarma, Rev. Mod. Phys. 80, 1083 (2008).
  * Raussendorf  _et al._ (2007) R. Raussendorf, J. Harrington, and K. Goyal, New Journal of Physics 9, 199 (2007).
  * Lin  _et al._ (2011) Y.-J. Lin, K. Jiménez-García, and I. B. Spielman, Nature 471, 83 (2011).
  * Brown  _et al._ (2018) D. J. Brown, A. V. H. McPhail, D. H. White, D. Baillie, S. K. Ruddell, and M. D. Hoogerland, Phys. Rev. A 98, 013606 (2018).
  * Blakie  _et al._ (2008) P. Blakie, A. Bradley, M. Davis, R. Ballagh, and C. Gardiner, Advances in Physics 57, 363 (2008).