Self-Supervised Learning of Brain Dynamics from Broad Neuroimaging Data
Abstract
Self-supervised learning techniques are celebrating immense success in natural language processing (NLP) by enabling models to learn from broad language data at unprecedented scales.
Here, we aim to leverage the success of these techniques for mental state decoding, where researchers aim to identify specific mental states (e.g., the experience of anger or joy) from brain activity.
To this end, we devise a set of novel self-supervised learning frameworks for neuroimaging data inspired by prominent learning frameworks in NLP.
At their core, these frameworks learn the dynamics of brain activity by modeling sequences of activity akin to how NLP models sequences of text.
We evaluate the frameworks by pre-training models on a broad neuroimaging dataset spanning functional Magnetic Resonance Imaging data from 11 , 980 experimental runs of 1 , 726 individuals across 34 datasets and subsequently adapting the pre-trained models to benchmark mental state decoding datasets.
The pre-trained models transfer well, generally outperforming baseline models trained from scratch, while models trained in a learning framework based on causal language modeling clearly outperform the others.
Introduction
In mental state decoding researchers aim to identify certain mental states (e g., deciding to accept or reject a gamble or the experience of happiness or fear) from brain activity [1].
To this end, researchers train predictive models to correctly identify (i e., decode) these mental states from measured brain activity.
At first sight, this approach seems straightforward, yet, researchers interested in training mental state decoding models are faced with the challenge that individual neuroimaging datasets, particularly functional Magnetic Resonance Imaging (fMRI) data, are often of high dimension and low sample size, such that individual samples can comprise many hundred thousand dimensions while a dataset includes only a few hundred samples for each of tens to hundreds of individuals.
In this setting, predictive models are prone to overfitting, thereby not generalizing well to new datasets, which severely limits any generalizable insights that can be gained from training these models about the studied mental states and brain activity.
In spite of the low sample size of individual datasets, 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 of pre-training in neuroimaging at scale by training models on broad, public neuroimaging data so that the resulting models can be adapted well to the datasets collected by individual researchers, without the need for much new data.
A wealth of empirical evidence has demonstrated that models pre-trained on large neuroimaging datasets generally achieve better mental state decoding performance on new data, while requiring less 36th Conference on Neural Information Processing Systems (NeurIPS 2022). training time and data, when compared to models trained from scratch (e.g., [6, 7, 8, 9]).
Yet, much of this work has either pre-trained models on large but homogenous datasets, such as data from many individuals who all perform the same few tasks at the same few acquisition sites [10, 7], which can limit the generalizability of the resulting models due to potential systematic biases in homogeneous datasets (e.g., specific to the acquisition site or experimental paradigm [11, 12, 13, 14, 15]), or required highly-preprocessed input data (e.g., statistical maps summarizing the measured sequences of brain activity [16]).
In addition, researchers often use standard supervised pre-training tasks, training models to decode specific mental states assigned to each sample of the data [6, 7, 10, 17].
While this kind of supervised training can be fruitful within individual neuroimaging datasets, it is difficult to extend to many datasets, as neuroimaging researchers usually do not adhere to standardized labeling schemes for mental states [18] when assigning labels to the mental states of their experiments.
Due to this lack of standardization, it is often unclear whether two datasets contain the same or distinct mental states (for a detailed discussion, see [19]).
Here, we aim to overcome these limitations by leveraging recent advances in self-supervised learning to pre-train deep learning (DL) models on broad neuroimaging data.
In particular, we take inspiration from natural language processing (NLP), where the application of self-supervised learning has led to unprecedented breakthroughs in the adaptive capabilities of models trained on public text data (e.g., [20, 21, 22, 23]), and test whether modeling sequences of brain activity akin to sequences of text enables the learning of generalizable representations of brain activity and its dynamics.
To enable this, we represent sequences of brain activity akin to the embedded representation of words in sequences of text in NLP [24, 25] by representing each time point as a vector that describes whole-brain activity at this point in time as the activity of a set of functionally-independent brain networks [26].
Contributions.
By representing sequences of brain activity similar to the embedded representation of text in NLP, we are able to devise novel self-supervised learning frameworks for neuroimaging data that are inspired by prominent learning frameworks in NLP, namely, sequence-to-sequence autoencoding [27], causal language modeling [20], and masked language modeling [21].
We pre-train DL models with these frameworks on a large-scale neuroimaging dataset comprising fMRI data from 11 , 980 experimental runs of 1726 individuals in 34 datasets.
To our knowledge, this represents one of the broadest neuroimaging datasets used for pre-training to date, spanning a wealth of experimental conditions and a diverse set of acquisition sites.
We evaluate the downstream adaptation performance of the pre-trained models in two benchmark mental state decoding datasets and show that models pre-trained in a learning framework based on causal language modelling clearly outperform the others, while all pre-trained models generally outperform baseline models trained from scratch.
To enable others to build on this work, we make our code, training data, and pre-trained models publicly available 1 .
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-dimensional volumes V ∈ R x × y ×
z that show the BOLD signal intensity for each spatial location of the brain (as indicated by the three spatial dimensions x , y , and z ).
However, due to the strong spatial correlations of brain activity, the measured BOLD signal is often summarized differently by representing it as a set Θ = { θ 1 , ..., θ n } of n brain networks (i.e., parcels) θ , each describing the BOLD signal for some subset of voxels v x,y,z ⊂ V with correlated activity patterns.
Here, we utilize such a functional parcellation of brain activity to represent the measured BOLD signal, namely, Dictionaries of Functional Modes (DiFuMo), which was recently proposed by [28] and learned across millions of fMRI volumes.
DiFuMo defines Θ through a sparse dictionary matrix D ∈ R p × n containing n networks, each defined by p values (one for each voxel v ∈ V , flattened over the three spatial dimensions).
The BOLD signal of volume V t is then described as a linear combination of these networks by the use of weights α t ∈ 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 ∈ R n × p .
We utilize a parcelation with n = 1024 networks, such that the resulting BOLD parcelation X ∈ R t × n describes the measured BOLD signal of each of these networks at each time point t of the input sequence S .
This representation of the input data is thereby analogous to the embedded representation of sentences in NLP, where each word is represented by some embedding vector [24, 25].
2.2 Self-supervised learning frameworks: modeling sequences of brain activity akin to sequences of text
We first establish several commonalities among the frameworks before describing them individually in more detail (for an overview, see Fig. 1).
2.2.1 Commonalities among frameworks
First, each framework takes as input a parcellated BOLD sequence X ∈ R t × n (with t time points for n networks) as well as the corresponding sampling frequency (i.e., repetition time or TR) at which the fMRI sequence was collected (typically in the range of 0 .
7 s to 4 s).
Second, all frameworks linearly project the parcellated BOLD sequence X into an embedding representation E X ∈ R t × e : E Xt,e = b e + (cid:80) n X t,n w n,e , where w ∈ R n × e and b ∈ R e indicate weights and biases, thereby reducing the dimensionality of the input: R t × n → R t × e .
Third, to account for any differences in the sampling frequency of the BOLD sequences, the frameworks add one of r learnable TR embeddings E TR ∈ R r × e to each time point t ∈ E X .
Each TR embedding corresponds to one possible time point k of the input.
We set r = 1500 with a time increment of 0 .
2 s between subsequent TR embeddings, such that k ∈ 0 , 0 . 2 , ..., 300 s. For example, for an input sequence of 21 s that was collected at a TR of 0 .
7 s, we first round each time point t ∈ 0 , 0 .
7 , ..., 21 to the closest k and then add the corresponding E TRk to E Xk .
Note that we always label time points relative to the beginning of an input sequence, such that time point k = 0 is assigned to the first sequence element.
Importantly, TR embeddings carry different information than the position embeddings commonly used in transformer models [29] because the same position in an input can correspond to different time points, depending on the sampling frequency at which the data were collected.
Fourth, all frameworks utilize standard sequence-to-sequence DL model architectures f ( · ) that take E in as input and predict a corresponding output sequence f ( E in )
∈ R t × e .
Fifth, during upstream learning, the output sequence f ( E in ) is linearly projected back to the dimensionality of the parcellated BOLD sequence X , such that ˆ X t,n = b n + (cid:80) e f ( E in ) t,e w e,n , where w and b again indicate weights and biases.
This step is crucial, as all proposed learning frameworks involve reconstructing X (or parts of X ) after some of its elements have been masked or it has been encoded in a lower-dimensional representation.
To measure reconstruction performance, we utilize the mean absolute difference between the parcellated BOLD sequence X and its reconstruction ˆ X :
L rec = 1 n ( J ) (cid:80) j ∈ J | X j − ˆ X j | , where J indicates the set of elements that is to be reconstructed and n ( J ) the number of elements in J .
Sixth, for downstream adaptation to a new mental state decoding task, the frameworks utilize a simple decoding head p ( · ) , composed of a dense hidden layer with e model units (one for each embedding 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 x belongs to mental state i .
Consequently, all frameworks treat the decoding task as a multinomial classification problem and use a standard cross entropy loss objective: L cls = − (cid:80) 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 neural machine translation models (e.g., [27]).
These models take as input some text sequence, encode this sequence in a lower-dimensional representation, and predict an output sequence from the lower-dimensional encoding (e.g., a translation to another language).
Accordingly, the 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 [30]) units.
The encoder e ( · ) takes E in as input and encodes it in a lower-dimensional representation h ∈ R e .
The decoder d ( · ) then takes h as input and predicts a corresponding output sequence d ( h ) ∈ R t × e , such that f a ( E in ) =
d ( h ) .
Our encoder is composed of a stack of bidirectional LSTM units, whereas the decoder is build from a stack of unidirectional LSTM units.
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 units in a stack (with weights w i and w i +1 ).
At the t -th step of an input sequence, the inputs and outputs of units i and ( i + 1) are defined 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 ∈ R e represents the input to LSTM i at step t , whereas h it ∈ R e and c it ∈
R e represent hidden and cell state respectively.
We define the sequence encoding h that is forwarded to the decoder as the mean of the final hidden states of the two LSTM units contained in the encoder’s final bidirectional LSTM unit.
Upstream learning.
During pre-training, we jointly train the encoder and decoder to reconstruct the parcellated BOLD sequence X by minimizing its mean absolute difference to the reconstructed sequence ˆ X (see section 2.2.1): min 1 t (cid:80) t i =1 1 n (cid:80) n j =1 | ˆ X i,j − X
i,j | .
We also apply teacher-forcing [31] by using E in in 50% of the cases as each next input for the decoder instead of its own preceding output.
Adaptation.
To adapt the autoencoding framework to a new decoding task, we forward h as input to 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 language modeling [20], where models are trained to predict the next word in a text sequence.
Specifically, these models receive some text sequence as input, represented as a matrix X ∈ R n × e with n words each encoded by an e -dimensional vector, where the last sequence element (representing the last word) is masked, for example, by replacing it with a learned mask embedding E msk ∈ R
e (e.g., "When it rains, it [msk]").
The model’s task is then to predict the masked word.
In causal language modeling, models are thereby solely concerned with the part of the input sequence that precedes the masked word.
To adapt causal language modeling to neuroimaging data, we first mask the last time point T of a parcellated BOLD sequence X by replacing it with a learnable mask embedding E msk before transforming X to E in by means of linear projection and the addition of TR embeddings E TR .
In line with recent advances in causal language modeling, we utilize a standard transformer decoder model f d ( · ) (based on the GPT architecture [20]) in this learning framework 2 .
As transformer models are not aware of the ordering of their input, we further add learnable position embeddings E pos ∈ R p × e to E in , each representing one of p possible positions in the input sequence (we set p = 512 for all transformer models).
Given E in , the transformer decoder model predicts a corresponding output sequence f d ( E in ) ∈ R t × e in the form of the hidden states at the output of its last layer.
Upstream learning.
In line with causal language modeling, the upstream objective of the CSM framework is to reconstruct the masked last time point T of the parcellated BOLD sequence X by minimizing the mean absolute difference to its reconstruction ˆ X T : min 1 n (cid:80) ni
{+=
+} | ˆ X T,i − X T,i | .
Adaptation.
During downstream adaptation, we attach a learnable classification embedding[-E cls 
-]
E cls ∈ R n to the end of BOLD parcellation X and forward the resulting prediction f
( E in ) T +1 to a corresponding decoding head p ( · ) (see section 2.2.1).
2.2.4 Sequence-BERT
The Sequence-BERT framework is based on recent advances in bidirectional masked language modeling, particularly BERT [21].
BERT is trained to jointly solve a masked-language-modeling 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 of words have been randomly masked.
BERT is then asked to I) predict the masked words (masked-language-modeling) and II) determine whether S b follows S a (next-sentence-prediction).
In contrast to causal language modeling, masked language modeling is considered as a bidirectional learning task, as the model is not just concerned with the parts of the input that precede a masked word but also the parts that follow it.
To adapt BERT to neuroimaging data, we first sample two input sequences X a ∈ R t a × n and X b ∈ R t b × n from the data (with lengths t a and t b ).
In 50% of the cases, the two sequences represent sequential parts of the same underlying BOLD sequence, while in the other 50% of cases, the two sequences are randomly sampled from two distinct BOLD sequences.
Due to the strong temporal autocorrelation of BOLD data, we randomly leave a gap of 1 - 5 TRs (generally corresponding to a gap of 1 to 20 s) between the two sequences when sampling them from the same underlying BOLD sequence (to encourage the models to consider the entire input sequences and not just their connecting time points).
In line with BERT, we next randomly mask time points of each sequence (with a probability of 20% ) by replacing them with a learnable mask embedding E msk .
We also attach a separation embedding E sep ∈ R n to the end of each sequence (to indicate their endpoints) and a classification embedding E cls to the beginning of S a (for the next-sentence-prediction task; see upstream learning below) 3 .
The resulting versions of X a and X b are then linearly projected to obtain respective BOLD embeddings E X a and E X b to which we further add TR and position embeddings E TR and E pos as well as sequence ID embeddings E id ∈ R 2 ×
e to indicate each of the two input sequences.
The resulting embedding representation E in =
[ { E X a t
[Equation]
E id 1 } t ∈ t a , { E X b t + E TRt + E post + E id 2 } t ∈ t b ] is then forwarded to a standard transformer encoder model f e ( · ) (based on the BERT architecture [21]), which outputs a corresponding output sequence f e ( E in ) ∈ R t × e in the form of the hidden states at the output of its final layer.
Upstream learning.
Similar to BERT, we define the upstream objective of the Sequence-BERT framework as the sum of two learning objectives:
In line with masked-language-modeling, we define the first objective as the reconstruction loss L rec for all masked time points J : L rec =n ( J ) (cid:80) j ∈ J 1 n (cid:80) ni
=1 | X j,i − ˆ X j,i | .
In line with next-sentence-prediction, we define the other objective as the binary cross entropy loss for a decoding head that receives f ( E in ) 1 (corresponding to the transformer output for classification embedding E cls ) as input and predicts probability p that X b follows X a : L cls = − y log p − (1 − y ) log(1 − p ) , where y indicates a binary variable that equalsif the two sequences come from the same fMRI run and 0 otherwise.
Adaptation.
During downstream adaptation, we forward the transformer’s output for classification embedding 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 of a sequence by trying to reconstruct the parcellated whole-brain BOLD signal for specific time points.
Yet, the dynamics of brain activity data can also be viewed differently by focusing on the interaction of network activities over time instead of on the distribution of whole-brain activity at individual time points.
Consequently, we created a variant of the Sequence-BERT framework that randomly (at a rate of 10% ) masks the activity of individual networks n in the parcellated BOLD sequence X by replacing them with a learnable network mask embedding E msk ∈ R t instead of masking whole-brain activity at individual time points.
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 reconstruction performance for the masked network time courses
J : L rec =
[Equation]
| .
Adaptation.
During downstream adaptation, Network-BERT forwards the transformer’s output for classification embedding f e ( E in ) 1 to a respective decoding head p ( · ) (see section 2.2.1).
2.3 Training details
We train all models with stochastic gradient descent and the ADAM optimizer (with β 1 = 0 .
9 , β 2 = 0 . 999 , and ϵ = 1 e − 8 ) [32], if not reported otherwise.
We also apply a linear learning rate decay schedule (with a warmup phase of 1% of the total number of training steps), gradient norm clipping 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 the autoencoding, CSM, Sequence-BERT, and Network-BERT frameworks, respectively (based on the learning rates used in corresponding NLP studies [20, 21, 27]) and randomly sample sequences X of 10 to 55 TRs from our upstream fMRI runs.
For the Sequence- and Network-BERT frameworks, each sequence is split into two sequences X ai and X bi of equal length (with a randomly sampled gap of 1 to 5 TRs between the two sequences) and X bi randomly exchanged (with 50% probability) with another X bj (where i ̸ = j ).
During downstream adaptation, we begin training with the pre-trained model parameters and then allow all parameters to be changed freely during training.
Input sequences are drawn from the downstream fMRI runs according to the on- and offset of each experimental trial (with durations generally ranging between 2 and 30 s), when accounting for the temporal delay of the hemodynamic response function.
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, while downstream adaptations were performed on compute nodes of the Texas Advanced Computing Center with one Nvidia 1080-TI GPU, 32 CPU threads, and 128 GB RAM memory.
Training times varied depending on model size but were generally in the range of 1 to 3 days for upstream learning and 0 .to 3 hours for downstream adaptations.
Experiments 3.1 Datasets
All unprocessed fMRI data used in this study are publicly available through OpenNeuro.org [3] and the Human Connectome Project (HCP [33]).
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 runs from 1 , 726 individuals and 34 datasets (see Fig.
2 and Appendix Table A1).
To our knowledge, this represents one of the broadest fMRI datasets used for pre-training in neuroimaging to date, spanning many acquisition sites and a diverse set of experimental conditions and domains.
A few prominent examples of included datasets are The Midnight Scan Club [34], Individual Brain Charting [35], Amsterdam Open MRI collection [36], BOLD5000 [37], and the Narratives collection [38].
We randomly split the upstream data into distinct training and evaluation datasets by randomly designating 5% of the fMRI runs of each included fMRI dataset as evaluation data (at a minimum ofruns per dataset) and using the rest of the runs for training.
At each evaluation step, we randomly sample 640 , 000 sequences from the evaluation dataset.
3.1.2 Downstream: benchmark datasets with many mental states
To evaluate the adaptation performance of the pre-trained models, we utilize two benchmark mental state decoding datasets that span a wide range of mental states, namely, task-fMRI data from the HCP [33] and the Multi-Domain Task Battery (MDTB [39]).
Our HCP dataset spans 100 participants and 19 distinct mental states across seven experimental tasks (see Appendix A.2). For each task, two fMRI runs were collected.
We also include two resting state fMRI runs per HCP participant, in which individuals simply rest in the fMRI without any particular task, increasing the total number of mental states to 20 .
The MDTB dataset spans 24 individuals and 26 experimental tasks, which participants performed across two scanning sessions with eight fMRI runs per session.
Due to the strong similarity of the experimental conditions within individual experimental tasks of the MDTB, we define each of its tasks as one target mental state (see Appendix A.2).
3.2 Hyper-parameter evaluation: larger embeddings but not deeper models
In our first experiment, we compared the performance of different hyper-parameter settings for the proposed frameworks.
Specifically, we evaluated
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 [40], we scaled the number of attention heads n α per layer of the transformer models with the embedding dimension, such that n α = e 64 .
We trained one model for each point of the resulting 9 -point grid of each framework for 200 , 000 training steps, using a mini-batch size of 256 random sequences.
Overall, model performances improved (in terms of validation loss) with larger embedding dimensions but not with the number of hidden layers (Fig. 3).
We therefore decided to continue all further analyses with models comprising 768 embedding dimensions and 2 hidden layers for the encoder and 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 our second experiment, we fully pre-trained one model in each framework according to our previously selected set of hyper-parameters (see section 3.2) by training models for 300 , 000 training steps at a mini-batch size of 768 sequences.
All models learned well over the course of their training, with gradually decreasing training and validation losses (Fig. 4).
Note that the transformer model trained in the CSM framework slightly overfitted the upstream data, whereas the gap in training and evaluation performance of the autoencoding framework results from the application of teacher-forcing during training (see section 2.2.2).
We also evaluated the reconstruction performance of the final models for different parts of the brain and found that they exhibited similar overall distributions of reconstruction error, with relatively higher errors in the posterior parietal, occipital, and cingulate cortices as well parts of the limbic system (see Appendix B.2).
3.4 Adapting pre-trained language models does not yield upstream performance gains
Given that the formatting of our input data is highly similar to the embedded representation of sentences in language modeling, where each word of a sentence is represented by an embedding vector [41, 21], we next tested whether adapting pre-trained language models to our upstream data would yield meaningful performance gains, when compared to training models from scratch.
To this end, we adapted pre-trained language model variants of GPT-2 [41] and BERT [21] (with 12 hidden layers and an embedding dimension of 768 ; as provided by Hugging Face’s Transformer library [42]) to our upstream dataset (see Appendix B.3).
The upstream validation performance of these two language models did not meaningfully improve upon the performance of our models trained from scratch (see section 3.3), at final upstream validation losses of 0 .
73 (CSM with GPT-2) and 0 .(Sequence-BERT with BERT) (Appendix Fig. B3).
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 (HCP and MDTB; see section 3.1.2).
To gauge the effectiveness of pre-training, we adapted the pre-trained models to varying sizes of the downstream datasets by randomly assigning individuals to distinct training, validation, and test datasets.
Specifically, for each evaluated dataset size, we first randomly selected 10 (HCP) and(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 number of individuals whose data we use for training ( 1 to 48 (HCP) and 11 (MDTB)) from the remaining pool of individuals.
We then adapted each pre-trained model with varying learning rates and numbers of training steps to the training data (see Appendix B.4) and used the model variant that achieved the highest final mental state decoding accuracy in the validation data for any further analyses (for an overview of validation decoding accuracies, see Appendix Fig.
B6 and B7).
For comparison, we also trained a linear baseline model (inline with the state-of-the-art in mental state decoding [43]) and a variant of our decoding head p ( · ) that we applied directly to the parcellated BOLD sequences X (see Appendix B.5). [Table]
[Table] The pre-trained models clearly outperformed the linear baseline model across all training dataset sizes, while their performances also scaled well with the number of individuals in the training data (Tablesand 2; at chance-levels of 5 . 0% (HCP) and
9% (MDTB)).
Overall, the transformer decoder model trained in the CSM framework clearly outperformed the other models by achieving the highest test decoding accuracies throughout all evaluated sizes of the training dataset.
The other pre-trained models performed on par with the decoding head p ( · ) in smaller training dataset sizes ( ≤ 6 (HCP) and ≤ 3 (MDTB) individuals), while clearly outperforming it in larger training datasets.
For the
HCP data, for which other reported decoding performances exist in the literature, the pre-trained models performed on par with models trained on much larger datasets (often comprising data of many hundred individuals [10, 7, 6, 44]), which generally report test decoding accuracies between 80% and 93% .
A feature ablation analysis of the models’ mental state decoding decisions in both datasets further revealed that the these are strongly dependent on the activity of the occipital and inferior temporal cortex as well as parts of the pre- and postcentral gyrus (see Appendix B.6).
We also tested whether the reported downstream model performances are stable over the nondeterministic aspects of their training [45] by replicating our adaptation analysis with different random seeds and splits of the data (see Appendix B.7).
This replication confirmed our results, yielding final test decoding accuracies that were not meaningfully different from our initial analysis (Appendix Tables A3, A4, A5, A6).
Conclusion
In this work, we propose a set of novel self-supervised learning frameworks for neuroimaging data inspired by prominent learning frameworks in NLP.
By the use of these frameworks, we are able to pre-train DL models on broad, public neuroimaging data, comprising many individuals, experimental domains, and acquisition sites.
The pre-trained models generalize well to benchmark mental state decoding datasets and generally outperform baseline models trained from scratch, while models pre-trained in causal sequence modelling clearly outperform the others.
We hope that this work will inspire others to explore pre-training in functional neuroimaging at scale, using self-supervised learning.
Limitations.
This work neglects a key aspect of mental state decoding, namely, the ability to draw inferences about the association between the decoded mental states and the underlying brain activity from the trained models.
First empirical evidence indicates that attribution methods from explainable artificial intelligence (XAI [46]) research are well-suited to provide insights in the mental state decoding decisions of DL models [19].
Similarly, this work does not provide any insights into how the proposed self-supervised learning frameworks compare to contrastive learning techniques, which have demonstrated empirical success in computer vision [47] and medical imaging [48].
Potential negative social impact.
We are currently not aware of any direct negative social impacts that could follow from pre-training DL models on public neuroimaging data that are de-identified and collected and shared with human consent in a way approved by institutional review boards. {+Acknowledgments. We gratefully acknowledge the support of NIH under No. U54EB (Mobilize), NSF under Nos. CCF
+}
(Beyond Sparsity), CCF1563078 (Volume to Velocity), and 1937301 (RTML); ARL under No.
W911NF-21-2-0251 (Interactive Human-AI Teaming);
{+ONR under No. N
+}
(Unifying Weak Supervision); ONR N00014-20-1-2480: Understanding and Applying Non-Euclidean Geometry in Machine Learning; N000142012275 (NEPTUNE); NXP,
Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASF, Accenture,
Ericsson, Qualcomm, Analog Devices, Google Cloud, Salesforce, Total, the HAI-GCP Cloud
Credits for Research program, the Stanford Data Science Initiative (SDSI), the Texas Advanced
Computing Center (TACC) at The University of Texas at Austin, and members of the Stanford
DAWN project: Facebook, Google, and VMWare.
The OpenNeuro data archive is supported by NIH grant R24MH117179.
The U.S. Government is authorized to reproduce and distribute reprints for
Governmental purposes notwithstanding any copyright notation thereon.
Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of NIH, ONR, or the U.S. Government.
A Data details A.1 Upstream datasets
Appendix Table A1 provides an overview of the datasets that we included in our upstream training [49, 50, 51, 52, 53, 54, 34, 55, 56, 37, 57, 58, 59, 60, 61, 62, 38, 35, 36, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72].
The unprocessed fMRI data of all datasets are publicly available (under a Creative Commons
CC0 license) through OpenNeuro.org [3] under the specified identifier (ID) of each dataset.
All fMRI data were de-identified and collected and shared with human consent in a manner approved by institutional review boards.
We did not use any personally identifiable data
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
A.2 Downstream mental states overview
We provide a brief overview of the mental states included in both downstream datasets below.
For any further details on the experimental procedures of the datasets, we defer the reader to the original publications [33] (HCP) and [39] (MDTB).
HCP.
Appendix Table A2 provides an overview of the mental states of each HCP experiment task.
[Table]
MDTB.
The MDTB dataset includes the following set of tasks, each representing one mental state in our analyses (as labeled by the original authors): CPRO, GoNoGo, ToM, actionObservation, affective, arithmetic, checkerBoard, emotionProcess, emotional, intervalTiming, landscapeMovie, mentalRotation, motorImagery, motorSequence, nBack, nBackPic, natureMovie, prediction, rest, respAlt, romanceMovie, spatialMap, spatialNavigation, stroop, verbGeneration, and visualSearch.
A.3 FMRI data preprocessing
We preprocessed all fMRI data with fMRIPrep (versions 20.2.0 and 20.2.3; a minimal, automated preprocessing pipeline for fMRI data [5]) using fMRIPrep’s default settings without FreeSurfer [73] surface preprocessing.
We then applied a sequence of additional minimal processing steps to fMRIPrep’s derivatives, which included I) spatial smoothing of the fMRI sequences with a 3mm full-width at half maximum Gaussian Kernel, II) detrending and high-pass filtering (at 0.008 s) of the individual voxel activity time courses, and III) basic confound removal by regressing out noise in the data related to head movement (as indicated by the six basic motion regressors x , y , z , roll, pitch, and yaw) as well as the mean global signal and mean signal for white matter and cerebrospinal fluid masks (as estimated by fMRIPrep).
Lastly, we parcellated each preprocessed fMRI run with the
DiFuMo
atlas (see section 2.2 of the main text) and standardized the resulting individual network time courses to have a mean of 0 and unit variance.
B Experiment details B.1 Hyper-parameter evaluation run for largest Sequence-BERT model
The largest model variant of the transformer encoder model that we trained in the Sequence-BERT framework during our hyper-parameter evaluation (see section 3.2 of the main text), with 12 hidden layers and an embedding dimension of 768 , did not meaningfully learn over the course of its training.
For better visibility of the other models’ performances, we decided to not include the model’s training run in Fig.
and are instead showing it in Appendix Fig.
B1.
B.2 BOLD reconstruction error in validation data
To evaluate the BOLD reconstruction performance of our pre-trained models (see section 3.3 of the main text) for different parts of the brain, we computed each models’ mean reconstruction error ( L rec ) in the upstream validation data for each network of our BOLD parcellation (see section 2.1 of the main text) and projected the average network reconstruction errors back to the voxel-level (as described in section 2.1 of the main text).
The four pre-trained models exhibit similar distributions of reconstruction error throughout the brain (Appendix Fig.
B2), with relatively higher errors in the posterior parietal, occipital, and cingulate cortices as well parts of the limbic system.
B.3 Adaptation of pre-trained language models
We adapted pre-trained language model variants of GPT-2 [41] and BERT [21] (as provided by
Hugging Face’s Transformer library [42]) to our upstream data in two training phases, using the CSM and Sequence-BERT frameworks, respectively (Appendix Fig.
B3):
First, we froze all parameters of the two language models and trained them for 20 , 000 training steps at a mini-batch size of 512 and a learning rate of 1 e − 4 , bringing all other parameters of the CSM and Sequence-BERT frameworks into sensible ranges.
After this warmup phase, we continued training both models for a total of 150 ,training steps of the same mini-batch size, using learning rates of 5 e − 4 (GPT-2) and 5 e − 5 (BERT) respectively, while allowing all model parameters to change freely.
B.4 Downstream adaptation of pre-trained models
For each downstream application of our pre-trained models, we evaluated two learning rates ( 1 e −and 5 e − 5 ) and different training lengths.
Specifically, we scaled the number of training steps N step according to the number of subjects N sub in the training data, such that: N step = 1000 + ηN sub , using two η -values ( 150 and 300 ) (see Appendix Fig. B6 and B7).
We only report test decoding accuracies for the model variant achieving the highest validation decoding accuracy in this 4-point grid search in the main text (see section 3.5 of the main text).
B.5 Baseline models
Linear baseline.
The linear baseline model forms a decoding decision in two steps: it first aggregates the time course of each network n ∈ X : a n = b n + (cid:80) t X t,n w t,n (with w and b indicating weights and biases) and then predicts a probability p ( a ) i that X belongs to mental state i from the distribution of aggregated time courses a ∈ R n : p ( a ) i
=
σ ( b i + (cid:80) n a n w n,i ) , where σ represents the softmax function and w and b indicate a second set of weights and biases.
Decoding head p ( · ) .
We also evaluated the performance of our decoding head p ( · ) (see section 2.2.1 of the main text) when applied directly to the parcellated BOLD data X ∈ R t × n .
To allow for the application of p ( · ) to inputs of varying length, we first averaged the signal of each network n over its time course ( X n = 1 t (cid:80) ti =1 X i,n ) before forwarding the time-averaged signal X to the decoding head to make a decoding decision for each mental state i : p ( X ) i .
Training.
Similar to the other frameworks, we trained both baseline models to minimize a standard cross-entropy loss with additional L 2 -regularization:
L cls = − (cid:80) i
y i log p i + λ (cid:80) i
w 2 i , where y i indicates a binary variable that equals 1 if i is the correct mental state and 0 otherwise, while λ scales the L 2 -regularization strength.
Hyper-parameter evaluation.
For each application of the baseline models, we evaluated two learning rates ( 1 e − 3 and 1 e − 4 ), three L 2 -regularisation strengths ( 0 . 1 , 1 , and 10 ), and two training lengths ( 1000 and 3000 training steps at a mini-batch size of 512 ; Note that the baseline models generally required fewer training steps to converge than the pre-trained models; see Appendix Fig.
B6 and B7) in a 12 -point grid search.
As for the pre-trained models, we only report the test decoding accuracy of the model variant achieving the highest validation accuracy in the main text (see section 3.5 of the main text for details on the data split and Appendix Fig.
B6 and B7 for an overview of model performances).
B.6 Downstream feature ablation analysis
To understand which parts of the brain were most relevant for the mental decoding decisions of the adapted models, we performed a feature ablation analysis of the test mental state decoding decisions of the best-performing models (that were adapted to the largest training dataset size of both downstream datasets; see Tables 1 and 2 of the main text).
Specifically, for each sample of the test datasets, we replaced the signal of individual networks with random Gaussian noise from a standard normal distribution and measured the effect on the model’s decoding prediction for the sample’s mental state label.
Note that we evaluated the predicted logits prior to their softmax scaling.
This analysis revealed that the decoding decisions of all pre-trained models are strongly dependent on the signal of the occipital and inferior temporal cortex as well as parts of the pre- and postcentral gyrus in both datasets (see Appendix Fig. B4 and B5).
we ablate the time course of individual networks and measure the resulting effect on the models’ prediction for the mental state associated with the sample.
Average changes in output are projected onto the inflated cortical surface of the FsAverage template [73].
B.7 Replication of downstream adaptation analysis
To test for the stability of the reported model performances over the non-deterministic aspects of their training (such as different random weight initializations or random shufflings of the data during training), we replicated our downstream adaptation analysis (see section 3.5 of the main text) with different random splits of the data and different random seeds per split.
This replication confirmed our initial results (compare Tables 1 and 2 of the main text with Appendix Tables A3 and A5).
We also tested whether the final test decoding accuracy of each model training run of the analysis was meaningfully different between our initial analysis (see Appendix Fig.
B6 and B7) and the replication (see Appendix Fig.
B8 and B9)
by computing the difference in final test decoding accuracy between each initial training run and its replication and testing the resulting distribution of test decoding accuracy differences against a mean of 0 in a two-sided t-test (Appendix Tables A4 and
A6).
The models’ final test decoding accuracies were not meaningfully different between our initial analysis and the replication, indicating strong stability of the models’ performance over the various non-deterministic aspects of their training.
[Table]
[Table]
[Table]
[Table]