SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Paul S. Scotti\*,1,2, Atmadeep Banerjee\*,2, Jimmie Goode†,2, Stepan Shabalin², Alex Nguyen¹, Ethan Cohen³, Aidan J. Dempster⁴, Nathalie Verlinde¹, Elad Yundler⁵, David Weisberg¹,², Kenneth A. Norman‡,¹, and Tanishq Mathew Abraham‡,2,6,7
$^{1}$ Princeton Neuroscience Institute  
 $^{2}$ Medical AI Research Center (MedARC)  
 $^{3}$ Ecole Normale Supérieure, PSL University  
 $^{4}$ University of Toronto  
 $^{5}$ Hebrew University of Jerusalem  
 $^{6}$ EleutherAI  
 $^{7}$ Stability AI  
Project Page: https://medarc.ai/mindeye/
We present MindEye, a novel fMRI-to-image approach to retrieve and reconstruct viewed images from brain activity. Our model comprises two parallel submodules that are specialized for retrieval (using contrastive learning) and reconstruction (using a diffusion prior). MindEye can map fMRI brain activity to any high dimensional multimodal latent space, like CLIP image space, enabling image reconstruction using generative models that accept embeddings from this latent space. We comprehensively compare our approach with other existing methods, using both qualitative side-by-side comparisons and quantitative evaluations, and show that MindEye achieves state-of-the-art performance in both reconstruction and retrieval tasks. In particular, MindEye can retrieve the exact original image even among highly similar candidates, indicating that its brain embeddings retain fine-grained image-specific information. This allows us to accurately retrieve images even from large-scale databases like LAION-5B. We demonstrate through ablations that MindEye's performance improvements over previous methods result from specialized submodules for retrieval and reconstruction, improved training techniques, and training models with orders of magnitude more parameters. Furthermore, we show that MindEye can better preserve low-level image features in the reconstructions by using img2img, with outputs from a separate autoencoder. All code is available on GitHub.
The problem of decoding environmental inputs and cognitive states from brain activity is fundamental to the field of neuroscience, where improved computational approaches allow for further understanding of brain mechanisms [1]. A neuroimaging methodology that has seen significant success in this domain is functional magnetic resonance imaging (fMRI), where neural activity is measured by detecting changes in blood oxygenation. fMRI decoding is already being used in
†Core contribution.
$\ddagger$  Joint senior authors.
Figure 1: Example images reconstructed from human brain activity corresponding to passive viewing of natural scenes. Reconstructed outputs from Versatile Diffusion [6] given CLIP fMRI embeddings generated by MindEye for Subject 1. See Figure 4 and Appendix A.4 for more samples.
real-time clinical domains [2] and has potential for novel mind reading applications in brain-computer interfaces. Previous works mapped fMRI activity to the embeddings of image generation models via relatively simple mappings, usually ridge regression [3-5]. Here we propose MindEye, a novel approach that involves mapping via large-scale multilayer perceptrons (MLPs), contrastive learning, and diffusion models to achieve state-of-the-art image reconstruction. See Figure 1 for select samples of reconstructions. $^{1}$
MindEye learns to map flattened spatial patterns of fMRI activity across voxels (3-dimensional cubes of cortical tissue) to the image embedding latent space of a pretrained CLIP [7] model. MindEye has an MLP backbone and 2 specialized submodules for retrieval and reconstruction. The retrieval submodule is contrastively trained and produces "disjointed CLIP fMRI" embeddings that have high cosine similarity with the corresponding image embeddings but differ in magnitude. To reconstruct images, we train a diffusion prior [8] from scratch to take in the outputs from the MLP backbone and produce aligned embeddings suitable as inputs to any pretrained image generation model that accepts CLIP image embeddings. In order to ensure that our reconstructions also match the original images' low-level features (e.g., color, texture, spatial position), we train a separate encoder that directly maps voxels to the embedding space of Stable Diffusion's [9] variational autoencoder (VAE), obtaining blurry image reconstructions that lack high-level semantic content but perform state-of-the-art on low-level image metrics. Combining the high-level "semantic" pipeline with the low-level "perceptual" pipeline in an img2img [10] setting allows MindEye to output state-of-the-art reconstructions across both low- and high-level image metrics.
In addition to image reconstruction metrics, our disjointed CLIP fMRI embeddings attain state-of-the-art performance on image retrieval and brain retrieval metrics. Image retrieval refers to finding the original seen image out of a pool of other images given a brain sample, while brain retrieval refers to finding the brain sample given an image. MindEye finds exact (top-1) matches in the pool of NSD test samples with  $>90\%$  accuracy for both image and brain retrieval, outperforming previous state-of-the-art [11, 4] which showed  $<50\%$  retrieval accuracies. These results demonstrate that MindEye brain embeddings possess fine-grained exemplar-level signal.
Our main findings are: (1) Specialized submodules for retrieval (using contrastive learning) and reconstruction (using a diffusion prior) enable a single model to achieve state-of-the-art results across both tasks even though the objectives exhibit a tradeoff. (2) Mapping to a deep MLP with a parameter count orders of magnitude higher than previous methods does not produce overfitting and instead directly benefits model performance. (3) A novel bidirectional version of mixup contrastive data augmentation further improves model performance in this low-sample setting. (4) State-of-the-art reconstructions for low-level image metrics can be obtained by independently mapping to Stable Diffusion's VAE latent space. (5) fMRI-to-image retrieval can find the exact original image even among highly similar candidates, suggesting that fine-grained image-specific information is contained in brain embeddings, thus allowing retrieval to be scaled up to large-scale databases like LAION-5B to output images without generative models.
<sup>1</sup> Images containing each subject's 982 test set reconstructions and retrievals are available on GitHub.
Figure 2: MindEye overall schematic. A high-level "semantic" pipeline maps voxels to CLIP embeddings for image reconstruction (outputs from a diffusion prior feed through generative models like Versatile Diffusion) or retrieval tasks (such as K-nearest neighbor querying of brain embeddings to the CLIP embeddings of LAION-5B images). A low-level "perceptual" pipeline maps voxels to the variational autoencoder used by Stable Diffusion to obtain blurry reconstructions, which are used as the initialization for subsequent diffusion-based image generation. The contrastive loss for the low-level pipeline is omitted for simplicity; see Appendix A.3.2 for details.
2 MindEye
MindEye consists of two pipelines (see Figure 2), a high-level (semantic) pipeline where fMRI voxels are mapped to the CLIP ViT-L/14 image space and a low-level (perceptual) pipeline where the voxels are mapped to the image embedding space of a VAE. Both pipelines follow a common structure: a residual MLP backbone followed by two task-specific submodules. For the high-level pipeline the submodules are an MLP projector and a diffusion prior. For the low-level pipeline the submodules are an MLP projector and a CNN decoder that performs  $4 \times$  upsampling. For both pipelines we observe that training the projector submodule with a contrastive loss and the second submodule with mean squared error (MSE) loss gives the best performance.
2.1 High-Level (Semantic) Pipeline
The high-level pipeline is the core of MindEye as it maps voxels to CLIP image space to be fed through pretrained image generation models. We refer to it as a "high-level" pipeline because CLIP embeddings are inherently more semantic than perceptual, given that CLIP image encoders were trained to maximize similarity with text captions (low-level features like color and object location are not typically preserved in these captions). MindEye can be used without the low-level pipeline, which simply aids to better preserve low-level image features during reconstruction.
The MLP backbone for our high-level pipeline maps flattened voxels to an intermediate space of size  $257 \times 768$ , corresponding to the last hidden layer of CLIP ViT/L-14 (see Appendix 1 for PyTorch model code). The backbone consists of a linear layer followed by 4 residual blocks and a final linear projector. The embeddings from the backbone are fed to an MLP projector and a diffusion prior in parallel. The whole model is trained end-to-end with the prior getting an MSE loss and the projector getting a bidirectional CLIP loss. The projector outputs can be used for retrieval tasks and the diffusion prior outputs can be used by generative models to reconstruct images.
Contrastive Learning: Contrastive learning is an effective method for learning representations across modalities by maximizing cosine similarity for positive pairs while minimizing similarity for negative pairs. Previous work has shown the potential benefits of using contrastive learning alongside neural data [12, 13]. CLIP [7] is an example of a multimodal contrastive model that maps
images and text captions to a shared embedding space. MindEye is trained to introduce fMRI as an additional modality to the embedding space of a pretrained CLIP model, keeping the CLIP image space frozen as done with locked-image text tuning (LiT) [14]. We use the CLIP loss [7] as our contrastive objective. This loss is bidirectional and helps improve both image and brain retrieval.
Recent work [15-18] has explored novel data augmentation techniques that offer several benefits like improving performance, increasing robustness, and reducing training data requirements. Mixup [15] is one such technique that trains models on synthetic data created through convex combinations of two datapoint-label pairs [19]. Kim et al. [20] introduce MixCo, an extension of mixup that uses the InfoNCE loss, and show that MixCo improves classification performance in a semi-supervised setting. Based on the same principle, we modify the bidirectional CLIP loss to use MixCo. While Kim et al. [20] observed that MixCo gives largest performance benefit for smaller models, we observe that it also helps large models in low data regimes.
To combine MixCo with CLIP loss, we mix voxels using a factor  $\lambda$  sampled from the Beta distribution with  $\alpha = \beta = 0.15$ .
$$
x _ {\operatorname {m i x} _ {i, k _ {i}}} = \lambda_ {i} \cdot x _ {i} + (1 - \lambda_ {i}) \cdot x _ {k _ {i}}, \quad p _ {i} ^ {*} = f \left(x _ {\operatorname {m i x} _ {i, k _ {i}}}\right), \quad p _ {i} = f \left(x _ {i}\right), \quad t _ {i} = \mathrm {C L I P} _ {\text {I m a g e}} \left(y _ {i}\right) \tag {1}
$$
Here,  $x_{i}$  and  $y_{i}$  represent the  $i$ -th fMRI sample and image respectively.  $k_{i} \in [1,N]$  is an arbitrary mixing index for the  $i$ -th datapoint and  $f$  represents the combined MLP and projector.  $p^{*}$ ,  $p$  and  $t$  are L2-normalized. The CLIP loss with MixCo is defined as:
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {B i M i x C o}} = - \sum_ {i = 1} ^ {N} \left[ \lambda_ {i} \cdot \log \left(\frac {\exp \left(\frac {p _ {i} ^ {*} \cdot t _ {i}}{\tau}\right)}{\sum_ {m = 1} ^ {N} \exp \left(\frac {p _ {i} ^ {*} \cdot t _ {m}}{\tau}\right)}\right) + (1 - \lambda_ {i}) \cdot \log \left(\frac {\exp \left(\frac {p _ {i} ^ {*} \cdot t _ {k _ {i}}}{\tau}\right)}{\sum_ {m = 1} ^ {N} \exp \left(\frac {p _ {i} ^ {*} \cdot t _ {m}}{\tau}\right)}\right) \right] \\ - \sum_ {j = 1} ^ {N} \left[ \lambda_ {j} \cdot \log \left(\frac {\exp \left(\frac {p _ {j} ^ {*} \cdot t _ {j}}{\tau}\right)}{\sum_ {m = 1} ^ {N} \exp \left(\frac {p _ {m} ^ {*} \cdot t _ {j}}{\tau}\right)}\right) + \sum_ {\{l | k _ {l} = j \}} (1 - \lambda_ {l}) \cdot \log \left(\frac {\exp \left(\frac {p _ {l} ^ {*} \cdot t _ {j}}{\tau}\right)}{\sum_ {m = 1} ^ {N} \exp \left(\frac {p _ {m} ^ {*} \cdot t _ {j}}{\tau}\right)}\right) \right] \tag {2} \\ \end{array}
$$
We term this bidirectional loss as BiMixCo. Here  $\tau$  is a temperature hyperparameter, and  $N$  is the batch size.
Recent works [21, 22] have shown that stopping mixup augmentation after a certain number of epochs leads to better classification performance. As per these findings, we stop using mixup and switch from a hard contrastive loss to a soft contrastive loss one-third of the way through training. This improves our reconstructions without harming our retrieval performance (see Table 4). BiMixCo gives the highest retrieval performance but slightly hurts reconstructions (Table 4) likely due to how the reconstruction task needs absolute targets for the mixed inputs (which we generate by mixing the original targets in the same ratio as the mixup inputs), causing a slight shift in the distribution of target embeddings. Our final schedule combines BiMixCo and soft contrastive loss to strike the best balance between retrieval and reconstruction performance in a single model.
Our soft contrastive loss is inspired by knowledge distillation [23], where the authors argue that the softmax probability distribution produced by a powerful teacher model acts as a better teaching signal for a student than hard labels. To generate the soft labels we take the dot product of CLIP image embeddings in a batch with themselves. The loss (with bidirectional component omitted for brevity) is calculated between CLIP-CLIP and Brain-CLIP matrices as:
$$
\mathcal {L} _ {\text {S o f t C L I P}} = - \sum_ {i = 1} ^ {N} \sum_ {j = 1} ^ {N} \left[ \frac {\exp \left(\frac {t _ {i} \cdot t _ {j}}{\tau}\right)}{\sum_ {m = 1} ^ {N} \exp \left(\frac {t _ {i} \cdot t _ {m}}{\tau}\right)} \cdot \log \left(\frac {\exp \left(\frac {p _ {i} \cdot t _ {j}}{\tau}\right)}{\sum_ {m = 1} ^ {N} \exp \left(\frac {p _ {i} \cdot t _ {m}}{\tau}\right)}\right) \right] \tag {3}
$$
Diffusion Prior: Using a diffusion model to align the outputs of a contrastive learning model was inspired by DALL-E 2 [8], where a "diffusion prior" was used to map CLIP text embeddings to CLIP
image space before using an unCLIP decoder to reconstruct images. Here we train our own diffusion prior from scratch to map CLIP fMRI embeddings to CLIP image space, which are then fed into a pretrained Versatile Diffusion model to generate image reconstructions. We modified an open-source implementation of the DALL-E 2 diffusion prior available on GitHub (see Appendix A.3.1). We used the same prior loss as Ramesh et al. [8]. Our total end-to-end loss is defined as:
$$
\mathcal {L} = \mathcal {L} _ {\text {B i M i x C o} | \text {S o f t C L I P}} + \alpha \cdot \mathcal {L} _ {\text {p r i o r}} \tag {4}
$$
We use  $\alpha = 0.3$  and switch from BiMixCo to SoftCLIP after one-third of the train cycle. All our models are trained on a single A100 GPU for 240 epochs with a batch size of 32. Despite a high parameter count, MindEye (including both high- and low-level pipelines) can be trained on a single A100 in less than 18 hours. This efficiency is due to the bulk of the parameters stemming from MLPs, which are faster to compute than transformers or CNNs.
The diffusion prior is critical for reconstruction because contrastive learning only incentivizes the CLIP fMRI embeddings to match the vector direction of the associated CLIP image embeddings. This generates disjointed embeddings as observed by Ramesh et al. [8]. Theoretically, multimodal contrastive learning will always produce disjointed embeddings because of the "modality gap" phenomenon whereby encoding modalities into a shared space restricts the effective embedding space to a narrow cone in geometric space [24]. We use pre-trained models that expect CLIP image embeddings as input, thus motivating our training of a diffusion prior to align disjointed embeddings.
To rectify this issue, the diffusion prior learns a distribution of CLIP image embeddings conditioned on CLIP fMRI embeddings. UMAP [25] plots of disjointed CLIP fMRI embeddings next to aligned CLIP fMRI embeddings in Appendix A.5 show how the diffusion prior addresses the disjointed embedding spaces problem. We observe that the prior's role cannot be fulfilled by simply adding MSE loss to the MLP projector in Table 4. This is because there is a tradeoff between reconstruction and retrieval objectives and a model cannot effectively learn a single embedding space that does well on both.
2.2 Low-Level (Perceptual) Pipeline
The low-level pipeline maps voxels to the embedding space of Stable Diffusion's VAE. The output of this pipeline can be fed into the VAE decoder to produce blurry image reconstructions that lack high-level semantic content but exhibit state-of-the-art low-level image metrics. We use img2img [10] to improve our final image reconstructions in terms of low-level metrics, with minimal impairment to high-level metrics, such that we start the diffusion process from the noised encodings of our blurry reconstructions rather than pure noise.
The MLP backbone for our low-level pipeline follows the same architecture as our high-level pipeline, except that the final outputs are of size (16, 16, 64). These are upsampled to (64, 64, 4) by a CNN upsampler. An MLP projector projects the backbone outputs to a 512 dimensional space where an auxiliary contrastive loss is applied. For more information on the low-level pipeline see Appendix A.3.2. See Appendix Figure 7 for example blurry reconstructions and Appendix Table 5 to see the effect of changing img2img strength on subsequent reconstruction metrics.
For all experiments, we used the Natural Scenes Dataset (NSD) [26], a public fMRI dataset containing the brain responses of human participants passively viewing natural scenes from MS-COCO [27]. By utilizing MS-COCO, this dataset provides measured brain responses to rich naturalistic stimuli, allowing us to study how well low- and high-level image features are reconstructed by MindEye. We used the same standardized train/test splits as other NSD reconstruction papers [3, 4, 28], training subject-specific models for each of 4 participants. We averaged across three same-image repetitions for the test set (leaving 982 test samples) but not the training set (24,980 training samples), similar to Takagi and Nishimoto [3]. For more information on NSD and data preprocessing see Appendix A.2; for single-trial and reduced dataset results see Appendix A.9 and Appendix A.10.
Figure 3: MindEye image retrieval. Given a pool of candidate images, nearest neighbor search in CLIP space enables searching for the original image based on brain activity. Top section depicts how, given 982 test NSD images (many containing very similar looking images, e.g., over a dozen zebras), MindEye top-1 performance is  $93.2\%$  for Subject 1. The ability to distinguish among confusable candidates suggests brain embeddings retain fine-grained, image-specific information. Bottom section depicts scaling up to the LAION-5B dataset (see Appendix A.4 for more examples). Even with billions of images, MindEye finds images similar to the original.
3.1 Image/Brain Retrieval
Image retrieval evaluations reveal the level of fine-grained image-specific information contained in the predicted brain embeddings. For example, if the model is given a dozen pictures of zebras and the brain sample corresponding to viewing one of those zebras, can the model correctly find the corresponding zebra? If the model can correctly deduce that the brain sample corresponds to an image of a zebra but cannot deduce the specific image amongst various candidates, this would suggest that category-level information but not exemplar-specific information is preserved in the CLIP fMRI embedding. MindEye not only succeeds in this zebra example but also demonstrates  $93.2\%$  overall accuracy for Subject 1 in finding the exact original image within the 982 test images (see Figure 3).
Although we use the full test dataset for retrievals in Figure 3, we followed the same procedure as Lin et al. [11] for calculating the retrieval metrics reported in Table 1. Brain retrieval performance was calculated according to the following procedure: for each test image, the image is converted to a CLIP image embedding and we compute the cosine similarity to both its respective ground truth disjointed CLIP fMRI embedding as well as 299 other randomly selected disjointed CLIP fMRI embeddings in the test set. For each test sample, success is determined if the cosine similarity is greatest between the ground truth CLIP embedding and its respective fMRI embedding (aka top-1 retrieval performance, chance  $= 1 / 300$ ). We average retrieval performance across all test samples and repeat the entire process 30 times to account for the variability in random sampling of batches. For image retrieval, the same procedure is used except image and brain samples are flipped such that the goal is to find the corresponding paired CLIP image embedding out of 300 possible CLIP embeddings in the batch. Lin et al. [11] refer to image retrieval as "forward retrieval" and brain retrieval as "backward retrieval" in their paper.
We can scale up image retrieval using a pool of billions of image candidates. In Figure 3 we show results querying the LAION-5B dataset [29] using our CLIP fMRI embeddings. The final layer CLIP ViT-L/14 embeddings for all 5 billion images are available at knn.laion.ai, and can be queried for K-nearest neighbor lookup via the CLIP Retrieval client [30]. For each test sample, we first retrieve 16 candidate images using this method (using a variant of MindEye that maps voxels to the final
Figure 4: Side-by-side comparison of reconstructions from fMRI-to-Image NSD papers. The same test set was used across papers. All reconstructions come from Subject 1.
layer of CLIP, see Appendix A.7). The best image is then selected based on having the highest CLIP embedding cosine similarity to the CLIP fMRI embedding. This image retrieval approach is especially well-suited for tasks involving fine-grained classification, and can be used as an alternative to image reconstruction without a generative model (evaluations in Table 1).
3.2 fMRI-to-Image Reconstruction
The diffusion prior outputs from MindEye are aligned CLIP fMRI embeddings that can be used with any pretrained image generation model that accepts latents from CLIP image space. We evaluate the outputs of MindEye reconstructions across several models including Versatile Diffusion [6], Stable Diffusion (Image Variations) [31], and Lafite [32, 11]. Here we report results from Versatile Diffusion since it yielded the best results, and we report results from the other models in Appendix A.7. We qualitatively compare our reconstructions side-by-side with outputs from other fMRI-to-image reconstruction models in Figure 4 and quantitatively compare against other models in Table 1, demonstrating state-of-the-art MindEye reconstructions.
For each subject, for each test brain sample, we output 16 CLIP image embeddings from MindEye and feed these embeddings through the image variations pipeline of Versatile Diffusion. This produces 16 image reconstructions per brain sample. For our reconstructions we use 20 denoising timesteps with UniPCMultistep noise scheduling [33] and start the denoising process from the noised output of our low-level pipeline (img2img). We then select the best of 16 reconstructions by computing last hidden layer CLIP embeddings and picking the image with the highest cosine similarity to the disjointed CLIP fMRI embedding. This automatic second-order selection was inspired by DALL-E 2 [8], which used a similar process of selecting the best of 2 generated samples.