BrainCodec: Neural fMRI codec for the decoding of cognitive brain states

In this section, we delineate the architecture of the model, focusing on the parameters without delving into the specifics of parameter selection. For further details on parameter choices, refer to A.2.

The critical hyperparameters are summarized in Table 5. As illustrated in Figure 1, BrainCodec is an AutoEncoder-based model consisting of three primary components: the Encoder, the Decoder, and the Residual Vector Quantization (RVQ).

The Encoder is primarily composed of ConvTransposedBlocks. Each ConvTransposedBlock consists of a Transposed Convolution Layer (Dumoulin and Visin [2016]) and a Residual Layer. In the Transposed Convolution Layer, operations are performed to expand in the temporal direction when the stride is greater than one. Additionally, a standard Convolution Layer is inserted to match the input and output dimensions, and an LSTM is applied to the output of the ConvTransposedBlock. Summarizing the transformations, the shape of the input fMRI data after applying DiFuMo (1024, 50) undergoes the following changes:

The Decoder, on the other hand, mirrors the Encoder’s structure. Each module is designed symmetrically with reversed filter sizes and strides, and the LSTM is applied first. Unlike the Encoder, the Decoder compresses in the temporal direction using standard Convolution Layers instead of Transposed Convolution Layers.

Finally, the RVQ is an extension of the widely used Vector Quantization (VQ) module, concatenating multiple VQ modules (eight in this case). Each VQ’s input is the difference between the input and output of the previous VQ. This hierarchical modeling allows different codebooks to capture distinct features, akin to frequency decomposition. In the audio domain, the first VQ module tends to capture low-frequency semantic features, enabling techniques like teacher forcing to explicitly learn these semantic features, as proposed in codecs such as the SpeechTokenizer (Zhang et al. [2023]). Similar trends are observed in our study, as demonstrated in Section 4.2.

In this section, we present the important hyperparameters considered in our study, accompanied by loss graph plots. It is important to note that while each individual experiment offers a valid comparison, the overall comparisons are not consistent across all experiments. Therefore, direct comparisons between separate experiments should be approached with caution. Additionally, since the trajectory of the loss typically stabilizes early in the training process, some experiments were halted at the initial stages.

Figure 6(a): Strides
Our experiments revealed that increasing the temporal dimension in the encoder leads to a greater reduction in reconstruction error. This is a logical outcome, as it essentially corresponds to a reduction in effective compression rate, making reconstruction easier. Considering this trade-off, we opted for a factor of 3 for temporal expansion. This choice is analogous to the approach used in Encodec, where the temporal dimension is compressed by a factor of 320, while the spatial dimension is compressed by a factor of 128, resulting in an effective compression rate of approximately 2.5. In BrainCodec, we used the spatial dimension compression to 128, as in Encodec, yielding an 8-fold compression. To achieve an effective compression rate closest to 2.5, we chose to expand the temporal dimension by a factor of 3.

Figure 6(b): Loss function
By default, Encodec employs L1 loss, which is a very common setting in the audio domain. However, this setting is not self-evident for fMRI data. In fact, in reconstructions using VAE, L2 loss is also often employed (Kim et al. [2021]). Therefore, we conducted a comparison and found that using L2 loss resulted in a lower reconstruction error for L2 loss, while L1 loss increased correspondingly, and vice versa. This outcome is quite natural. Given that it is unclear which loss function is better based on reconstruction error alone, we compared them in terms of downstream performance. As shown in the figure, L2 loss contributed slightly more to improvement. L2 loss imposes a heavier penalty on outliers, which could be beneficial for the noisy nature of fMRI data that may contain many outliers.

