SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jingyuan Sun†‡
KU Leuven
jingyuan.sun@kuleuven.be
Zijiao Chen
Shaonan Wang
Chinese Academy of Science
shaonan.wang@nlpr.ia.ac.cn
Mingxiao Li†
mingxiao.li@kuleuven.be
Yunhao Zhang
zhangyunhao2021@ia.ac.cn
Marie-Francine Moens
sien.moens@kuleuven.be
Decoding visual stimuli from neural responses recorded by functional Magnetic Resonance Imaging (fMRI) presents an intriguing intersection between cognitive neuroscience and machine learning, promising advancements in understanding human visual perception. However, the task is challenging due to the noisy nature of fMRI signals and the intricate pattern of brain visual representations. To mitigate these challenges, we introduce a two-phase fMRI representation learning framework. The first phase pre-trains an fMRI feature learner with a proposed Double-contrastive Mask Auto-encoder to learn denoised representations. The second phase tunes the feature learner to attend to neural activation patterns most informative for visual reconstruction with guidance from an image auto-encoder. The optimized fMRI feature learner then conditions a latent diffusion model to reconstruct image stimuli from brain activities. Experimental results demonstrate our model's superiority in generating high-resolution and semantically accurate images, substantially exceeding previous state-of-the-art methods by  $39.34\%$  in the 50-way-top-1 semantic classification accuracy. The code implementations will be available at https://github.com/soinx0629/vis_dec_neurips/.
Reconstructing visual stimuli from neural imaging data represents a promising interdisciplinary area between cognitive neuroscience and machine learning [1]. A system capable of accurately decoding neural responses to visual input can help illuminate the complex mechanisms underlying the brain's visual perception [2, 3]. Furthermore, it can elucidate the relationships between human visual systems and computational vision models [4, 5, 6]. Such a system also holds the potential to assist patients, particularly those with motor disabilities, in expressing their thoughts and intentions through brain signals.
Functional Magnetic Resonance Imaging (fMRI), as a non-invasive method to measure brain activity, has been extensively utilized to decipher perceptions from neural responses [5, 7, 8]. However,
Corresponding Author
despite its utility, achieving reliable and robust visual reconstructions from fMRI recordings remains a significant challenge [9]. Primarily, the fMRI data is inherently noisy. The recorded signals encompass not only the specific responses elicited by visual stimuli but also incorporate additional sources of noise arising from various cognitive, physiological processes, and scanner operations [10, 11]. These noises can obscure the neural activation patterns associated with the stimulus, thereby complicating the direct decoding of visual information from the fMRI data. Moreover, the process by which a visual stimulus arouses a neural response is dynamic, intricate and multifaceted. It involves multiple stages of neural processing [12, 13], from the initial perception in the retina to higher-order cognitive operations within the brain's visual and associative regions [14, 15]. The resulting fMRI signal is a highly convolved representation of these distinct processing stages, rather than a linear one [16]. It is thus non-trivial to reverse this process and disentangle the convolved representations to achieve visual decoding.
Existing methods for vision decoding tend to struggle with such challenging complexity. Previous pioneering work has relied on traditional statistical approaches, such as ridge regression, to map fMRI signals back to the corresponding stimuli [17, 18]. However, these oversimplified methods often fall short of capturing the non-linear relationship between the stimulus and the neural responses. More recently, deep learning methods such as Generative Adversarial Networks (GAN) [9, 19] and latent diffusion models (LDMs) [7, 5] have been adopted to model the non-linearity and yield better results. But the difficulty of disentangling vision-related brain activities from noises still hinders these methods from decoding images with optimal accuracy.
To navigate these challenges, we propose a double-phase fMRI representation learning framework. In Phase 1, we pre-train the fMRI feature learner on large-scale unlabeled fMRI data with a novel Double-contrastive Masked Auto-encoder (DC-MAE). The DC-MAE helps discern common patterns of brain activities shared among populations over individual noises. In Phase 2, we further tune the feature learner on the fMRI-image parallel data with an image auto-encoder. The pixel-level guidance from the image auto-encoder teaches the fMRI auto-encoder to attend to brain signals that are most informative for image reconstruction. The trained fMRI feature learner is then used to condition an LDM to reconstruct image stimuli from brain activities. Experimental results demonstrate that the proposed model generates high-resolution and semantically accurate images, outperforming the previous state-of-the-art method by  $39.34\%$  in 50-way-top-1 accuracy. Our research paves the way for further exploration of the potential of decoding tasks.
2.1 Visual Decoding from fMRI
Driven by the substantial potential, recent years have witnessed a growing interest in reconstructing visual experiences from fMRI data. This task has been examined in various contexts, including explicitly presented visual stimuli [4, 20], perceived emotions [21] and even imagined content [1]. Though studies on this task keep emerging, the challenges presented by the low signal-to-noise ratio (SNR) of fMRI data and the high complexity of brain visual representations still exist. In the initial stage in this field, efforts to identify or reconstruct visual images from fMRI primarily utilized handcrafted features [17] and traditional regression models [18, 22]. Nonetheless, oversimplified methods can not fully account for the intricate patterns of brain visual representations, generating blurry and semantically meaningless images.
Recent research has shifted towards artificial neural network representations and deep generative models [23]. Typically, such models mapped fMRI signals to image features and fine-tuned pretrained generative models like Generative Adversarial Networks (GANs) [9, 19] or Diffusion Models [7, 24, 25, 26, 27] to generate images from the mapped features. For example, [28] decoded fMRI data to hierarchical image features extracted by a pre-trained VGG and fed the predicted features to a GAN. But [28], like other parallel work [29, 1], used fMRI directly for training and decoding. Without explicitly denoising the fMRI representations, though these works outperformed traditional regression-based models, they still fell short by generating implausible images. In contrast, our framework contains an individual phase to learn denoised fMRI representations with DC-MAE. [7] adopted a naive MAE to learn fMRI representations. But compared with [7], our framework further introduces pixel-level guidance from the image auto-encoder to help disentangle the vision-related neural activities from background noises. Our model is also not similar to other works [30, 31] which
directly map fMRI to the image feature space and inevitably face data incompatibility between the two very different modalities. We encourage the emergence of a cross-modality representation space by guiding the fMRI auto-encoder with a pre-trained image auto-encoder. Experimental results prove the superiority of our methods over these related baselines.
2.2 Diffusion Probabilistic Models
Diffusion probabilistic models have been empirically established as powerful generative models for image generation surpassing GANs [32] in terms of both diversity and fidelity. The diffusion model was first proposed in [33] and further improved by [26, 34]. The quality of generated image could also be enhanced by training-free methods [35, 36]. The initial diffusion models were primarily applied in the pixel space. They achieve impressive success in generating images of high quality, but also suffer from significant drawbacks such as prolonged inference time and substantial training costs [26]. To address these two issues, [27] introduce the latent diffusion model (LDM), also referred to as stable diffusion (SD). This approach adopts a pretrained Vector Quantized Generative Adversarial Network (VQGAN) [37] or Variational auto-encoder (VAE) [38] to construct a latent image space, within which optimization and evaluation are performed. The LDM not only generates images of high quality, but also alleviates the computational burden. Moreover, the incorporation of a cross-attention mechanism within the attention block of the diffusion UNet model permits the LDM model to offer a broad spectrum of controls in image synthesis. This includes textual controls [39, 40, 41], controls over images in various domains [42, 43] such as depth maps, sketch maps, or candy maps. Such versatility and adaptability give the LDM model a substantial advancement in the field of image synthesis.
In this subsection, we provide a concise analysis of the fMRI data and brain visual representations that motivate the design of our method.
First, the fMRI recordings are inherently noisy, subject to various sources of physiological and scanner-related noises [44]. FMRI records not only responses to visual stimuli but also signals from other cognitive processes. Second, fMRI quantifies changes in the blood-oxygen-level-dependent (BOLD) signal. Adjacent voxels are often found to display similar magnitudes, suggesting the spatial redundancy of fMRI data [2]. Third, neural responses to identical stimuli can exhibit significant divergence [45, 46] across populations.
Considering these analyses altogether, we propose a double-phase fMRI representation learning framework. In Phase 1, we pre-train an MAE with a contrastive loss to learn fMRI representations from unlabeled data. The masking which sets a certain portion of the input data to zero targets the spatial redundancy of fMRI data. The calculation of recovering the original data from the remaining after masking suppresses noises. Optimization of the contrastive loss discerns common patterns of brain activities over individual variances. After pre-training in Phase 1, we tune the fMRI auto-encoder with an image auto-encoder. We expect the pixel-level guidance from the image auto-encoder to support the fMRI auto-encoder in disentangling and attending to brain signals related to vision processing. After FRL Phase 1 and Phase 2, we apply the representation learned by the fMRI auto-encoder as conditions to tune the LDM and generate the image from the brain activities.
3.2 fMRI Representation Learning (FRL)
Phase 1: Pre-training Double-Contrastive Masked Auto-Encoder (DC-MAE) We introduce a method termed the "Double-Contrastive Masked Auto-Encoder". DC-MAE has been specifically designed to pre-train fMRI representations from unlabeled data, as inspired by previous work in visual contrastive learning [47]. As shown in Figure 1, the DC-MAE consists of an encoder  $E_{F}$  and a decoder  $D_{F}$ .  $E_{F}$  takes a masked version of the fMRI signal as the input, and  $D_{F}$  is trained to predict the unmasked fMRI. The term "Double-Contrastive" implies that the model optimizes contrastive losses and engages in two separate contrasting processes during the representation learning of an fMRI example.
Figure 1: The Phase 1 and Phase 2 of the fMRI Representation Learning framework. After the fMRI feature learner is pre-trained in Phase 1, it will be tuned with an image auto-encoder in Phase 2. The definitions of depicted variables in this figure are detailed in Section 3.1.
Specifically, in the first contrasting, each sample  $v_{i}(i \leq n)^{1}$  within a batch of  $n$  fMRI examples  $v$  undergoes two separate instances of random masking. This procedure yields two distinct masked versions of  $v_{i}$ , referred to as  $v_{i}^{m_{1}}$  and  $v_{i}^{m_{2}}$ . These masked versions establish a positive sample pair for the first contrasting. Subsequently, a 1D convolutional layer, characterized by a stride that matches the patch size, tokenizes  $v_{i}^{m_{1}}$  and  $v_{i}^{m_{2}}$  into embeddings. These embeddings are then fed independently into the same fMRI encoder  $E_{F}$ . The decoder  $D_{F}$ , in turn, receives each of the encoded latent representations as input and generates predictions  $v_{i}^{dm_{1}}$  and  $v_{i}^{dm_{2}}$ . The first contrastive loss, denoted as the cross-contrastive loss, is thus calculated through an InfoNCE loss [48] as follows:
$$
\mathcal {L} _ {C} = - \log \frac {\exp \left(v _ {i} ^ {d m _ {1}} \cdot v _ {i} ^ {d m _ {2}} / \tau\right)}{\exp \left(v _ {i} ^ {d m _ {1}} \cdot v _ {i} ^ {d m _ {2}} / \tau\right) + \sum_ {j \neq i} \exp \left(v _ {i} ^ {d m _ {1}} \cdot v _ {j} ^ {d m _ {1}} / \tau\right)} \tag {1}
$$
In the second contrasting, an unmasked original image  $v_{i}(i \leq n)$  and its corresponding masked image  $v_{i}^{m}$  form an inherent positive sample pair. Here,  $v_{i}^{dm}$  denotes the predicted image output from the decoder  $D_{F}$ . The second contrastive loss, referred to as the self-contrastive loss, is computed through as:
$$
\mathcal {L} _ {S} = - \log \frac {\exp \left(v _ {i} ^ {d m} \cdot v _ {i} / \tau\right)}{\exp \left(v _ {i} ^ {d m} \cdot v _ {i} / \tau\right) + \sum_ {j \neq i} \exp \left(v _ {i} ^ {d m} \cdot v _ {j} ^ {d m} / \tau\right)}. \tag {2}
$$
Optimizing the self-contrastive loss  $\mathcal{L}_S$  can also achieve mask-reconstruction. For both  $\mathcal{L}_C$  and  $\mathcal{L}_S$ , the negative examples  $v_{j}(j\neq i)$  are sourced from the same batch of instances.  $\mathcal{L}_C$  and  $\mathcal{L}_S$  are optimized jointly as follows:
$$
\mathcal {L} = \gamma_ {C} \mathcal {L} _ {C} + \gamma_ {S} \mathcal {L} _ {S} \tag {3}
$$
In this equation, the hyper-parameters  $\gamma_{C}$  and  $\gamma_{S}$  regulate the weight of each loss.
<sup>1</sup>Note that this is a 1-dimensional vector, not a time series, as we have averaged the data across the time dimension. This results in a spatial pattern of fMRI signal over the visual cortex for each picture viewed by the subject. We then employ a 1D convolutional model to transform this 1D spatial pattern of the fMRI signal,  $v_{i}$ , into an embedding.
Phase 2: Tuning with Cross Modality Guidance Considering the relatively low signal-to-noise ratio (SNR) and highly convolved nature of fMRI recordings, it is important for the fMRI feature learner to attend to the brain activation patterns most relevant for visual processing and most informative for reconstruction.
So as shown in Figure 1, after pre-training in Phase 1, the fMRI auto-encoder is tuned to reconstruct fMRI with the aid of an MAE for image, and vice versa in Phase 2. Specifically, denote one image from a batch of  $n$  samples as  $u_{i}(i\leq n)$  and the fMRI recorded neural responses to  $u_{i}$  as  $v_{i}$ .  $u_{i}$  and  $v_{i}$  go through patching and random masking to be  $u_{i}^{m}$  and  $v_{i}^{m}$ .  $u_{i}^{m}$  and  $v_{i}^{m}$  are respectively fed in the image encoder  $E_{I}$  and fMRI encoder  $E_{F}$ , getting  $h_{E_I}^{u_i} = E_I(u_i^m)$  and  $h_{E_F}^{v_i} = E_F(v_i^m)$ . For reconstructing the fMRI  $v_{i}$ ,  $h_{E_I}^{u_i}$  and  $h_{E_F}^{v_i}$  are merged with a cross-attention module as follows:
$$
\begin{array}{l} Q _ {I} ^ {u} = W ^ {Q _ {I}} h _ {E _ {I}} ^ {u _ {i}} + b ^ {Q _ {I}}; K _ {F} ^ {v} = W ^ {K _ {F}} h _ {E _ {F}} ^ {v _ {i}} + b ^ {K _ {F}}; V _ {F} ^ {v _ {i}} = W ^ {V _ {F}} h _ {E _ {F}} ^ {v _ {i}} + b ^ {V _ {F}} \\ C A _ {F} \left(Q _ {I} ^ {u _ {i}}, K _ {F} ^ {v _ {i}}, V _ {F} ^ {v _ {i}}\right) = \operatorname {s o f t m a x} \left(\frac {Q _ {I} ^ {u _ {i}} \left(K _ {F} ^ {v _ {i}}\right) ^ {T}}{\sqrt {d _ {k}}}\right) V _ {F} ^ {v _ {i}} \tag {4} \\ \end{array}
$$
$W$  and  $b$  denote the weights and biases of corresponding linear layers.  $\sqrt{d_k}$  is the scaling factor and  $d_{k}$  is the dimension of the key vectors.  $CA$  is the abbreviation of cross-attention.  $CA_{F}(Q_{I}^{u_{i}},K_{F}^{v_{i}},V_{F}^{v_{i}})$  is then added to  $h_{E_F}^{v_i}$  and fed into the fMRI decoder to reconstruct  $v_{i}$  as  $v_{i}^{d}$ .
$$
v _ {i} ^ {d} = D _ {F} \left(h _ {E _ {F}} ^ {v _ {i}} + C A _ {F} \left(Q _ {I} ^ {u _ {i}}, K _ {F} ^ {v _ {i}}, V _ {F} ^ {v _ {i}}\right)\right) \tag {5}
$$
A similar computation is also conducted in the image auto-encoder. The output  $h_{E_I}^{u_i}$  of the image encoder  $E_I$  is merged with the output of  $E_F$  through a cross-attention module  $CA_I$ , and then used to decode an image  $u_i$  as  $u_i^d$ :
$$
u _ {i} ^ {d} = D _ {I} \left(h _ {E _ {I}} ^ {u _ {i}} + C A _ {I} \left(Q _ {F} ^ {v}, K _ {I} ^ {u _ {i}}, V _ {I} ^ {u _ {i}}\right)\right) \tag {6}
$$
The fMRI and image auto-encoders are trained jointly by optimizing the following loss:
$$
L = \gamma_ {F} \left(v _ {i} - v _ {i} ^ {d}\right) ^ {2} + \gamma_ {I} \left(u _ {i} - u _ {i} ^ {d}\right) ^ {2} \tag {7}
$$
3.3 Image Generation with Latent Diffusion Model (LDM)
After the fMRI feature learner is trained through FRL Phase 1 and Phase 2, we use its encoder  $E_{F}$  to condition a LDM to generate images from brain activities. As displayed in Figure 2, the diffusion model consists of a forward diffusion process and a reverse denoising process. The forward process gradually degrades an image to a normal Gaussian noise by incrementally introducing Gaussian noise of variable variance. This process can be mathematically formulated as  $q(x_{t}|x_{t - 1}) = \mathcal{N}(x_{t},\sqrt{\alpha} x_{t - 1},(1 - \alpha_{t})I)$ , where  $t$  denotes the temporal step and  $\alpha$  encompasses predefined noise schedule parameters. During the training phase, the diffusion model, specifically the U-Net [49] model, is optimized with the loss function below to learn the noise  $\epsilon_t$  added at each time step in the forward process.
$$
\mathcal {L} _ {t} ^ {\text {s i m p l e}} = E _ {t, x _ {0}, \epsilon_ {t} \sim \mathcal {N} (0, 1)} [ \| \epsilon_ {t} - \epsilon_ {\theta} (z _ {t}, t) \| _ {2} ^ {2} ] \tag {8}
$$
In the backward process, an image is synthesized by progressively eliminating noise from a randomly initialized standard Gaussian noise. The Latent Diffusion Model (LDM) [27] performs both forward and backward processes in a low-dimensional latent space which is constructed using a pre-trained VQ-VAE model [38]. This approach significantly mitigates the computational complexity while simultaneously maintaining the generated image's quality. We leverage the LDM as the backbone of our image generation model.
Decoding fMRI to a natural image can be seen as conditional image generation task. Given the relatively low SNR inherent in fMRI, combined with the limited quantity of fMRI-to-image data pairs, training an fMRI-to-image generation model from scratch presents substantial challenges. Consequently, the objective of this phase is to harness the fMRI to extract image-related knowledge from a pre-existing conditional image generation model.
Our approach extracts visual knowledge from the pretrained label-to-image LDM to generate image with the conditioning of fMRI. We incorporate the fMRI information into the LDM via a cross-attention operation similar to equation (4), as proposed in [27]. To further enhance the guidance provided by the conditional information, following the methodology of previous research [7], we
Figure 2: [a] Demo of the forward and backward processes of the diffusion model. [b] The forward process of the diffusion model which progressively corrupts an image with Gaussian noise. [c] In the backward process, the diffusion model, conditioned on our pretrained fMRI encoder, gradually denoises white noises to generate the image.
implement both cross-attention conditioning and time steps conditioning [50]. During training, given an image  $u$  and fMRI  $v$ , along with VQGAN encoder  $E_{G}$  and the fMRI encoder  $E_{F}$  (which is trained in FRL Phase 1 and 2), we freeze the LDM and finetune the fMRI encoder using the following loss:
$$
\mathcal {L} _ {t} ^ {\text {s i m p l e}} = E _ {t, u _ {0}, \epsilon_ {t} \sim \mathcal {N} (0, 1)} [ \| \epsilon_ {t} - \epsilon_ {\theta} (\phi (E _ {G} (u)) _ {t}, t, E _ {F} (v)) \| _ {2} ^ {2} ] \tag {9}
$$
Here,  $\phi(E_g(u))_t = \sqrt{\overline{\alpha}_t} E_g(u) + \sqrt{1 - \overline{\alpha}_t} \epsilon$  and  $\overline{\alpha}_t$  is the noise schedule of diffusion model. During the inference phase, the process begins with a standard Gaussian noise at time step  $T$ . The LDM follows the backward process sequentially to progressively denoise the hidden representation, conditioning on the given fMRI information. Upon reaching time step zero, the VQGAN decoder  $D_G$  is used to convert the hidden representations to an image. The forward and backward process of the LDM are detailed in Figure 2[b,c].
4.1 fMRI Datasets
HCP The Human Connectome Project (HCP) originally serves as an extensive exploration into the connectivity of the human brain. It offers an open-sourced database of neuroimaging and behavioral data collected from 1,200 healthy young adults within the age range of 22-35 years. Currently, it stands as the largest public resource of MRI data pertaining to the human brain, providing an excellent foundation for the pre-training of brain activation pattern representations. Of the subjects involved, 1113 underwent scanning via a Siemens Skyra Connectom scanner for 3T MR, while a Siemens Magnetom scanner for 7T MR was utilized for the remaining 184. For the scope of this paper, we will predominantly focus on the data derived from the more populated 3T dataset.
GOD The Generic Object Decoding (GOD) Dataset is a specialized resource developed for fMRI-based decoding. It aggregates fMRI data gathered through the presentation of images from 200 representative object categories, originating from the 2011 fall release of ImageNet. The training session incorporated 1,200 images (8 per category from 150 distinct object categories). In contrast, the test session included 50 images (one from each of the 50 object categories). It is noteworthy that the categories in the test session were unique from those in the training session and were introduced in a randomized sequence across runs. On five subjects the fMRI scanning was conducted.
BOLD5000 The BOLD5000 dataset is a result of an extensive slow event-related human brain fMRI study. It comprises 5,254 images, with 4,916 of them being unique. This makes it one of the most comprehensive publicly available datasets in the field. The dataset's principal advantage is its high
diversity, enabling the capture of the complexity and variability inherent in natural visual stimuli. The images in BOLD5000 were selected from three popular computer vision datasets: ImageNet, COCO, and Scenes. ImageNet provided 1,916 images primarily focusing on singular objects. Meanwhile, COCO contributed 2000 images featuring multiple objects, and Scenes contributed 1000 images depicting hand-crafted indoor and outdoor scenes. Four participants labeled CSI1 through CSI4, were involved in this study and underwent scanning via a 3T Siemens Verio MR scanner equipped with a 32-channel phased array head coil.
4.2.1 FMRI Representation Learning (FRL)
For both FRL Phase 1 and Phase 2, the fMRI auto-encoder is the same ViT-based masked auto-encoder (MAE). We employed an asymmetric architecture for the fMRI auto-encoder, in which the decoder is considerably smaller with 8 layers than the encoder with 24 layers. We used a larger embedding-to-patch size ratio, specifically a patch size of 16 and an embedding dimension of 1024 for our model. We used random sparsification (RS) as a form of data augmentation, randomly selecting and setting  $20\%$  of voxels in each fMRI to zero.