Self-Supervised Learning Of Brain Dynamics From Broad Neuroimaging Data
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Self-supervised learning techniques are celebrating immense success in naturallanguage processing (NLP) by enabling models to learn from broad language data atunprecedented scales.
Here, we aim to leverage the success of these techniques formental state decoding, where researchers aim to identify speciﬁc mental states (suchas an individual’s experience of anger or happiness) from brain activity.
To this end,we devise a set of novel self-supervised learning frameworks for neuroimaging databased on prominent learning frameworks in NLP.
At their core, these frameworkslearn the dynamics of brain activity by modeling sequences of activity akin tohow NLP models sequences of text.
We evaluate the performance of the proposedframeworks by pre-training models on a broad neuroimaging dataset spanningfunctional Magnetic Resonance Imaging (fMRI) data from 11 , 980 experimentalruns of 1726 individuals across 34 datasets and subsequently adapting the pretrained models to two benchmark mental state decoding datasets.
We show that thepre-trained models transfer well, outperforming baseline models when adapted tothe data of only a few individuals, while models pre-trained in a learning frameworkbased on causal language modeling clearly outperform the others.
Introduction
In mental state decoding (1), researchers aim to identify certain mental states (e g., deciding to acceptor reject a gamble or the experience of happiness or fear) from brain activity.
To this end, researchers train predictive models to correctly identify (i e , decode) these mental states from measured brainactivity.
At ﬁrst sight, this approach seems straightforward, yet, researchers interested in trainingmental state decoding models are faced with the challenge that individual neuroimaging datasetsare generally of high-dimension and small sample-size, such that individual samples often comprisemany hundred thousand dimensions while each dataset includes only a few hundred samples for eachof tens to hundreds of individuals.
In this setting, predictive models are prone to overﬁtting, therebynot generalizing well to new individuals or datasets, which severely limits any generalizable insights that can be gained from training these models about the studied mental states and brain activity.
Neuroimaging research has recently entered a big data era, as individual researchers more frequently share their collected datasets publicly (2; 3).
In addition, several efforts have been made to standardize the data structure (4) and preprocessing (5) of neuroimaging data.
Together, these developments open up new possibilities for the application pre-training in neuroimaging at scale, by training models onlarge and diverse sets of public neuroimaging data so that the resulting models can be adapted well tothe datasets of individual researchers, without the need for much new data.
A wealth of empirical evidence has demonstrated that models pre-trained on large neuroimagingdatasets generally achieve better mental state decoding performance on new data, while requiring less training time and data, when compared to models trained from scratch (e.g., 6; 7; 8)
Yet, much of this work has either pre-trained models on large but homogenous datasets, such as data from manyindividuals who all perform the same few tasks at the same few acquisition sites (e.g., by using datafrom the Human Connectome Project; 9; 7) or trained models on highly-preprocessed input data(e.g.
, statistical maps summarizing the measured sequences of brain activity; 10), both limiting thedownstream generalizability of the pre-trained models.
Here, we aim to overcome these limitations by leveraging recent advances in self-supervised learningto pre-train DL models on broad, minimally-preprocessed, and public neuroimaging data.
Particularly,we take inspiration from natural language processing, where the application of self-supervised learninghas led to unprecedented breakthroughs in the adaptive capabilities of models trained on public textdata (e.g., 11; 12; 13; 14).
Speciﬁcally, we test whether modeling sequences of brain activity akinto sequences of text enables the learning of generalizable representations of the dynamics of brainactivity.
To this end, we represent sequences of brain activity akin to the embedded representationof words in sequences of text, by describing each time point as a vector in which each dimensionreﬂects the activity of a speciﬁc functional brain network at this point in time.
Contributions.
We devise a set of novel self-supervised learning frameworks for neuroimaging databased on prominent learning frameworks in NLP, namely, sequence-to-sequence autoencoding (15),causal language modeling (11), and masked language modeling (12).
We pre-train DL models withthese frameworks on a large-scale, minimally-preprocessed upstream dataset comprising functional
Magnetic Resonance Imaging (fMRI) data from 11 , 980 experimental runs of 1726 individuals across34 datasets.
To our knowledge this represents one of the broadest datasets used for pre-training inneuroimaging to date, spanning a wealth of experimental conditions and domains as well as a diverseset of acquisition sites.
We then evaluate the downstream adaptation performance of the pre-trainedmodels in two benchmark mental state decoding datasets and show that models pre-trained in a causalsequence modeling framework clearly outperform the others, while all pre-trained models outperforma baseline model trained from scratch.
To enable others to build on this work, we make our trainingdata and pre-trained models publicly available (see Appendix A.4).
Methods overview
2.1 Input data: from brain volumes to brain networks
FMRI data are conventionally represented in four dimensions, such that the measured blood-oxygenlevel-dependent (BOLD) signal is indicated in temporal sequences S
= { V 1 , ..., V t } of 3-dimensionalvolumes V 2 R x ⇥
y ⇥ z that show the BOLD signal intensity for each spatial location of the brain (asindicated by the three spatial dimensions x , y , and z ).
However, due to the strong spatial correlationsof brain activity, the measured BOLD signal can also be summarized differently by representing it asa set ⇥ =
{ ✓ 1 , ..., ✓ n } of n brain networks ✓ , each describing the BOLD signal for some subset ofvoxels v x,y,z ⇢ V that exhibit correlated activity patterns.
Here, we utilize such a functional parcellation of brain activity to represent the measured BOLDsignal, namely, Dictionaries of Functional Modes (DiFuMo), which were recently proposed by (16)and learned across millions of fMRI volumes.
DiFuMo deﬁnes ⇥ through a sparse dictionary matrix D 2 R p ⇥ n containing n networks, each deﬁned by p values (one for each voxel v 2 V whenﬂattened over the three spatial dimensions).
The BOLD signal of volume V t is then described asa linear combination of these networks, by the use of weights ↵ t 2 R n , such that V t
= D ↵ t and↵ t = D † V t , where D † indicates the pseudo-inverse of D : D † =
( D T D )  1 D T 2 R n ⇥ p .
We utilize aparcelation with n = 1024 networks, such that the resulting BOLD parcelation X 2 R t ⇥ n describesthe measured BOLD signal of each of these networks at each time point t of the input sequence S . Note that this representation of the input
data
is analogous to the embedded representation of sentences in
natural language processing
, where
each
word
v
is represented as some n -dimensional
vector v 2 R n .
2.2 Self-supervised learning frameworks: modeling sequences of brain activity akin tosequences of text
In the following, we describe our proposed set of self-supervised learning frameworks, which arebased on prominent learning frameworks in NLP (for an overview, see Fig. 1).
Before going into thedetails of each framework, we will establish several commonalities among them.
2.2.1 Commonalities among frameworks
[Figure] First, each framework takes as input a parcelated BOLD sequence X 2 R t ⇥ n (with t time pointsfor n networks) as well as the corresponding sampling frequency (i.e., repetition time or TR) of theunderlying fMRI sequence (typically in the range of 0 .
7 s to 4 s).
Second, all frameworks linearly project the parcelated BOLD sequence X into an embeddingrepresentation E X 2 R t ⇥ e : E Xt,e = b e +
P n w n,e X t,n (where w 2 R n ⇥ e and b 2 R e indicateweights and biases), thereby reducing the dimensionality of the input: R t ⇥ n !
R t ⇥ e , if e < n .
Third, to account for any differences in the sampling frequency of the fMRI input sequences, theframeworks add one of r learnable TR embeddings E TR 2 R r ⇥ e to each time point t in the BOLDembedding E X .
Each TR embedding thus corresponds to one possible time point k of the input.
We set r = 1500 with a time increment of 0 .
2 s between any subsequent TR embeddings, such thatk 2 0 , 0 .
2 , ..., 300 s.
For example, for a sequence of 21 s with a TR of 0 .
7 s, we ﬁrst round each timepoint t 2 0 , 0 .
7 , ..., 21 to the closest k and then add the corresponding E TRk to E Xk .
Note that wealways label time points relative to the beginning of the input sequence, such that time point 0 isalways assigned to the ﬁrst sequence element.
Fourth, all frameworks utilize some standard DL model f ( · ) that takes E in as input and predicts acorresponding output sequence f
( E in ) 2 R t ⇥ e .
Fifth, during upstream learning, the output sequence f ( E in ) is linearly projected back to the dimensionality of the parcelated BOLD sequence X , such that ˆ X t,n = b n +
P e w e,n f ( E in ) t,e , where wand b again indicate weights and biases.
This step is crucial, as all proposed learning frameworksinvolve reconstructing X (or parts of X ) after some of its elements have been masked or it hasbeen encoded in a lower-dimensional representation.
To measure reconstruction performance, weutilize the mean absolute difference between the parcelated BOLD sequence X and the reconstructed
BOLD sequence ˆ X : L rec = 1 n ( J )
P j 2 J | X j  ˆ X j | , where J indicates the set of elements that is tobe reconstructed and n ( J ) the number of elements in J .
Sixth, for downstream adaptation to a new mental state decoding task, the frameworks utilize asimple decoding head p ( · ) , composed of a dense hidden layer with e model units (one for eachembedding dimension, with tanh activation) as well as a softmax output layer with one model unit i for each considered mental state in the data.
p ( x ) i therefore indicates the probability that input xbelongs to mental state i .
Consequently, all frameworks treat the decoding task as a multinomialclassiﬁcation problem and use a standard cross entropy loss objective: L cls =
P i
y
i log p ( x ) i ,where y i indicates a binary variable that equals 1 if i is the correct mental state and 0 otherwise.
2.2.2 Autoencoding
The autoencoding framework is inspired by recurrent sequence-to-sequence neural machine translation models (e.g., 15).
These models take as input some text sequence, encode this sequence ina lower-dimensional representation, and predict an output sequence from the lower-dimensionalencoding (e.g., a translation of the original input text to another language).
Accordingly, our autoencoding framework uses a standard recurrent autoencoder architecture f a ( · ) ,with an encoder and decoder part, each comprised of a stack of long short-term memory (LSTM;17) units (see Fig. 1).
The encoder e ( · ) takes E in as input and encodes it in a lower-dimensionalrepresentation h 2 R e .
The decoder d ( · ) then takes h as input and predicts a correspondingoutput sequence d ( h ) 2 R t ⇥ e , such that f a ( E in ) =
d ( h ) .
Our encoder is composed of a stackof bidirectional LSTM units, whereas the decoder is build from a stack of unidirectional LSTMunits.
Note that we apply residual connections between the LSTM units in a stack: Let LSTM i and LSTM ( i +1) be the i -th and ( i + 1) -th
LSTM layer in a stack (with weights w i and w i +1 ).
At thet -th step of an input sequence, the input and outputs of layers i and ( i + 1) are deﬁned as: h it , c it =
LSTM i ( h it  1 , c it  1 , x i  1 t  1 ; w i ) ;
x it = h it + x i  1 t  1 ; h i +1 t
+1 ,
c i +1 t
+1 =
LSTM i +
( h i +1 t , c i +1 t , x it ; w i +1 ) ,where x it 2 R e represents the input to LSTM i at step t , while h
it 2 R e and c it 2 R e represent the LSTM ’s hidden and cell states respectively.
For the encoder, we average the hidden states h
z 2 i t andh y 2 i t of the two LSTM units z and y contained in any bidirectional LSTM layer i at each time point t :h it = 0 .
5( h z 2 i t + h
y 2 i t ) .
The encoded sequence embedding h that is forwarded to the decoder is thendeﬁned as the mean of the ﬁnal hidden states of the two LSTM units contained in the encoder’s ﬁnallayer.
Upstream learning.
During pre-training, we jointly train the encoder and decoder to reconstructthe parcelated BOLD sequence X by minimizing its mean absolute difference to the reconstructed BOLD sequence ˆ X (see section 2.2.1): min 1 t
P ti =1 1 e
P ej =1 | ˆ X i,j  X i,j | .
Importantly, we applyteacher-forcing (18) during upstream learning by using E in in 50% of the cases as each next inputfor the decoder, instead of using the decoder’s own preceding output.
Adaptation.
To adapt the autoencoding framework to a new decoding task, we forward h as inputto a corresponding decoding head p ( · ) (see section 2.2.1), thereby omitting the decoder d ( · ) .
2.2.3 Causal sequence modeling (CSM)
The causal sequence modeling (CSM) framework is inspired by recent advances in causal languagemodeling (11), where models are trained to predict the next word in a sequence of text.
Speciﬁcally,these models receive some text sequence as input, e.g., represented as a matrix X 2 R n ⇥ e with nwords each encoded by an e -dimensional vector, where the last sequence element (representing thelast word) has been replaced by some some learned mask embedding E msk 2 R e
(e.g., "When itrains, it [msk]").
The task of a model is then to predict the word that was replaced by E msk .
In causallanguage modeling, models are therefore solely concerned with the part of the input sequence thatprecedes the masked word.
To adapt causal language modeling to neuroimaging data, we mask the last time point T of theparcelated BOLD sequence X by replacing it with a learnable mask embedding E msk .
As forthe autoencoding framework, X is then transformed to E in by means of linear projection and theaddition of TR embeddings E TR .
In line with recent advances in causal language modeling, we utilizea standard transformer decoder model f d ( · ) (based on the GPT architecture; 11) in this learningframework.
As transformer models are not aware of the ordering of their input, we further addlearnable position embeddings E pos 2 R p ⇥ e to E in , each representing one of p possible positions inthe input sequence (we set p = 512 for all transformer models).
Given E in , the transformer decodermodel predicts a corresponding output sequence f d ( E in ) 2 R t ⇥ e in the form of the hidden states atthe output of its last layer.
Upstream learning.
In line with causal language modeling, the upstream objective of the CSMframework is to reconstruct the masked last time point T of the parcelated BOLD sequence X byminimizing the mean absolute difference to its reconstruction ˆ X T : min 1 e
P ei =1 | ˆ X T,i  X T,i |
Adaptation.
During downstream adaptation, we attach a learnable classiﬁcation embedding E cls
R e to the end of BOLD embedding E X and forward the resulting prediction f
( E in ) T +1 to acorresponding decoding head p ( · ) (see section 2.2.1).
Note that we use the same learnable TRembedding for all embeddings that are attached to an input sequence during training.
2.2.4 Sequence-BERT
The Sequence-BERT framework is based on recent advances in bidirectional masked languagemodeling, particularly BERT (12).
In NLP, BERT is trained to jointly solve a masked-languagemodeling and next-sentence-prediction task.
To this end, BERT receives two masked sentences S a(e.g., "The sky is [msk].") and
S b (e.g., "The [msk] is shining.") as input, in which some fraction ofwords have been randomly masked.
BERT is then asked to I) predict the masked words (maskedlanguage-modeling) and II) determine whether S b follows S a (next-sentence-prediction).
In contrastto causal language modeling, masked language modeling is considered as a bidirectional learningtask, as the model is not just concerned with the parts of the input that precede a masked word butalso the parts that follow it.
To adapt the BERT framework to neuroimaging data, we ﬁrst sample two input sequences X a {+
+} R t a ⇥ n and X b 2 R t b ⇥ n from the data (with lengths t a and t b ).
In 50% of the cases, the two sequencesrepresent sequential parts of the same underlying fMRI sequence, while in the other 50% of cases, thetwo sequences are randomly sampled from two distinct fMRI sequences.
Due to the strong temporalautocorrelation of BOLD data, we randomly leave a gap of 1 - 5 TRs between the two sequences whensampling them from the same underlying fMRI sequence, to encourage the models to consider theentire input sequences and not just their connecting time points.
As for BERT, we next randomlymask time points of each sampled sequence (with a probability of 20% ), by replacing them with alearnable mask embedding E msk .
The masked versions of X a and X b are then linearly projectedto obtain respective BOLD embeddings E X a and E X b to which we then add TR and positionembeddings E TR and E pos as well as a sequence ID embedding E id 2 R 2 ⇥ e (to indicate each of thetwo input sequences).
In line with BERT, we lastly also attach a separation embedding E sep 2 R eto the end of each sequence (to indicate their endpoints) and a classiﬁcation embedding E cls to thebeginning of S a (for the next-sentence-prediction task; see upstream learning below).
The resultingembedding representation E in =
[ E cls , { E X a t
[Equation]
E post + E id 2 } t 2 t b , E sep ] is then forwarded to a standard transformer encoder model f e ( · ) (based onthe BERT architecture; 12), which outputs a corresponding output sequence f e ( E in ) 2 R t ⇥ e in theform of the hidden states at the output of its ﬁnal layer.
Upstream learning.
Similar to BERT, we deﬁne the upstream objective of the Sequence-BERTframework as the sum of two learning objectives:
In line with masked-language-modeling, theﬁrst objective is deﬁned as the reconstruction loss L rec for all masked time points J : L rec =1 n ( J )
P j 2 J 1 e
P ei =1 | X j,i  ˆ X j,i | .
In line with next-sentence-prediction, the other objective is deﬁnedas the binary cross entropy loss for a decoding head that receives f ( E in ) 1 (corresponding to thetransformer output for classiﬁcation embedding E cls ) as input and predicts probability p ( f ( E in ) 1 )that X b follows X a : L cls =
y log( p ( f ( E in ) 1 ))
(1  y ) log(1  p ( f ( E in ) 1 )) , where y indicatesa binary variable that equals 1 if the two sequences come from the same fMRI run and 0 otherwise.
Adaptation.
During downstream adaptation, we forward the transformer’s output for classiﬁcationembedding f e ( E in ) 1 as input to a corresponding decoding head p ( · ) (see section 2.2.1).
2.2.5 Network-BERT
All so far presented frameworks model their input on the level of the individual time points in asequence, by trying to reconstruct the parcelated whole-brain BOLD signal for speciﬁc time points.
Yet, the dynamics of brain activity data can also be viewed differently, by focusing on the interaction ofnetwork activities over time instead of focusing on the distribution of network activities at individual time points.
We therefore created a variant of the Sequence-BERT framework that randomly (at arate of 10%) masks the activity of individual networks n in the parcelated BOLD sequence X byreplacing them with a learnable network mask embedding E msk 2 R t , comprising one value foreach possible time point in the input (see Fig. 1).
With the exception of the masking procedure, the
Network-BERT and Sequence-BERT frameworks are identical.
Upstream learning.
During upstream learning, Network-BERT uses the same cross-entropy loss
L cls as Sequence-BERT and adds it to a variant of the reconstruction loss that measures reconstructionperformance for masked network time courses
J : L rec = 1 n ( J )
P j 2 J 1 t P ti =1 | X i,j  ˆ X i,j | .
Adaptation.
During downstream adaptation, Network-BERT forwards the transformer’s output forthe classiﬁcation embedding f e ( E in ) 1 to a respective decoding head p ( · ) (see section 2.2.1).
2.3 Training details
If not reported otherwise, we train all models with stochastic gradient descent and the ADAMoptimizer (with  1 = 0 .
9 ,  2 = 0 . 999 , and ✏ = 1 e  8 ; (19)).
We further use a linear learning ratedecay schedule (with a warmup phase of 1% of the total number of training steps), gradient normclipping at 1 . 0 , and L 2 -regularisation (weighted by 0 . 1 ).
During upstream learning, we set the maximum learning rates to 2 e  4 , 5 e  4 , 1 e  4 , and 1 e  4 for theautoencoding, CSM, Sequence-BERT, and Network-BERT frameworks respectively (based on thelearning rates used in corresponding NLP studies; 11; 12; 15) and randomly sample input sequencesof 10 to 55 TRs from the fMRI ﬁles contained in our upstream dataset.
For the two BERT frameworks,each sequence S i is split into two sequences S ai and S bi of equal length (with a randomly sampledgap of 1 to 5 TRs between the two sequences), while S bi is then randomly exchanged (with 50%probability) with another S bj (where i 6 = j ).
For each upstream training run, we then randomly dividethe upstream data into distinct training and validation datasets by randomly designating %5 of thefMRI run ﬁles of each upstream dataset as validation data (at a minimum of 2 ﬁles per dataset) andusing the rest of the ﬁles for training.
At each evaluation step, we randomly sample 640 , 000 fromthe validation dataset to evaluate model performances.
During downstream adaptation, we begin training with the respective pre-trained model parametersand then allow all parameters to be freely changed during the adaptation.
Other optimization andarchitecture hyper-parameters are evaluated and reported in the experiments described below.
Compute resources.
Upstream learning was performed on Google Compute Engine n1-highmem64 nodes with four Nvidia Tesla P100 GPUs, 64 CPU threads, and 416 GB RAM memory, whiledownstream adaptation was performed on compute nodes of the Texas Advanced Computing Centerwith one Nvidia 1080-TI GPU, 32 CPU threads, and 128 GB RAM memory.
Training times varieddepending on model size but were generally in the range of 1 to 3 days for upstream learning and 0 .
5to 3 hours for downstream adaptations.
Experiments
3.1 Datasets
All unprocessed fMRI data used in this study are publicly available through OpenNeuro.org (3) andthe Human Connectome Project (HCP; 20).
For an overview of our preprocessing, see Appendix A.3. 3.1.1 Upstream: 11 , 980 fMRI runs from 1726 individuals across 34 datasets
Our upstream dataset comprises 11 , 980 fMRI run ﬁles from 1726 individuals across 34 datasets(for an overview, see Fig.
2 and Appendix Table A.1).
To our knowledge, this represents one of thebroadest fMRI datasets used for pre-training in neuroimaging to date, spanning many acquisition sitesand a diverse set of experimental conditions and domains.
A few prominent examples of includeddatasets are The Midnight Scan Club (21), Individual Brain Charting (22), Amsterdam Open MRI
3.1.2 Downstream: 2 benchmark datasets with 20 and 26 mental states
To evaluate downstream adaptation performance, we use two benchmark mental state decoding datasets, which span a wide range ofmental states, namely, task-fMRI data from the HCP (20) and the Multi-Domain Task Battery (MDTB; 26).
Our HCP dataset includesdata from 100 individuals performing two fMRI runs for each ofseven tasks, with a total of 19 distinct mental states across the tasks(see Appendix A.2).
We also include two resting state fMRI runs perindividual, in which individuals simply rest in the fMRI without anyparticular task, thereby increasing the total number of mental statesto 20 for the HCP data.
The MDTB data include 24 individualsperforming 26 tasks in two scanning sessions with eight fMRI runsper session.
Due to the strong similarity of experimental conditionswithin individual tasks of the MDTB, we deﬁned each of its tasksas one target mental state (see Appendix A.2).
During downstreamadaptation, we sample sequences from both datasets according totheir experiment event ﬁles, which deﬁne an on- and offset for eachtrial in an fMRI run (with durations generally ranging between 2 and30 s, at a TR of 0 .
72 (HCP) and 1 s (MDTB)).
For the HCP’s restingstate data, we randomly sampled sequences between 10 and 30 s.
3.2 Hyper-parameter evaluation: larger embeddings but not deeper models
In our ﬁrst experiment, we compare the performance of different hyper-parameter settings for theproposed frameworks.
Speciﬁcally, we evaluate three different embedding dimensions e ( 192 , 384 ,and 768 ; see section 2.2.1) and three different numbers of hidden layers for the used DL models ( 2 ,4 , and 6 hidden layers for the encoder and decoder parts of the autoencoding framework and 4 , 8 ,and 12 hidden layers for the transformer models).
In line with other work (27), we scale the numberof attention heads n ↵ per layer of the transformer models with the embedding dimension, such thatn ↵ = e 64 .
We then train one model for each point of the resulting 9 -point grid of each framework for200 , 000 training steps with a mini-batch size of 256 random samples per step.
Overall, model performances improve (in terms of validation loss) with larger embedding dimensionsbut not with the number of hidden layers (see Fig. 3).
To this end, we decided to continue all furtheranalyses with models comprising 768 embedding dimensions and 2 hidden layers for the encoderand decoder parts of the autoencoding framework and 4 hidden layers for the transformer models.
3.3 Upstream performance: models learn well in all frameworks
In experiment two, we then fully pre-trained one model in each framework according to our previouslyselected set of hyper-parameters (see section 3.2), by training models for 300 , 000 training steps ata mini-batch size of 768 .
All models learned well over the course of their training, with graduallydecreasing training and validation losses (see Fig. 4).
Note that the transformer model trained in CSM slightly overﬁtted the upstream data while the gap in training and validation performance of theautoencoding framework results from our application of teacher-forcing during training (see section2.2.2).
We also evaluated the reconstruction performance of the ﬁnal models for different parts of thebrain (see Appendix A.6) and found that the frameworks exhibited similar overall distributions ofreconstruction error, with relatively higher errors in the posterior parietal, occipital, and cingulatecortices as well parts of the limbic system.
3.4 Adapting pre-trained language models does not yield meaningful upstream performancegains
In our third experiment, we tested whether adapting pre-trained language models to our upstreamdata would yield meaningful performance gains, when compared to the models trained from scratch.
To this end, we adapted pre-trained variants of GPT-2 (28) and BERT (12) (with 12 hidden layersand an embedding dimension of 768 ; as provided by Hugging Face’s Transformer library (29))
toour upstream dataset (see Appendix A.7). However, the upstream validation performance of these models did not meaningfully improve upon the performance of our models trained from scratch, atﬁnal upstream validation losses of 0 .
73 (CSM with GPT-2) and 0 .
80 (Sequence-BERT with BERT).
3.5 Downstream adaptation: CSM outperforms other frameworks
Lastly, in our fourth experiment, we evaluated the downstream adaptation performance of our pretrained models in two benchmark mental state decoding datasets (see section 3.1.2).
To gauge theeffectiveness of pre-training, we adapted the pre-trained models to varying sizes of the downstreamdatasets, comprising the data from 1 to 48 (HCP) and 11 (MDTB) individuals.
Importantly, for eachtraining run, we randomly split the data into training, validation, and test datasets, by randomlyassigning individuals to one of these three datasets.
Speciﬁcally, for each run, we ﬁrst randomlyselected 10 (HCP) and 3 (MDTB) individuals whose data we use for validation and 20 (HCP) and 9(MDTB) other individuals whose data we use for testing, before randomly sampling the given numberof individuals whose data we use for training from the remaining pool of individuals.
We then adaptthe pre-trained models to the training data at a mini-batch size of 256 with varying learning rates andnumbers of training steps (see Appendix A.9 for an overview of our downstream hyper-parameterevaluation).
Here, we only report ﬁnal test decoding accuracies for those models that achieved thehighest ﬁnal validation decoding accuracy (see Appendix Fig.
A4 and A5 for individual validation accuracies).
In line with the state-of-the-art in mental state decoding (see 30), we also evaluated theperformance of a linear baseline model on the both datasets, which we trained from scratch each time(see Appendix A.8).
The pre-trained models clearly outperformed the linear baseline model in their test decoding accuracyacross the different sizes of the training dataset (at chance-levels of 5 . 0% (HCP) and 3 . 9% (MDTB)),while the models’ performances also scaled well with the number of individuals in the training data(Fig. 5).
For the HCP data, for which other reported decoding performances exist, our pre-trainedmodels perform on par with models trained on much larger fractions of the data (often comprisingthe data of many hundred individuals; 9; 7; 6; 31), which generally report test decoding accuracies inthe range of 80 to 93% .
Overall, the transformer decoder model trained with the CSM frameworkclearly outperformed the other models by achieving the highest test decoding accuracy throughout allevaluated sizes of the training dataset.
Note that model performances were very stable across the different random splits of the data (see Fig.5 and Appendix Fig. A4 and A4), strongly suggesting stability in the models’ performance towardsdifferent possible splits of the data.
In addition, we tested whether model performances are alsostable over other non-deterministic aspects of the training by replicating our downstream adaptationanalysis for the MDTB dataset with a different set of random seeds. Again, model performanceswere very stable across the different random seeds, yielding almost identical results in the replicationwhen compared to our original analysis (see Appendix A.10).
Conclusion
We propose a set of novel self-supervised learning frameworks for neuroimaging data based onprominent counterparts in natural language processing.
By the use of these frameworks, we areable to pre-train DL models across broad, public neuroimaging data, comprising many individuals,experimental domains, and acquisition sites.
We show that the pre-trained models generalize well totwo benchmark neuroimaging datasets in a mental state decoding analysis, outperforming baselinemodels even when trained on small subsets of the data.
To enable others to build on this work, wemake our training data and pre-trained models publicly available (see Appendix A.4).
Limitations and potential negative social impact.
While the proposed frameworks perform wellin training models that generalize to new neuroimaging datasets, this work currently neglects anotherkey aspect of mental state decoding, namely, the ability to draw inferences about the associationbetween decoded mental states and brain activity from the trained models.
First empirical evidenceindicates that attribution methods from explainable artiﬁcial intelligence (XAI; 32) research arewell-suited to provide insights in the mental state decoding decisions of DL models (for a detaileddiscussion, see 33).
Yet, further research is needed to understand their limits in mental state decoding.
We are currently not aware of any direct negative social impacts that could follow from pre-training DL models on public neuroimaging data, which are de-identiﬁed and collected and shared withhuman consent in a way approved by respective institutional review boards.