Figure 6(c): Discriminator
This is likely the most significant factor contributing to the differences observed. We concluded not to use the traditionally well-utilized discriminator. The reason, as illustrated in the figure, is that it exacerbated the reconstruction error. This issue was not resolved by simply adjusting the coefficients between the loss functions. This instability is believed to stem from the inherently noisy nature of fMRI data. Reconstruction error, by definition, calculates loss based on the difference across all data points between input and output. Consequently, when inputs like fMRI data, which are significantly affected by noise, are processed, the learning progresses towards replicating that noise in the output. While discriminators generally aim to enhance fidelity, employing them leads to learning that replicates a similar noise output from noise inputs, essentially attempting to reproduce the noise distribution. Therefore, this compatibility issue between the reconstruction error, which aims to fit to one realization of noise distribution, and the use of discriminators led to unstable learning outcomes. To improve upon this, approaches such as the introduction of noise-robust GANs (Kaneko and Harada [2020]) might be considered, though this remains a topic for future research.

Additionally, as shown in Figure 6(d), changing the number of ConvBlocks confirms that with fewer layers, overfitting occurs rapidly. Furthermore, altering the dimensions of the codebook (Figure 6(e)) did not result in significant differences, so we used the default 128 dimensions. Lastly, we also present the results of changing the number of codebooks (Figure 6(f)). Naturally, the reconstruction error decreases as the number of codebooks increases. However, increasing the number of codebooks reduces the difference from continuous value compression and makes the analysis significantly more complex. Therefore, we adopted the default of 8 codebooks for our proposed method.

Here, we present the reconstruction errors and actual examples of reconstructions for BrainCodec, BrainCodec large, and BrainVAE, which were used in Section 4.1.

Figure 7 shows the reconstruction errors. BrainVAE, which uses continuous latent representations, and BrainCodec large, which uses a very high number of codebooks (128), have significantly lower reconstruction errors compared to BrainCodec. However, as discussed in Section 4.1, there are still cases where improvements depend on the task. This indicates that further improvements are possible, though it is unclear whether this should focus on further reducing reconstruction errors or introducing discriminators. This remains a topic for future research.

Figure 8 displays the fMRI images reconstructed by each method alongside the ground truth images. The images reconstructed by BrainCodec are noticeably rough and fail to capture finer details. In contrast, the reconstructions by BrainCodec large and BrainVAE, as reflected in their lower loss, manage to preserve finer details more accurately.

While these qualitative results are not conclusive, it is noteworthy that discretizing with only 8 codebooks, as in BrainCodec, leads to rougher reconstructions. This suggests that using a higher number of codebooks, as in BrainCodec large, or continuous latent representations, as in BrainVAE, may be more effective for capturing detailed brain activity in fMRI data.

By using BrainCodec, we can convert fMRI data into a token sequence. This approach allows us to treat fMRI data similarly to text information, enabling the full utilization of NLP frameworks. Specifically, we perform next token prediction using the obtained token sequence as input and train the model with cross-entropy loss (see Fig 9). We refer to the CSM trained in this manner as TokenCSM. This method is also commonly used in the field of audio area (Wang et al. [2023], Copet et al. [2024]).

Here, we describe how to train CSM in a token-based manner by leveraging BrainCodec’s ability to tokenize fMRI data. This technique is well-known for handling raw data that is difficult to work with directly, such as audio (Wang et al. [2023], Łajszczak et al. [2024]).

Using the Encoder and RVQ of BrainCodec, DiFuMo-transformed fMRI data with a shape of (1024, 50) can be converted into a token sequence with a shape of (8, 150). Typically, in Transformer training with text, the token sequence shape is a one-dimensional array like (150,). The hierarchical tokenization introduced by RVQ must be managed appropriately. MusicGen (Copet et al. [2024]) discusses how to input these hierarchical token sequences into a Transformer. It shows that the "delay pattern", which inputs the token sequences offset by one token per hierarchy, best balances efficiency and performance. However, our experiments found that the "flatten pattern", which simply arranges the token sequences in a single line, yielded the best results (see Fig 10(b)).

Figure 9 illustrates the overall process of TokenCSM’s upstream learning. Note that, since the input and output are token sequences, the training objective is cross-entropy loss rather than reconstruction error. For downstream learning, as described in Section 2.2.2, we append a CLS token at the end and use this CLS token as the feature for solving the classification task.

