Reconstructing the Mind's Eye: fMRI-to-Image with Contrastive Learning and Diffusion Priors

Paul S. Scotti; Atmadeep Banerjee; Jimmie Goode; Stepan Shabalin; Alex Nguyen; Ethan Cohen; Aidan J. Dempster; Nathalie Verlinde; Elad Yundler; David Weisberg; Kenneth A. Norman; Tanishq Mathew Abraham

Source-linked AI summary

Reconstructing the Mind's Eye: fMRI-to-Image with Contrastive Learning and Diffusion Priors

Paul S. Scotti, Atmadeep Banerjee, Jimmie Goode, Stepan Shabalin, Alex Nguyen, Ethan Cohen, Aidan J. Dempster, Nathalie Verlinde, Elad Yundler, David Weisberg, Kenneth A. Norman, Tanishq Mathew Abraham

arXiv:2305.18274v2May 29, 2023cs.CVcs.AIq-bio.NC

TL;DR

MindEye addresses the challenge of decoding viewed images from fMRI activity. It maps brain activity into CLIP space with parallel retrieval and reconstruction modules, achieving state-of-the-art reconstruction and retrieval, including exact-image retrieval above 90% on NSD test samples. The approach also supports fine-grained retrieval among similar images and large-scale databases, while remaining limited by subject-specific training and natural-scene scope.

Problem
Existing fMRI-to-image methods used relatively simple mappings, motivating improved approaches for decoding environmental inputs and cognitive states from brain activity.
Method
MindEye uses an MLP backbone with specialized contrastive retrieval and diffusion-prior reconstruction modules that map fMRI activity into CLIP image-embedding space.
Results
>90% top-1 accuracy was achieved for both image and brain retrieval on NSD test samples, compared with <50% for previous state-of-the-art methods, alongside state-of-the-art reconstruction performance.
Takeaways & Limitations
MindEye brain embeddings retain fine-grained image-specific information, enabling exact-image retrieval among similar candidates and retrieval from large-scale databases such as LAION-5B.
Takeaways & Limitations
The model does not decode across subjects, requires up to 40 hours of MRI data per participant, and is limited to natural scenes such as those in the dataset.

Abstract

from arXiv · show

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.

1 Introduction

MindEye maps fMRI activity to CLIP image embeddings using large MLPs, contrastive learning, and diffusion models for image reconstruction and retrieval. It reports state-of-the-art performance, including exact-image retrieval above 90% on NSD test samples.

Motivation: MindEye addresses fMRI decoding for understanding brain mechanisms and potential mind-reading applications by mapping brain activity to image-model embeddings.Previous approaches typically used simple mappings such as ridge regression.
Approach: MindEye maps flattened voxel activity to CLIP image-embedding space with an MLP backbone and specialized retrieval and reconstruction submodules.The retrieval submodule uses contrastive training, while reconstruction uses a diffusion prior.
Results: >90% top-1 accuracy was achieved for both image and brain retrieval on NSD test samples, compared with <50% for previous state-of-the-art methods.Image retrieval selects the original viewed image from candidate images; brain retrieval selects the matching brain sample for an image.
Findings: MindEye’s main findings attribute performance gains to task-specialized submodules, large MLP parameter counts, bidirectional mixup contrastive augmentation, and separate VAE-latent mapping.The authors also report exact-image retrieval among highly similar candidates and scaling retrieval to databases such as LAION-5B.
Approach: MindEye’s semantic pipeline supports reconstruction through generative models such as Versatile Diffusion and retrieval through nearest-neighbor searches over image embeddings.Its low-level perceptual pipeline maps voxels to a Stable Diffusion VAE space for blurry reconstructions that initialize later image generation.

2 MindEye

MindEye maps fMRI activity into CLIP and VAE image-embedding spaces through an MLP backbone and task-specific retrieval and reconstruction modules. Its training combines contrastive objectives, diffusion-based alignment, and augmentation strategies, while a low-level pipeline and img2img help preserve perceptual features.

