[Paper Review] MindEye2: Shared-Subject Models Enable fMRI-To-Image With 1 Hour of Data

Paul S. Scotti, Mihir Tripathy+9 more|arXiv (Cornell University)|Mar 17, 2024

Functional Brain Connectivity Studies5 citations

TL;DR

MindEye2 pretrains a single model across multiple subjects and fine-tunes it on just 1 hour of held-out-subject fMRI data to achieve state-of-the-art fMRI-to-image reconstructions and retrieval, using a shared-subject latent space and fine-tuned unCLIP/SDXL pipeline.

ABSTRACT

Reconstructions of visual perception from brain activity have improved tremendously, but the practical utility of such methods has been limited. This is because such models are trained independently per subject where each subject requires dozens of hours of expensive fMRI training data to attain high-quality results. The present work showcases high-quality reconstructions using only 1 hour of fMRI training data. We pretrain our model across 7 subjects and then fine-tune on minimal data from a new subject. Our novel functional alignment procedure linearly maps all brain data to a shared-subject latent space, followed by a shared non-linear mapping to CLIP image space. We then map from CLIP space to pixel space by fine-tuning Stable Diffusion XL to accept CLIP latents as inputs instead of text. This approach improves out-of-subject generalization with limited training data and also attains state-of-the-art image retrieval and reconstruction metrics compared to single-subject approaches. MindEye2 demonstrates how accurate reconstructions of perception are possible from a single visit to the MRI facility. All code is available on GitHub.

Motivation & Objective

• Motivate practical fMRI-to-image reconstruction with minimal subject-specific data.
• Develop a shared-subject functional alignment to map diverse brains into a common latent space.
• Integrate a unified pipeline mapping fMRI to CLIP space and reconstructing images via a fine-tuned unCLIP/SDXL model.
• Demonstrate strong retrieval and captioning capabilities alongside high-fidelity reconstructions.

Proposed method

• Pretrain a single model on fMRI data from 7 subjects, then fine-tune on a held-out 8th subject.
• Use a linear ridge mapping to map voxel responses to a 4096-d shared-subject latent space, followed by a residual MLP backbone to OpenCLIP ViT-bigG/14 embeddings.
• Train a diffusion prior to map fMRI latents to OpenCLIP image space, with a retrieval submodule trained contrastively.
• Introduce a low-level submodule to preserve fine-grained structure and aid SDXL reconstruction.
• Fine-tune Stable Diffusion XL unCLIP to accept image embeddings instead of text and refine outputs with base SDXL using caption guidance.

Figure 1: MindEye2 overall schematic. MindEye2 is trained using samples from 7 subjects in the Natural Scenes Dataset and then fine-tuned using a target held-out subject who may have scarce training data. Ridge regression maps fMRI activity to an initial shared-subject latent space. An MLP backbone

Experimental results

Research questions

• RQ1 Can a shared-subject latent space enable high-quality fMRI-to-image reconstructions with very limited data from a new subject?
• RQ2 Does multi-subject pretraining improve generalization to held-out subjects compared to single-subject models trained from scratch?
• RQ3 How do alignment, diffusion priors, and submodules contribute to reconstruction and retrieval performance?
• RQ4 What is the impact of including image captioning as auxiliary guidance on reconstructions?
• RQ5 How do refined reconstructions compare to unrefined ones in objective metrics versus human preference?

Key findings

• Pretraining across 7 subjects and fine-tuning on 1 hour of new-subject data yields state-of-the-art reconstruction and retrieval metrics.
• The shared-subject linear alignment to a 4096-d latent space plus a unified pipeline improves generalization with limited data.
• Fine-tuning SDXL unCLIP to accept CLIP image embeddings enables high-fidelity reconstructions closely matching ground-truth images.
• Predicted image captions provide useful conditional guidance during final refinement and improve semantic fidelity.
• Refined reconstructions are preferred by human raters over unrefined results, despite some metrics favoring unrefined outputs.
• MindEye2 with 1 hour of data performs comparably to single-subject models trained with ~40x more data in certain metrics.

Figure 3: Reconstructions from different model approaches using 1 hour of training data from NSD.