SlowGuess/NeurIPS_2023_OCR huggingface dataset:
FRL Phase 1 In Phase 1, we train the masked ViT-based fMRI auto-encoder with contrastive loss. For GOD subject 1,4,5 and BOLD5000 CSI 1,2, self-contrastive  $(\gamma_{s})$  and cross-contrastive  $(\gamma_{c})$  loss weights are both 1. The masking ratio is 0.5. For GOD subject 2,3 and BOLD5000 CSI 3,4,  $\gamma_{s} = 1$  and  $\gamma_{c} = 0.5$ , masking ratio is 0.75. We set the batch size to 250 and train for 140 epochs on one NVIDIA A100 GPU. We train with 20-epoch warming up and an initial learning rate of  $2.5\mathrm{e - }4$ . We optimize with AdamW and weight decay 0.05.
FRL Phase 2 In Phase 2, we tune the fMRI autoencoder jointly with an image auto-encoder, which is a pre-trained ViT-based MAE released by [51]. The image auto-encoder has a 12-layer encoder with a 768 hidden size and a 6-layer decoder with a 512 hidden size. We set the batch size to be 16 and train for 60 epochs. We train with 2-epoch warming up. The initial learning rate is 5.3e-5. We optimize with AdamW and weight decay 0.05. We freeze the parameters of the decoder of the image-autoencoder and only tune the encoder. The two phases in total take about 12 hours to run. After the two phases, we only save the checkpoint of the fMRI encoder which has 15.16M parameters
4.2.2 Fine-tuning LDM
In this stage, we jointly optimize the parameters of LDM cross-attention heads and the fMRI encoder, while keeping other parameters of LDM unchanged. Given an fMRI-image pair, we first use the pre-trained VQ encoder to encode the image to obtain the latent representation which is further used as an objective to guide the joint training of the fMRI encoder and LDM cross-attention heads. During training, the fMRI data passes through the fMRI encoder trained using FRL, producing a patchified representation. This representation is then projected into key and value representation of cross-attention modules in the UNet of LDM. Furthermore, it is added to the time embedding to conduct double conditioning. The training follows the regular training pipeline of the diffusion model, where the model is optimized to learn to predict the Gaussian noise added to the image latent representation at each time step with the guidance of the given conditioning information. Here, we use the output of the fMRI encoder as the conditioning information. We conduct training with the following parameters: the batch size of 5, diffusion steps of 1000, the AdamW optimizer, a learning rate of  $5.5e - 5$ , and an image resolution of  $256 \times 256 \times 3$ .
4.3 Baseline Models and Evaluation Metric
Baseline Models We juxtapose our proposed model with recently published benchmarks: IC-GAN [52], Self-supervised auto-encoder (SS-AE) [53], and DC-LDM [7]. The IC-GAN is based on an Instance-Conditioned GAN, whereas the SS-AE utilizes cycle consistency and perceptual losses to reconstruct images from fMRI brain recordings. The DC-LDM employs a double-conditioned LDM, demonstrating superior performance on the GOD and BOLD5000 datasets. These benchmarks reflect prevalent methodologies in visual reconstruction.
Evaluation Metric Visual decoding prioritizes semantic consistency. Therefore, our evaluation metric is the  $n$ -way top- $k$  accuracy, which aligns with the precedent set in the literature [53, 7]. We
employ the pre-trained ImageNet-1K classifier [54] as a semantic correctness evaluator. During the evaluation, both the generated image and the corresponding ground truth image are fed into the classifier. Semantic correctness is then determined based on whether the top- $k$  classification among  $n$  randomly selected classes aligns with the ground-truth classification. Further details can be referenced in the appendix.
5.1 Reconstruction Results
In this section, we present a comparative analysis of our results with preceding studies, which include DC-LDM, IC-GAN, and SS-AE. We have evaluated our model on both GOD and BOLD5000 datasets. Following the setting of previous work [7], we first display in detail the results in Figure 3 for GOD subject 3 and BOLD5000 subject CSI 1 which achieve the best performance in their respective dataset. It is noteworthy that the original DC-LDM implementation utilized test set fMRI data for tuning, a setting not adopted by IC-GAN and SS-AE. To maintain a fair comparison, we first prohibit the use of test set fMRI data by the DC-LDM. As shown in Figure 3[a], our method surpasses the previous models by a large margin. For instance, our model achieves an accuracy exceeding that of DC-LDM and IC-GAN by over  $39.34\%$  (calculated by  $(25.080 - 17.999) / 25.080 \times 100\% \approx 39.34\%$ , comparing our model's accuracy 25.080 against the DC-ldM's accuracy 17.999) and  $66.7\%$  respectively. The improvement of our model over the baselines is significant. All the significant results have p-value  $< 0.01$  with paired t-tests.
Figure 3: Reconstruction results. [a] Top-1 (a1) and top-5 (a2) classification accuracy of our model and other baselines on GOD subject 3. [b] Samples of reconstructed images and their ground truth from GOD subject 3's data. [c] Samples of reconstructed images and their ground truth from the BOLD5000 CSI 1's data.
To further investigate the quality of the images produced by different models, we randomly select 5 samples from the GOD test set and exhibit the generated images in Figure 3[b]. The SS-AE model can only generate a general layout, while the remaining three models are cable of producing images with a semantic meaning similar to that of ground truth image. Our model generates images of superior quality that exhibits a higher degree of semantic consistency with the ground truth images. In Figure 3[c], we compare images generated by our model and DC-LDM. (IC-GAN and SS-AE were not trained on BOLD5000 in their original paper). Our model achieves 50-way-top-1 accuracy of 25 on CSI1. Though both models can produce high-quality images, the image generated by our model bears a semantic meaning more consistent with the ground truth images.
To check if our model may reliably reconstruct brain activities on different subjects, we further evaluate it against DC-LDM on all the other 4 subjects of the GOD dataset. The bar charts in Figure 4[a] below show the results in 50-way-top-1 classification accuracy. Our model substantially outperforms the previous state-of-the-art method (DC-LDM) on all GOD subjects. To achieve DC-LDM's reported performance in its original paper [7], this method need signals from test set fMRI
data. To ensure a fair evaluation, we banned DC-LDM from tuning on the test set in the comparison of the Figure 3. But we show here that, our model still largely exceeds DC-LDM on four GOD subjects even after DC-LDM is tuned on the test set fMRI data.
Figure 4: Reconstruction performance of our model and other baselines on GOD subjects 1, 2, 4 and 5, measured by 50-way-top-1 classification accuracy
5.2 Ablation Study Results
We train the FRL Phase with different settings of hyper-parameters. The trained fMRI encoders then guide the LDM in generating images. We use the 50-way-top-1 accuracy of LDM image generation to measure the influence of hyper-parameter values. Due to space limit, in this section we mainly focus on FRL Phase 2 where direct interaction between the image and fMRI data takes place which may more highly influence the final reconstruction performance.
Ablation parameter	ID	fMRI rec. loss weight	Image rec. loss weight	fMRI mask ratio	fMRI mask ratio	Dec. layer num	Accuracy
fMRI and image reconstruction loss weight	1	0	1	0.75	0.75	6	12.52
2	1	0	0.75	0.75	6	16.64
3	0.5	1.5	0.75	0.75	6	17.68
4	1	1	0.75	0.75	6	19.64
5	1.5	0.5	0.75	0.75	6	19.76
6	0.25	1.5	0.75	0.75	6	20.8
Ablation parameter	ID	fMRI rec. loss weight	Image rec. loss weight	fMRI mask ratio	Image mask ratio	Dec. layer num	Accuracy
fMRI and image mask ratio	7	1	1	0.5	0.75	6	15.52
8	1	1	0.5	0.5	6	17.72
9	1	1	0.5	0.25	6	19
4	1	1	0.75	0.75	6	19.64
10	1	1	0.75	0.5	6	21.04
Ablation parameter	ID	fMRI rec. loss weight	Image rec. loss weight	fMRI mask ratio	Image mask ratio	Dec. layer num	Accuracy
number of fMRI and image decoder layer	11	1	1	0.75	0.5	2	17.44
12	1	1	0.75	0.5	4	17.76
4	1	1	0.75	0.5	6	19.64
13	1	1	0.75	0.5	8	17.44
Best parameter	14	0.25	1.5	0.75	0.5	6	25.08
Table 2: Ablation study of FRL Phase 2 on GOD subject 3. Cells with colored shades denote the hyper-parameters to be tuned in one ablation group. For example, cells with yellow shades denote that fMRI and image mask ratio are the parameters to be tuned while other parameters are kept the same.
Reconstruction Loss Weight In the FRL Phase 2, the feature learner is tuned by optimizing the joint loss of fMRI and image reconstruction, as in equation (7). So we first focus on the influence of the fMRI and image reconstruction loss weights, namely  $\gamma_{F}$  and  $\gamma_{i}$ . The results are reported in Table 2 (ID1-6). In experiment ID 1 and 2, we set one of the weights to 0 to evaluate the necessity of using the corresponding loss. We find that jointly optimizing the fMRI and image reconstruction loss is necessary to achieve optimal task performance. We adjust the  $\gamma_{F}$  and  $\gamma_{i}$  and find that setting in ID 6 where  $\gamma_{F} = 0.25$  and  $\gamma_{i} = 1.5$  yields the highest performance.
Mask Ratio In the FRL Phase 2, both the fMRI and image auto-encoders are based on MAE. The mask ratio of the input is an important setting for these models. We demonstrate the influence of the mask ratio setting in Table 2 (ID 4 and ID 7-10). We find that applying a higher mask ratio on fMRI data and a lower mask ratio on image data generally leads to better performance. LDM conditioned by the fMRI encoder which is trained with an fMRI mask ratio of 0.75 and image mask ratio of 0.5 leads to the highest reconstruction accuracy.
Decoder Layers Following [7, 51, 55], for both the fMRI and the image auto-encoder, we build asymmetric architectures where the decoder is much smaller than the encoder. In Table 2 (ID 11-13), we report the results of tuning decoder depth. We apply the same depth for the two decoders. We find that a moderate decoder depth of 6 produces optimal results. Neither a too shallow nor a too deep decoder improves reconstruction performance on GOD subject 3.
The experimental findings indicate that, with the proposed fMRI representation learning framework and a pre-trained LDM, we can achieve a degree of visual reconstruction of human brain activities, substantially outperforming baselines. Nonetheless, our analysis also uncovers some limitations within our model. Primarily, our model appears to be slightly affected by a categorical bias issue. We hypothesize that this may stem from the inherent bias present in the dataset used to train the LDM. Additionally, while our model demonstrates proficiency in capturing high-level semantics, it sometimes fails to reconstruct specific details of an image. A plausible explanation could be the concurrent imagination of multiple objects by the participants during data collection, which inevitably results in a noisy fMRI feature. Contrasting with general image generation, which is typically characterized by a focus on diversity, visual decoding underscores the importance of consistency, thereby necessitating a reduction in bias during the generation process. As such, the exploration of techniques to minimize the influence of data bias, as well as methods to enhance the reconstruction of image details when generating images from fMRI data, would hold significant academic value and interest.
In this work, we propose a double-phase fMRI representation learning framework (FRL) with an LDM to reconstruct visual experiences from brain activities. The FRL denoises fMRI features by contrastive masked modeling in Phase 1. And in Phase 2 it learns to disentangle and attend to brain activation patterns most informative for visual decoding with the guidance of an image MAE. With the conditioning of the optimized fMRI feature learner, we show that an LDM generates images of high quality, largely outperforming the previous state-of-the-art. Extensive ablation studies further verify the effectiveness of each component that we propose.
8 Ethical Statements
We used preprocessed data from publicly available datasets. The fMRI data that we train with have been processed and do not contain any data that can be directly linked to the participants' identities. The collection procedure of the fMRI undergoes strict ethical review as stated in their original paper.
This work is funded by the CALCULUS project (European Research Council Advanced Grant H2020-ERC-2017-ADG 788506) and the Flanders AI Impuls Programme - FLAIR.
A.1 Reconstruction Performance
We provide the performance of our model on BOLD5000 subjects 1, 2, 3, and 4 in Table A.1. Following previous work [7], all results are presented in 50-way-top-1 classification accuracy.
BOLD 5000	CSI 1	CSI 2	CSI 3	CSI 4
OURS	25	18.69	16.14	18.98
Table A.1 Reconstruction performance of our model on BOLD5000 subject CSI 1-4, measured by 50-way-top-1 classification accuracy.
A.2 Examples of Reconstructed Images
Figures A.2 and A.3 present images generated by our model using fMRI data from GOD and BOLD5000 datasets, respectively. We generated all images at a resolution of  $256 \times 256 \times 3$  using 250 PLMS steps. More samples can be generated using our code base in the supplementary materials. The code will be open-sourced with the camera ready version of this paper.
A.2.1 Reconstructed Images from GOD Dataset
GOD Subject 1
GOD Subject 2
GOD Subject 4
GOD Subject 5
Figure A.2: Randomly selected reconstructed images from GOD subject 1, 2, 4 and 5. For each subject, the upper line shows the ground truth images while the lower line shows the reconstructed images by our method.
A.2.2 Reconstructed Images from BOLD5000 Dataset
BOLD5000 CSI 1
Figure A.3: Randomly selected reconstructed images from BOLD5000 CSI 1-4. For each subject, the upper line shows the ground truth images while the lower line shows the reconstructed images by our method.
A.3 Evaluation Metrics
We use the common N-trial, n-way top-1 semantic classification as the main evaluation metrics. This evaluation method is summarized in the algorithm below:
Algorithm 1 Iterative Reasoning Module
Input: pre-trained image classifier  $F$  , generated image  $\hat{x}$  , corresponding ground truth (GT) image  $x_{gt}$    
Output: success rate  $sr\in [0,1]$    
for trail  $= 1$  to  $N$  do  $y_{gt} = F(x_{gt})$  get the prediction of GT image pred  $= F(\hat{x})$  get the output probabilities of generated image  $p = \{p_g,p_{y_1},\dots,p_{y_{n - 1}}\}$  generate probabilities set contains  $n - 1$  randomly selected from pred and  $y_{gt}$  Success if arg min  $= y_{gt}$    
end for   
return sr = number of success / N