NeuralOOD: Improving Out-of-Distribution Generalization Performance with Brain-machine Fusion Learning Framework

TL;DR

NeuralOOD tackles the OOD generalization gap in vision by fusing visual features with cognitive priors from the human brain. The BMFL framework combines a frozen image encoder (DINOv2), a pre-trained brain fMRI predictor, a brain transformer, and a cross-attention fusion module, with a Pearson correlation coefficient-based regularization to align modalities. Empirical results on ImageNet-1k and six curated OOD datasets show BMFL outperforming strong baselines like DINOv2, with ablations confirming the benefits of cross-attention fusion and ROI-based brain features. This approach demonstrates the practical value of brain-informed multimodal fusion for robust, real-world vision systems.

Abstract

Deep Neural Networks (DNNs) have demonstrated exceptional recognition capabilities in traditional computer vision (CV) tasks. However, existing CV models often suffer a significant decrease in accuracy when confronted with out-of-distribution (OOD) data. In contrast to these DNN models, human can maintain a consistently low error rate when facing OOD scenes, partly attributed to the rich prior cognitive knowledge stored in the human brain. Previous OOD generalization researches only focus on the single modal, overlooking the advantages of multimodal learning method. In this paper, we utilize the multimodal learning method to improve the OOD generalization and propose a novel Brain-machine Fusion Learning (BMFL) framework. We adopt the cross-attention mechanism to fuse the visual knowledge from CV model and prior cognitive knowledge from the human brain. Specially, we employ a pre-trained visual neural encoding model to predict the functional Magnetic Resonance Imaging (fMRI) from visual features which eliminates the need for the fMRI data collection and pre-processing, effectively reduces the workload associated with conventional BMFL methods. Furthermore, we construct a brain transformer to facilitate the extraction of knowledge inside the fMRI data. Moreover, we introduce the Pearson correlation coefficient maximization regularization method into the training process, which improves the fusion capability with better constrains. Our model outperforms the DINOv2 and baseline models on the ImageNet-1k validation dataset as well as six curated OOD datasets, showcasing its superior performance in diverse scenarios.

Figures (5)

• Figure 1: DNNs can achieve or even surpass human recognition capabilities when presented with common data. However, their recognition performance significantly deteriorates when faced with OOD data, in contrast to the resilient recognition abilities exhibited by humans in such scenarios.
• Figure 2: (A). Brain-machine Fusion model. Our learning process can be divided into four steps: image feature extraction, generation of brain fMRI via image features, fMRI feature extraction and multimodal modeling. The fused multimodal representation is used to predict the corresponding category via a linear classifier. (B). The frozen weighted DINOv2 model is applied as the image encoder, from which image [CLS] token and patched image tokens are extracted. (C). The training process and architecture of the brain encoder. The NSD dataset is used to train the model, the image patch tokens are extracted from the image encoder with the input of the NSD image stimuli. The brain encoder is optimized with mean square error (MSE) loss calculated between the predicted fMRI and the raw fMRI data.
• Figure 3: The reported result from adeli2023predicting of the correlation between the predicted fMRI and the ground truth fMRI.
• Figure 4: (A). The fMRI feature extractor.Inspired by ViT, we propose the brain transformer to extract knowledge from fMRI data. We slice the fMRI into several patches using convolutional operations, then add learnable position embeddings along with the [CLS] token. These fMRI patch tokens and [CLS] token are the input transformer encoder layer. (B). Brain-visual fusion module. The cross-attention mechanism has proven its strong capability in multimodal fusion in generative models. In the proposed BMFL framework, the [CLS] token is used as query of the cross-attention, while another modal's patch tokens are used as the key and value. The fused features are then concatenated with the global average pooling features of image patches to form the input for the linear classifier.
• Figure 5: Two self-created and four open-source OOD datasets based on the ImageNet-1K dataset.