For TokenCSM training, the parameters for upstream learning were the same as CSM (see Section 3.1), except the batch size was 32 and the learning rate was $7\times 10^{-4}$. This training lasted about 56 hours. For downstream learning, the only change made compared to the first method was the batch size, which was decreased to 64. This phase of training took approximately 8 hours. This configuration results in approximately 100 million parameters. Compared to the tens of billions of parameters commonly used in NLP (Touvron et al. [2023]) and the hundreds of millions to billions used in speech models (Łajszczak et al. [2024]), this is still small.

Here, we introduce the loss curves for TokenCSM when hyperparameters are adjusted. Specifically, we illustrate the effects of varying the number of layers and token rearrangement patterns.

Figure 10(a) shows the loss curves when the number of layers is varied. It was observed that increasing the number of layers consistently reduces the loss. Naturally, increasing the number of layers, and hence the parameter size, necessitates a corresponding increase in data volume. Since the data volume was fixed in this study, we selected 24 layers, the maximum value shown here.

Figure 10(b) depicts the changes in loss when the token rearrangement patterns "delay" and "flatten" are applied. As shown, the "flatten" pattern significantly outperforms the "delay" pattern. However, it is important to note that the "flatten" pattern uses approximately 8 times more token length than "delay," resulting in increased computational resources and time.

Moreover, regardless of the method, the absolute loss value remains around 6. Initially, the loss was approximately 6.8, so this represents only a reduction of 0.8. In text-to-speech systems, which operate on similar principles, the cross-entropy loss typically ranges from 3 to 4. This indicates that the current method has considerable room for improvement. The most effective and reliable way to achieve this improvement is to increase both the data volume and parameter size. Collecting more data remains a significant future challenge for this field.

The downstream training for TokenCSM is conducted in exactly the same manner as the training method described for CSM in Section 2.2.2.

Table 6 shows the downstream performance of TokenCSM on the HCP and MDTB datasets. the performance significantly deteriorated from the baseline in both the HCP and MDTB datasets under the current settings. This is believed to be due to the sheer volume of data and the size of the model. While the data used in this study was limited to about 10,000, with a parameter size of 100M, the initial CSM-based method, GPT-2 (Radford et al. [2019]), had a parameter size of 1.5B and utilized 40GB of data from vast amounts of webpages. If the goal is solely to improve downstream task performance, such large models and datasets might not be necessary. However, the absolute values of the loss in upstream learning for TokenCSM were not only large but also showed minimal reduction (Appendix B.3). Expanding the fMRI dataset size by tens of times in the future and applying such large-scale methods will be an important direction for further improving accuracy.

A brief overview of the mental states included in both downstream datasets is provided below. This information is also cited from Thomas et al. [2022]. For additional details on the experimental procedures of these datasets, readers are referred to the original publications: Van Essen et al. [2013] for HCP, King et al. [2019] for MDTB, and Tomoya Nakai [2020] for Over100.

HCP. Table 8 provides an overview of the mental states of each HCP tasks.

MDTB. The MDTB dataset comprises the following set of tasks, each representing a distinct 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.

Over100. In the Over100 dataset, tasks related to natural behaviors are designed as follows: PressRight, PressLeft, PressLR, RestOpen, RestClose, EyeBlink, RateTired, RateConfidence, RateSleepy, ImagineFuture, and 93 other tasks.

To conduct the analysis described in Section 4.2.1, we have constructed a new resting-state fMRI dataset. Following the methodology outlined by Thomas et al. [2022], we gathered data from OpenNeuro.org. The details of this dataset are presented in Table LABEL:tab:resting_state_dataset. The numbers of subjects and runs listed reflect those that were successfully preprocessed, excluding any that encountered errors during this stage.

In addition to these datasets, we utilized the "SRPBS Traveling Subject MRI Dataset," which comprises 411 scans of 3T MRI imaging data from 9 traveling subjects collected across 9 sites. For more details, refer to Tanaka et al. [2021]. Note that due to redistribution restrictions, our shared dataset does not include this specific dataset. From this dataset, we obtained a total of 410 runs from the 9 subjects.

