UniBrain: A Unified Model for Cross-Subject Brain Decoding

TL;DR

Cross-subject brain decoding from fMRI is challenged by subject-specific heterogeneity. The paper proposes UniBrain, a fully unified, parameter-free-per-subject model with a group-based extractor, a Transformer-based mutual assistance embedder, and a bilevel feature alignment to uncover cross-subject commonalities. It leverages voxel grouping to standardize signal length, transformer-based cross-subject embeddings, and CLIP-guided alignment to guide diffusion-based stimulus reconstruction. On the NSD benchmark, UniBrain achieves competitive performance with far fewer parameters and introduces a cross-subject OOD generalization benchmark to promote population-level generalization. The work provides public code and advocates for datasets and benchmarks focusing on cross-subject commonalities to advance generalized brain decoding.

Abstract

Brain decoding aims to reconstruct original stimuli from fMRI signals, providing insights into interpreting mental content. Current approaches rely heavily on subject-specific models due to the complex brain processing mechanisms and the variations in fMRI signals across individuals. Therefore, these methods greatly limit the generalization of models and fail to capture cross-subject commonalities. To address this, we present UniBrain, a unified brain decoding model that requires no subject-specific parameters. Our approach includes a group-based extractor to handle variable fMRI signal lengths, a mutual assistance embedder to capture cross-subject commonalities, and a bilevel feature alignment scheme for extracting subject-invariant features. We validate our UniBrain on the brain decoding benchmark, achieving comparable performance to current state-of-the-art subject-specific models with extremely fewer parameters. We also propose a generalization benchmark to encourage the community to emphasize cross-subject commonalities for more general brain decoding. Our code is available at https://github.com/xiaoyao3302/UniBrain.

Figures (4)

• Figure 1: Current brain decoding methods either train separate models for each subject (a) or train a partially unified model with subject-specific modules (b). Such methods greatly limit the generalization of the model with little scalability and struggle to capture cross-subject commonalities. In this paper, we for the first time propose a unified model for fMRI decoding cross subjects which requires no subject-specific modules (c), which significantly reduces the number of parameters while capturing the potential pattern across different subjects, and can be directly adopted for brain decoding on new subjects.
• Figure 2: An overview of our proposed unified model for cross-subject brain decoding (UniBrain), which includes a unified extractor and a unified mutual assistance embedder. The unified extractor includes three group-based extractors to extract semantic, geometric and mutual global representations from fMRI signals, respectively. The unified mutual assistance embedder includes two Transformer-based embedders to decode coarse semantic and geometric embeddings, respectively, and one Transformer-based mutual embedder to decode fine semantic and geometric embeddings. The mutual embedder takes as input the mutual global representation, the coarse semantic and geometric embeddings. The decoded embeddings can be used as the guidance of the Versatile Diffusion model for stimulus reconstruction. Note that the CLIP models only provide supervision during the training process, which are not utilized during the inference stage. Besides, the weights of the CLIP models and the Versatile Diffusion model are frozen.
• Figure 3: Visualizations of different fMRI decoding methods, including our newly proposed UniBrain method, our UniBrain method using subject-specific parameters and MindBridge, on the NSD dataset. All of the fMRI signals are from subject 1.
• Figure 4: Visualizations of the generalization of our newly proposed UniBrain on the NSD dataset, including both the successful cases (row 1 and row 2), and the failure cases (row 3). The model is trained using subjects 1, 2, 5, and tested on subject 7.