MedARC-AI / MindEye_Imagery

MIT License
4 stars 0 forks source link

Train brain encoding model for shared-subject latent space #13

Closed reesekneeland closed 1 week ago

reesekneeland commented 4 months ago

Right now brain retrieval tasks are limited by samples we have matched fMRI pairs for. We are experimenting with using existing brain encoding models (GNet) to generate synthetic data that can be used as input to the ridge layer of ME2 (see separate github issue for this work), but another idea is to train an entirely new encoding model that takes an image as input, and produces as output an embedding in the ME2 shared subject space.

This encoding model will have to be tied to a specific checkpoint of ME2 to work, but the idea is that we can skip over the noisy space of the brain and get a better synthetic library if we map directly to the shared subject latent space instead of to raw voxels.

We don't have prepared code to retrain GNet at the moment (the existing code is messy and will take too long to get working again), so I would suggest as a first step picking another good encoding model architecture that has code to train (I might be able to help find one if needed), and training it to map to the shared subject embeddings. Because we are bypassing the subject-specific layer of ME2, we can use the shared-subject space embeddings for all subjects to train this model.

Once we have a model that can map images -> shared-subject space embeddings of ME2, we can feed arbitrary images through ME2 to create a synthetic brain retrieval library, and we should experiment with using the same pool of subject1 training images as a retrieval pool, as well as other subjects or other image distributions.