Data used in the preparation of this work were obtained from the DecNef Project Brain Data Repository (https://bicr-resource.atr.jp/srpbsts/) gathered by a consortium as part of the Japanese Strategic Research Program for the Promotion of Brain Science (SRPBS) supported by the Japanese Advanced Research and Development Programs for Medical Innovation (AMED)

As previously mentioned, we utilized the datasets distributed by Thomas et al. [2022] for both upstream and downstream learning. However, for the resting state dataset, we prepared new data ourselves. In doing so, we adopted the same preprocessing methods as Thomas et al. [2022]. The details are outlined below.

All fMRI data were preprocessed using fMRIPrep (versions 20.2.0 and 20.2.3; a minimal, automated preprocessing pipeline for fMRI data (Esteban et al. [2019])) with its default settings, excluding FreeSurfer (Fischl [2012]) surface preprocessing. Subsequently, we applied a series of additional minimal processing steps to the fMRIPrep derivatives, including 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 Hz) of the individual voxel activity time courses, and iii) basic confound removal by regressing out noise 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 signals for white matter and cerebrospinal fluid masks (as estimated by fMRIPrep). Finally, each preprocessed fMRI run was parcellated using the DiFuMo atlas (see section 2.2 of the main text) and the resulting individual network time courses were standardized to have a mean of 0 and unit variance.

Here, we present the loss curves of upstream learning for the CSM models ultimately used in downstream learning.

Figure 11(a) shows the loss curves for the original CSM model, identical to the one used in the previous study by Thomas et al. [2022] (labeled as Raw in the figure), and for CSM combined with various codec models: CSM+BrainCodec, CSM+BrainCodec large, and CSM+BrainVAE. As shown in the figure, the Raw model exhibits the highest loss, while CSM+BrainCodec has the lowest. This result aligns with our motivation, demonstrating the beneficial effect of compression methods in facilitating model training. The Raw model is noisier and thus more difficult to predict accurately, leading to poorer generalization of CSM. CSM+BrainCodec large and CSM+BrainVAE, while able to reconstruct higher-resolution images due to lower compression ratios, result in increased loss during training.

Figure 11(b) presents the results for TokenCSM. Unlike CSM, TokenCSM operates on a token basis, preventing the use of BrainCodec large, which has a high number of codebooks, and BrainVAE, which is not a discretization method. Additionally, raw data cannot be tokenized, leaving us with a single result from tokenization using BrainCodec. While there are no comparisons here, the noteworthy aspect is the absolute value of the loss. The cross-entropy loss only decreases to around 6, which is quite large compared to other fields. For a detailed discussion on this, refer to B.3.

As explained in Section 2.1, CSM is a model that predicts the current time-step data using past data. By using this model, we can input a portion of fMRI data and output the continuation. In this study, we used half of the fMRI data in an evaluation dataset as a prompt (condition) and generated the subsequent fMRI data.

Figure 12 shows examples of these generated sequences. When visualizing, the first two steps (the top two images in each figure) are used as the prompt. The subsequent images are generated data. Additionally, the ground truth fMRI data is displayed for comparison (left column).

First, compare Figures 12(b) and 12(a). These figures show the results for CSM, which is the CSM trained using only raw fMRI data. As demonstrated by the high loss shown in Section D.1, the generated data significantly deviates from the ground truth.

Next, compare Figures 12(d) and 12(c). These show the results for CSM+BrainCodec. Note that the ground truth data (Figure 12(d)) is already reconstructed data. At the initial step following the prompt, the generated data closely matches the ground truth. However, as the sequence progresses, the generated data increasingly diverges from the ground truth. This divergence is likely due to the autoregressive nature of the CSM, where errors accumulate during sequential generation. Alternatively, this divergence could represent a plausible progression of fMRI data, but this was not verified in this study.

Finally, compare Figures 12(f) and 12(e). These figures show the results for CSM+BrainVAE. The generated results are more detailed than those using BrainCodec due to the use of BrainVAE. However, discrepancies from the ground truth are noticeable even from the step immediately following the prompt.

In conclusion, given the current data volume and model configurations, using BrainCodec is likely to provide a level of learning capable of generating plausible sequences. Of course, this conclusion cannot be drawn from a single example. For instance, generating data for downstream learning and performing classification on this generated data could objectively compare the quality of the generated sequences. This remains a subject for future research.