Architecture: MindEye uses two pipelines: a high-level semantic route to CLIP ViT-L/14 space and a low-level perceptual route to Stable Diffusion’s VAE space.Both routes use a residual MLP backbone with task-specific submodules; the low-level route adds CNN upsampling.
High-level pipeline: The high-level backbone feeds an MLP projector and diffusion prior in parallel, supporting retrieval from projector outputs and reconstruction from prior outputs.The model is trained end-to-end with bidirectional CLIP loss for the projector and MSE loss for the diffusion prior.
Contrastive learning: Contrastive learning aligns fMRI with frozen CLIP image embeddings by maximizing positive-pair cosine similarity and minimizing negative-pair similarity.The bidirectional CLIP objective improves both image and brain retrieval.
Training: Bidirectional Mixup contrastive augmentation mixes fMRI samples and images, then switches to soft contrastive loss one-third through training to improve reconstructions without harming retrieval.BiMixCo provides the highest retrieval performance but slightly hurts reconstructions, motivating the training schedule.
Diffusion prior: The diffusion prior learns a distribution of CLIP image embeddings conditioned on CLIP fMRI embeddings because contrastive learning aligns directions but leaves a modality gap.Adding MSE directly to the projector does not replace the prior because reconstruction and retrieval objectives trade off.

3 Results

MindEye achieves strong image and brain retrieval and state-of-the-art reconstruction, including exact-image retrieval among confusable candidates. Ablations attribute these gains to specialized components, larger MLPs, and improved contrastive training.

Image/Brain Retrieval: 93.2% top-1 accuracy identifies Subject 1’s exact original image among 982 NSD test images, including highly similar candidates.The result indicates that predicted brain embeddings retain fine-grained, image-specific information.
Image/Brain Retrieval: >90% accuracy is achieved for both image and brain retrieval on NSD test samples, compared with <50% for previous state-of-the-art methods.Image retrieval finds the viewed image from candidate images, whereas brain retrieval finds the brain sample given an image.
Image/Brain Retrieval: MindEye scales retrieval to LAION-5B by querying billions of CLIP image embeddings and selecting the candidate most similar to the CLIP fMRI embedding.The method first retrieves 16 candidates using nearest-neighbor search, then selects the best candidate by cosine similarity.
fMRI-to-Image Reconstruction: MindEye reconstructions achieve state-of-the-art performance across quantitative comparisons and side-by-side evaluations against other fMRI-to-image models.Reconstructions use aligned CLIP fMRI embeddings with pretrained image-generation models, including Versatile Diffusion.
Ablations: Four residual blocks with skip connections improve retrieval, while a 7× smaller CLS-only model performs substantially worse than models predicting all 257 CLIP embeddings.The ablation supports benefits from larger parameter counts and preserving fine-grained information across the full CLIP representation.
Ablations: CLIP loss improves image retrieval over InfoNCE, MixCo benefits both retrieval directions, and SoftCLIP combined with BiMixCo gives the best image retrieval performance.A separate diffusion prior is also needed because joint contrastive and MSE training struggles to satisfy the tradeoff between retrieval and reconstruction objectives.

4 Related Work

fMRI decoding progressed from linear classification of visual information to generative reconstruction of increasingly complex images. MindEye differs from prior NSD approaches through its architecture, objectives, and use of a diffusion prior.

4 Related Work: Earlier fMRI studies decoded spatial position, orientation, and coarse image category using linear classifiers.Later generative methods mapped brain activity into model latent spaces to reconstruct handwritten digits, faces, and natural scenes.
4 Related Work: Generative approaches expanded fMRI decoding from classification toward reconstruction of handwritten digits, human faces, and natural scenes.
4 Related Work: Compared with Lin et al., MindEye uses a convolution-free contrastive projector, CLIP loss, a diffusion prior, and no fine-tuned GAN.Lin et al. used a convolutional model, InfoNCE, a fine-tuned Lafite GAN, and mappings to CLIP image and text space.

5 Conclusion

MindEye reconstructs natural scenes from fMRI while preserving semantic and perceptual similarity through parallel high-level and low-level pipelines. Its current scope remains limited by participant-specific training, extensive data collection, and natural-scene inputs.

5 Conclusion: MindEye achieves state-of-the-art reconstruction of natural scenes presented to people during MRI scanning.The reconstructions retain semantic and perceptual similarity to the original images.
5 Conclusion: Specialized contrastive-retrieval and diffusion-reconstruction submodules learn mappings for both tasks in parallel.A diffusion prior translates brain embeddings into pretrained CLIP space, while the high-level and low-level pipelines support semantic and perceptual reconstruction.
5 Conclusion: MindEye can select ground-truth images from nearly 1,000 possibilities, demonstrating fine-grained image-specific information in its brain latents.The method achieves this retrieval capability using CLIP as a teacher despite fewer than 30,000 training samples per participant.
5 Conclusion: Current models do not decode across subjects, require up to 40 hours of MRI data per participant, and are limited to natural scenes such as MS-COCO.Other image distributions would require additional data collection and specialized generative models.

