Reconstruindo o Olho da Mente: fMRI-para-Imagem com Aprendizado Contrastivo e Priors de Difusão

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

May 29, 2023

Autores: 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

cs.AI

Resumo

Apresentamos o MindEye, uma nova abordagem de fMRI-para-imagem para recuperar e reconstruir imagens visualizadas a partir da atividade cerebral. Nosso modelo é composto por dois submódulos paralelos especializados para recuperação (usando aprendizado contrastivo) e reconstrução (usando um prior de difusão). O MindEye pode mapear a atividade cerebral de fMRI para qualquer espaço latente multimodal de alta dimensionalidade, como o espaço de imagens CLIP, permitindo a reconstrução de imagens usando modelos generativos que aceitam embeddings desse espaço latente. Comparamos de forma abrangente nossa abordagem com outros métodos existentes, utilizando tanto comparações qualitativas lado a lado quanto avaliações quantitativas, e mostramos que o MindEye alcança desempenho de ponta tanto em tarefas de reconstrução quanto de recuperação. Em particular, o MindEye pode recuperar a imagem original exata mesmo entre candidatos altamente semelhantes, indicando que seus embeddings cerebrais retêm informações específicas da imagem em nível detalhado. Isso nos permite recuperar imagens com precisão mesmo em bancos de dados de grande escala, como o LAION-5B. Demonstramos por meio de ablações que as melhorias de desempenho do MindEye em relação a métodos anteriores resultam de submódulos especializados para recuperação e reconstrução, técnicas de treinamento aprimoradas e modelos treinados com ordens de magnitude mais parâmetros. Além disso, mostramos que o MindEye pode preservar melhor as características de baixo nível das imagens nas reconstruções ao usar img2img, com saídas de um autoencoder separado. Todo o código está disponível no GitHub.

English

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.