image generated on title prompt by chatGPT

Image reconstruction from Brain signals

Reconstructing images from fMRI signals of subjects viewing the visuals, utilizing ridge regression and trained diffusion models.

I was curious about trying out diffusion models for different applications, and that’s when I stumbled upon this interesting use case in a few research papers. It was about using them for reconstructing images from fMRI signals, and it immediately caught my attention.

• High-Resolution Image Reconstruction With Latent Diffusion Models From Human Brain Activity — Yu Takagi, Shinji Nishimoto. https://github.com/yu-takagi/stablediffusionreconstruction
• Reconstructing the Mind’s Eye: fMRI-to-Image with Contrastive Learning and Diffusion Priors — Paul S. Scotti et al. https://github.com/MedARC-AI/MindEyeV2
• Natural scene reconstruction from fMRI signals using generative latent diffusion-Furkan Ozcelik, Rufin VanRullen. https://github.com/ozcelikfu/brain-diffuser

Okay, now this write-up, I will detail my findings and experiments with this. I started by reading these papers and trying to recreate some of them. All the papers have the source code repo attached, so I thought to just run it and see.

Use case

We’ve got fMRI signals paired with the images people were looking at when the signals were recorded. The goal is to build a model or pipeline that can take these signals and recreate the images they saw. Sounds like magic, right?

Solution

Starting without these initial findings, I thought it would involve complex training of an image generator model such as a diffusion model/GANs from scratch. It did not hit me that we can somehow leverage pre-trained diffusion models for this and don’t even have to fine-tune or change any weights of the diffusion models themselves.

But going through these papers, showed me an effortless way of using pretrained-diffusion models and regressors to do it. let’s see how.

Overall flow of brain-diffuser

This image gives us an overall idea. It explains ideas from brain-diffuser paper. The brain-diffuser process begins with fMRI signals from brain activity, which are input into a Variational Autoencoder (VAE) to generate a rough base image reconstruction. These fMRI signals are also then processed by a CLIP model to extract corresponding text and vision embeddings. These embeddings, along with the base image, are fed into a diffusion model, which leverages the combined information to refine and produce a high-resolution reconstruction that aligns closely with the original visual stimulus.

Now we know that the VAE cannot directly generate base image reconstructions from fMRI signals, nor can CLIP directly extract image embeddings from them. But we can do a trick, here regressors are employed as intermediaries. Here’s the step-by-step process:

Step 1: Train VAE regressors

Imagine a model capable of generating latent embeddings (from a Variational Autoencoder) corresponding to fMRI signals. These latent embeddings can then be used as input to the VAE decoder to produce an image that represents the embeddings, and by extension, the original fMRI signals. In this way, we only have to train this model, and not change anything in the VAE model, we can pre-trained weights. (VAE trained on image data for reconstruction of the input image)

Yes, we can do that by using simple Ridge regression models. We create these models to map, the fMRI signals input to Z latent embeddings.

But now how to train them? we don’t have fMRI signals to Z latent embeddings of image pairs as a dataset. So, we do this:

1. Get this from the Natural Scenes Dataset(NSD): fMRI signal and corresponding images pair.
2. Feed images to (pre-trained)VAE encoder, get corresponding Z latent embedding of image. Now we have fMRI signal and Z-latent embedding pairs.
3. Train a ridge regression model using this pair, such that

Z-latent embeddings = ride_regression_model(fMRI)

Training VAE regressor

Now that we have a model which can get Z-latent embeddings from fMRI signals, let's use it.

4. Get predicted Z latent embedding of the test fMRI signals.

5. Using the predicted Z latent embeddings, input it to the (pre-trained). And get the image reconstructed.

Using trained VAE regressor during reconstruction

This will serve as a base image for further diffusion. While this base image retains low-resolution features such as colors, and basic shapes, it gets fine detailed in the following processes.

Instead of normal VAEs, we use Very Deep VAEs. Because, we need many parameters in the bottlenet of the AE, to represent natural images. Hence we use this model, which is sort of like a U-Net. It’s also by OpenAI, check it out here https://github.com/openai/vdvae.

Step 2: Train CLIP regressors

This is similar to the previous step. Here we are training models to map fMRI signals to CLIP embeddings, both vision and textual. CLIP is a contrastive model, that fuses image and corresponding text forming an embedding space. It, therefore, can do a lot of back-and-forth conversion between images and text. It’s one of the well-known models by OpenAI, you can read more about it here, https://openai.com/index/clip/. I think it’s used heavily in their SORA models as well.

Training steps are similar:

1. Get this again from the Natural Scenes Dataset(NSD): fMRI signal and corresponding images pair.
2. Feed images to (pre-trained)CLIP-vision model, and get corresponding vision embeddings. Similarly, send the captions of the images to the (pre-trained)CLIP-vision model and also get their corresponding text embeddings. Now we have fMRI signal and corresponding CLIP vision embeddings and text embeddings.
3. Train two ridge regression model using these pairs, such that

CLIP vision embeddings = ridge_regression_model1(fMRI)

CLIP text embeddings = ridge_regression_model2(fMRI)

training CLIP regressors

We directly use these embeddings from the regressors in the diffusion process.

using trained CLIP regressors

Step 3: Diffusion

So, so far for the fMRI signals we have ways,

• to get the “Base image” from the regressor model and VDVAE decoder.
• to get CLIP “Vision embeddings” and “Text embeddings” from regressors trained using CLIP.

So, let’s see how to use them in getting or adding more details to the base image. Enter Diffusion.

Diffusion has been one of the main techniques used for image generation for a while. It derives concepts from psychics, and I don’t know how exactly, though but it works. Recent image and video generators, using latent diffusion models, or Stable diffusion models. Read more about it here, https://jalammar.github.io/illustrated-stable-diffusion/.

Then there are diffusion generators, that can be conditioned on different inputs. That is, we can give it text, images, or other modals as inputs and get an image reconstructed as output. One such model is the Versatile Diffusion model, https://github.com/SHI-Labs/Versatile-Diffusion. As the name suggests it’s versatile in the sense that it can do many transformations. Image-to-image, text-to-image, image-to-text, and many other transformations. In the base, it uses latent diffusion, but it can take many modals.

For our case, we will start the diffusion with our base image(from the VDVAE decoder), and condition it with CLIP vision and text embeddings(it integrates these embeddings. It uses a Cross-attention mechanism during the reverse diffusion process, to add the embeddings into the diffusion U-Net model.

entire pipeline, based on brain-diffuser.

so, these are the steps:

1. Input fMRI signals to the three ridge regressors.

z = model1(fMRI)

vision_embeddings = model2(fMRI)

text_embeddings = model3(fMRI)

2. Put the predicted z from model1, into VDVAE decoder, giving us the base image.

Base_image = VDVAE_decoder(z)

3. Start diffusion in Versatile diffusion model(pre-trained) with the base image. Condition it on both vision and text embeddings.

4. At the end of the diffusion process, we can get the reconstruction images.

Results

These are the results I got, which is great that I did’nt do any specific training of the diffusion models at all. (Which would take days to train, on heavy GPU requirements and a soul).

You check my code from this repo, https://github.com/sarabesh/Neural-Recon/tree/main, which has a decent enough readme to get you running this. But you would have to download the NSD dataset, which is close to 100GB, after getting permission. Dataset page: https://naturalscenesdataset.org/. This pipeline is heavily based on brain-diffuser: https://github.com/ozcelikfu/brain-diffuser/tree/main. The other papers mentioned in the beginning also do similar pipelines, with different models. From the results in the paper, this seems to be the advanced one using Stable diffusion XL for reconstruction. https://github.com/MedARC-AI/MindEyeV2

I think all three papers are worth a read if you are interested in this. I am just amazed by this application of pre-trained diffusion models and CLIP, perhaps by using them we can similar applications in different domains needing image reconstructions. This approach highlights the power of combining latent embeddings, regressors, and advanced generative models to bridge complex data representations and meaningful outputs. By refining each step, we can unlock the full potential of these models for innovative applications. Thank you for reading!

References

• High-Resolution Image Reconstruction With Latent Diffusion Models From Human Brain Activity Authors: Yu Takagi, Shinji Nishimoto GitHub Repository: https://github.com/yu-takagi/stablediffusionreconstruction
• Reconstructing the Mind’s Eye: fMRI-to-Image with Contrastive Learning and Diffusion Priors Authors: Paul S. Scotti et al. GitHub Repository: https://github.com/MedARC-AI/MindEyeV2
• Natural Scene Reconstruction From fMRI Signals Using Generative Latent Diffusion Authors: Furkan Ozcelik, Rufin VanRullen GitHub Repository: https://github.com/ozcelikfu/brain-diffuser
• Natural Scenes Dataset Website: https://naturalscenesdataset.org/
• Very Deep VAE GitHub Repository: https://github.com/openai/vdvae
• Versatile Diffusion GitHub Repository: https://github.com/SHI-Labs/Versatile-Diffusion