6 Open Research: 100% Transparent Volunteer-Driven Science

MindEye was developed as an openly accessible, volunteer-driven research project with public code, discussions, and recorded meetings.

6 Open Research: 100% Transparent Volunteer-Driven Science: MindEye was openly developed through volunteer contributions in the MedARC Discord server.
6 Open Research: 100% Transparent Volunteer-Driven Science: The source code remained accessible through a public GitHub repository throughout the project.
6 Open Research: 100% Transparent Volunteer-Driven Science: Research discussions occurred in public Discord channels, and weekly video calls were recorded and shared publicly.The project continues inviting global contributions to MedARC Neuroimaging & AI Lab projects.

7 Author Contributions

Author contribution details are provided separately in Appendix A.1.

7 Author Contributions: Detailed author contributions are documented in Appendix A.1.
7 Author Contributions: Appendix A.1 is the paper’s designated location for author contribution information.
7 Author Contributions: The author-contribution section directs readers to Appendix A.1 for further details.

A.1 Author Contributions

The authors describe their respective contributions to MindEye development, including project leadership, model development, manuscript drafting, codebase work, and creation of specialized components.

PSS devised the project, led the team, developed the models, and drafted the manuscript.
AB contributed to model development and manuscript drafting, including the low-level pipeline, BiMixCo, soft CLIP loss, and DALL-E 2 prior modifications.
JG developed models, tracked and compared variants, and significantly contributed to the MindEye codebase.

A.2 Additional Dataset Information

The study uses the Natural Scenes Dataset, a large 7-Tesla fMRI dataset of participants viewing natural-scene images repeatedly across many scanning sessions.

The Natural Scenes Dataset is a public 7-Tesla fMRI dataset containing responses from participants passively viewing images.Participants spent up to 40 hours in the MRI machine.
The dataset contains 9,000–10,000 unique MS-COCO natural-scene images, each presented for three seconds and shown three times across 30–40 sessions.
Models are trained separately for participants 1, 2, 5, and 7, who completed all scanning sessions.
The resulting subject-specific dataset has 24,980 training samples and 2,770 test samples, reduced to 982 averaged test samples after combining repetitions.The three same-image repetitions are averaged for testing but not training.

A.3 MindEye Architecture

MindEye’s implementation uses an MLP backbone and projector, with the diffusion-prior details described separately in Appendix A.3.1.

The paper provides PyTorch code for the MindEye MLP backbone and projector in Algorithm 1.
The implementation separates the MLP backbone and projector from the modified DALL-E 2 diffusion prior.
Specific modifications to the open-source DALL-E 2 diffusion-prior implementation are discussed in Appendix A.3.1.

A.3.1 Modifications from DALL-E 2 Diffusion Prior

MindEye combines high-level CLIP-space reconstruction with low-level VAE-space reconstruction, using task-specific modules, diffusion-prior alignment, and img2img refinement.

A.3.1 Modifications from DALL-E 2 Diffusion Prior: The diffusion prior denoises 257 CLIP image embeddings from backbone embeddings, a timestep embedding, and noised CLIP embeddings.Unlike the DALL-E 2 prior, MindEye directly predicts denoised CLIP embeddings without learnable queries.
A.3.2 Low-Level Pipeline: Mapping to Stable Diffusion Variational Autoencoder: The low-level pipeline maps voxels to a 16 × 16 × 64 latent space and uses a CNN upsampler to produce (64, 64, 4) Stable Diffusion VAE embeddings.Targets are created by encoding 512 × 512 upsampled NSD images with the Stable Diffusion VAE encoder.
A.3.2 Low-Level Pipeline: Mapping to Stable Diffusion Variational Autoencoder: The low-level contrastive submodule provides an auxiliary loss and distills VICRegL ConvNext-XXL representations rather than CLIP representations.
A.3.1 Modifications from DALL-E 2 Diffusion Prior: UMAP comparisons show that the diffusion prior aligns MindEye embeddings with CLIP image space better than the MLP backbone or projector.The diffusion-prior space has the lowest Euclidean distance to CLIP image space.
A.3.2 Low-Level Pipeline: Mapping to Stable Diffusion Variational Autoencoder: Img2img starts diffusion from noised low-level reconstructions rather than pure noise, improving low-level metrics with minimal impairment to high-level metrics.The final MindEye uses an img2img strength of 0.15.