As mentioned in Section 4.1.2, the performance of the Linear model was highly influenced by weight decay, likely due to the model’s small size. In Section 4.1, we could only include part of the results due to space constraints. Here, we present the complete results.

First, consider the results for the HCP dataset shown in Table 10. It is important to note that as we selected the best method for each approach to draw the previous conclusions, the conclusions remain unchanged here. One noticeable point we can find from this table is the clear advantage of the compression methods. For instance, in the Linear model without any codec, the best result is achieved with a weight decay of 1. This is natural since the model is very small, the data is noisy, and some weight decay helps prevent overfitting by acting as a regularization term. On the other hand, the Linear+BrainCodec model achieves the best performance with the smallest weight decay. This is because BrainCodec compresses and reconstructs the data, retaining important information and making it easier for the model to learn. Hence, in the HCP dataset, the codec method proves to be a very powerful tool. We can see the same trend with MDTB dataset.

In contrast, the results for the Over100 dataset shown in Table 10 are somewhat different. For example, while the smallest weight decay yields the best results in the baseline, it is not the best for BrainCodec, which contrasts with the trends observed in HCP and MDTB. This is likely due to the higher difficulty of the dataset. Since Over100 involves classifying over 100 classes, it is much more challenging than other datasets, and even with BrainCodec, there is a risk of overfitting. Therefore, the smallest weight decay is not optimal. In contrast, in the baseline models, there is little difference between weight decay values of 1 and 0.1, which suggests that the difficulty of the problem is so high that overfitting is not yet an issue.

Thus, neuroscience datasets exhibit significantly different characteristics, and improved generalization performance can be expected through the collection of even larger datasets in the future.

In Section 4.2.1, we used UMAP to visualize the codebooks learned by BrainCodec from resting state and task-related brain activities. This was done to verify the neuroscientific knowledge that brain activity during resting state encompasses task-related brain activity. We noted that similar patterns were observed in the second and subsequent codebooks. Here, we present the complete visualization of all codebooks in Figure 13. As expected, the overall conclusions remain unchanged.

To quantify the similarity between codebooks, we employed the Sinkhorn distance (Cuturi [2013]), calculated using an L2 distance cost matrix that underwent Min-Max normalization, allowing codebook-wise absolute comparisons. The Sinkhorn algorithm was then applied to compute the Sinkhorn distance.

The results, displayed in Fig 14, indicate that when comparing Task-codebook and Rest-codebook, the overall Sinkhorn distance is smaller than a reference (distances between Random-codebooks and distances to the Random-codebook). RVQ inherently clusters low-frequency components, meaning that major variations typically aggregate in the lower codebooks (A.1). The close distances between the first and second codebooks may suggest that the primary variations in resting state fMRI and task fMRI are relatively similar. Additionally, the fact that higher codebooks have distances closer to the Random-codebook distances indicates the separation capability of RVQ.

However, since the comparison is made with Gaussian noise, the smaller distance does not necessarily conclude that resting state fMRI and task fMRI are similar. More appropriate comparisons are needed, which is a subject for future research.

In Section 4.2.2, we evaluated BrainCodec’s reconstruction capability by illustrating reconstructed images for a specific HCP task fMRI. We demonstrated that these reconstructions make brain activity more discernible compared to the original or VAE-based methods. Here, we will show that the previous example is not an isolated success by presenting more cases using HCP data from different tasks.

For consistency, we set the brain slicing positions to align with the results published in Barch et al. [2013], which analyzes HCP task fMRI data. For example, the previous MOTOR task is depicted in Fig. 7 of this paper. In this paper, brain activity is visualized using traditional methods that average data over time and different runs, allowing for detailed identification of activity locations in the motor cortex. However, this traditional approach does not reveal the progression of brain activity starting from the visual cortex, as seen in Section 4.2.2. Therefore, when reviewing the results of the reconstructed images, please refer to both the traditional method as cited in Barch et al. [2013] and the averaged fMRI data, which, despite being a simple average, takes into account temporal variations. In the following sections, we describe the results respecting the traditional methods; however, it should be noted here that there is consistency between the results of the traditional methods and our averaged fMRI data throughout.

