UMBRAE: Unified Multimodal Decoding of Brain Signals

Weihao Xia¹ Raoul de Charette² Cengiz Öztireli³ Jing-Hao Xue<sup>1

1</sup>University College London ²Inria ³University of Cambridge

Multimodal Brain Decoding from our UMBRAE applied on the pratrained MLLMs with given prompts.

[Paper]      [Code]      [BrainHub]

Abstract

We address prevailing challenges in the brain-powered research, departing from the observation that the literature hardly recover accurate spatial information and require subject-specific models. To address these challenges, we propose UMBRAE, a unified multimodal decoding of brain signals. First, to extract instance-level conceptual and spatial details from neural signals, we introduce an efficient universal brain encoder for multimodal-brain alignment and recover object descriptions at multiple levels of granularity from subsequent multimodal large language model (MLLM). Second, we introduce a cross-subject training strategy that maps subject-specific features to a common feature space. This allows a model to be trained on multiple subjects without extra resources, even yielding superior results compared to subject-specific models. We further show that it supports weakly-supervised adaptation to new subjects, with only a fraction of the total training data. Experiments demonstrate that UMBRAE not only achieves superior results in the newly-introduced grounding tasks but also outperforms methods in well established tasks. To assess our method, we construct and share with the community a comprehensive brain understanding benchmark BrainHub.

Keywords: Multimodal Brain Decoding, Efficient Brain Encoder, Cross-Subject Training, Weakly-Supervised Adaptation, Large Language Model

Method

UMBRAE decodes multimodal explanations from brain signals. (1) We introduce a universal brain encoder for multimodal-brain alignment and recover conceptual and spatial details by using multimodal large language models. (2) We introduce cross-subject training to overcome unique brain patterns of different individuals. This allows brain signals from multiple subjects to be trained within the same model. (3) Our method supports weakly-supervised subject adaptation, enabling the training of a model for a new subject in a data-efficient manner. (4) For evaluation, we introduce BrainHub, a brain understanding benchmark, based on NSD and COCO.

Overview of UMBRAE. Our brain encoder includes subject-specific tokenizers and a universal perceive encoder. Brain signals from multiple subjects are mapped into a common feature space, enabling cross-subject training and weakly-supervised subject adaptation. The brain encoder learns to align neural signals with image features. During inference, the learned encoder interacts with MLLMs and performs brain understanding tasks according to given prompts.

Demo Examples

UMBRAE decodes multimodal explanations from brain signals.

Brain Captioning

Brain Captioning aims at textually describing the primary content of a given brain response. Here are the brain captioning comparison (on S1). The Baselines include SDRecon, BrainCap, and OneLLM. UMBRAE-S1 refers to our model trained only with subject S1, while UMBRAE denotes the model with cross-subject training. The left is the reference image.

Below are the Brain Captioning results on different subjects. “COCO” is the ground truth caption in the COCO dataset. We excerpt the first caption if there are multiple captions for the same image. Shikra-w/img is the result using the ground truth images as input. Results for all four subjects (S1, S2, S5, and S7) are from our cross-subject UMBRAE model.

Brain Grounding

Brain Grounding is the counterpart of visual grounding and seeks to recover spatial locations from brain signals by inferring coordinates. Here are the Brain Grounding results from UMBRAE using different task prompts. Reference images are visual stimuli for input brain responses and are just used here for visualization.

Below are the UMBRAE Brain Grounding results of different subjects.

Brain Retrieval

The forward retrieval (top) is to find the correct image embedding given a brain embedding. Conversely, the backward retrieval (bottom) aims to locate the correct brain embedding given an image embedding.

Visual Decoding

Despite that our method is not specifically designed for the task of visual decoding (fMRI-to-image reconstruction), our predicted textual and visual cues can be used for the final image reconstruction by using a variety of pretrained image generation models. These models include text-to-image SD, SD-XL, and Kandinsky3, a layout-to-image GLIGEN, and a multiple-condition Versatile Diffusion.

BrainHub

For evaluation, we introduce BrainHub, a brain understanding benchmark, based on NSD and COCO.

There are 982 test images, 80 classes, 4,913 captions, and 5,829 boundingboxes. For grounding evaluation, we further group the 80 classes of COCO into 4 salience categories according to their salience in images: Salient (S), Salient Creatures (SC), Salient Objects (SO), and Inconspicuous (I). The illustration shows the statistics and mapping of our categories, w.r.t. COCO classes.

BibTex

@article{xia2024umbrae,
  title={UMBRAE: Unified Multimodal Decoding of Brain Signals},
  author={Xia, Weihao and de Charette, Raoul and Öztireli, Cengiz and Xue, Jing-Hao},
  booktitle={arxiv preprint:arxiv 2404.07202},
  year={2024}
}

Acknowledgements: The project template is taken from DreamBooth.