Below, we present reconstructed images for other tasks: Language, Social, and Relational tasks.

Fig 15 presents the reconstruction results for the Language task. The Language task is divided into "story" and "math" components, and here we show the results for the "story" component. In the "story" task, participants first listen to a short story adapted from Aesop’s fables. Following this, they answer multiple-choice questions related to the story. For more detailed information about this task, please refer to Binder et al. [2011].

Upon examining the figure, it is evident that the reconstructions by BrainVAE closely resemble the originals, failing to effectively suppress noise. In contrast, the reconstructions by BrainCodec reveal strong reactions in the lower temporal regions that are not visible in the originals or BrainVAE reconstructions. These brain activities, observable only with BrainCodec, are crucial for characterizing the "story" task, as seen in Fig 8 of Barch et al. [2013]. Additionally, using only the first half of the codebooks exhibits similar trends, suggesting that, as discussed in Section 4.2.2, this approach may be sufficient. As in Section 4.2.2, we report the L1 distance between the Mean and each fMRI data (similarly for other tasks introduced below). The order is the same: original, BrainVAE, BrainCodec (0-7), and BrainCodec (0-3), with the results as follows: $0.794\pm 0.059,0.824\pm 0.042,0.380\pm 0.047,0.362\pm 0.038$. Indeed, similar to the MOTOR task, it can be confirmed that the reconstructions by BrainCodec replicate the Mean more accurately.

Next, Fig 16 shows the results for the Social task. In this task, participants watch a video displaying several shapes moving randomly on a screen. They are then asked to determine whether the shapes appear to be engaging in social interactions ("mental") or moving randomly without such interactions ("rnd"). For this reconstruction, we used the "mental" examples, where the shapes appear to be interacting. For further details on this task, please refer to Wheatley et al. [2007].

In this example, as well, the difference between the original and BrainVAE is minimal. In fact, BrainVAE appears to weaken the signal strength, suggesting that this might be a challenging example for VAE methods. On the other hand, the reconstructions by BrainCodec clearly identify areas of strong signal. Notably, in this example, there is widespread activation throughout the brain. While this might initially seem like excessive activation, it is consistent with the results shown in Fig 9 of Barch et al. [2013]. Similarly, the same patterns are observed when using only the first half of the codebooks. Similarly, the L1 distances between the Mean and each fMRI data are as follows: $0.767\pm 0.038,0.841\pm 0.028,0.358\pm 0.037,0.343\pm 0.032$. This further confirms the performance of BrainCodec numerically.

Finally, Fig 17 shows the results for the relational task. In this task, six types of shapes and six types of textures are used as stimuli. In the relational processing condition ("relation"), two pairs of objects are presented at the top and bottom of the screen. Participants first determine the dimension (shape or texture) in which the upper pair differs, and then decide whether the lower pair differs in the same dimension. In the control condition ("match"), two objects are shown at the top, one object at the bottom, and a word ("shape" or "texture") in the center. Participants must determine whether the bottom object matches either of the top objects in the specified dimension. For this reconstruction, we used data from the "relation" condition. For a more detailed explanation, please refer to Smith et al. [2007].

In this example, BrainVAE allows for relatively clear observation of brain activity, such as the responses in the visual cortex and the lower occipital regions (t=3.60). However, in all cases, BrainCodec exhibits these activities more distinctly. Notably, brain activity near the parietal region, reported in Fig 10 of Barch et al. [2013], is not observable in the original or BrainVAE reconstructions. This activity, however, is clearly visible in the BrainCodec reconstruction (t=5.76). Similarly, the L1 distances between the Mean and each fMRI data are as follows: $0.765\pm 0.038,0.909\pm 0.025,0.353\pm 0.033,0.334\pm 0.028$. The same results were confirmed.

The above are additional concrete examples demonstrating the superiority of our BrainCodec. These examples were chosen for their clarity. The overall effectiveness of this method, which can be viewed as a denoising technique, should ideally be evaluated quantitatively—a task for future research. Our code and trained BrainCodec models are publicly available, and we encourage you to apply them to your own fMRI data to assess their denoising performance.