Mapping Deep Networks to Human Visual Brain Poster

Paper context

Paper title: Brain Decodes Deep Nets: Mapping Visual Models to Brain Regions Abstract: A conference poster exploring how deep neural networks (CLIP, SAM, DiNOv2) map to human visual brain regions, analyzing layer alignment, model scaling effects, and catastrophic forgetting after fine-tuning. Paper body (method & results): Brain Decodes Deep Nets Huzheng Yang James Gee* Jianbo Shi* University of Pennsylvania https://huzeyann.github.io/brain-decodes-deep-nets i-th brain voxel (scalar) Selected feature layer space scale channel j k l m : Factorized (layer/space/scale) and topology-smooth feature selection : Channel clustering Brain encoding model Feature extraction layer space scale channel flatten Figure 1. Visualize Deep Networks in the Brain. The training objective of the brain encoding model is to predict the brain’s fMRI signal in response to an image stimulus. 3D visual brain surface is flattened into 2D for better visualization. 1⃝Image features are extracted from a pre-trained network. 2⃝Feature selection for each voxel is randomly initialized and learned using the brain encoding training objective. The selection is factorized in the layer/space/scale axis; the topological constraint improves selection smoothness and confidence. 3⃝ Linearized brain encoding model. 4⃝After training, linear weights are used to cluster channels. We use the resulting brain-to-network mapping together with the known knowledge of the brain to answer the question “how do deep networks work?”. Abstract We developed a tool for visualizing and analyzing large pre-trained vision models by mapping them onto the brain, thus exposing their hidden inside. Our innovation arises from a surprising usage of brain encoding: predicting brain fMRI measurements in response to images. We report two findings. First, explicit mapping between the brain and deep-network features across dimensions of space, layers, scales, and channels is crucial. This mapping method, Fac- torTopy, is plug-and-play for any deep-network; with it, one can paint a picture of the network onto the brain (liter- ally!). Second, our visualization shows how different train- ing methods matter: they lead to remarkable differences in hierarchical organization and scaling behavior, growing with more data or network capacity. It also provides in- sight into fine-tuning: how pre-trained models change when adapting to small datasets. We found brain-like hierarchi- cally organized network suffer less from catastrophic for- getting after fine-tuned. 1. Introduction The brain is massive, and its enormous size hides within it a mystery: how it efficiently organizes many specialized modules with distributed representation and control. One clue it offers is its feed-forward hierarchical organization (Figure 2). This hierarchical structure facilitates efficient computation, continuous learning, and adaptation to dy- namic tasks. Deep networks are enormous, containing billions of pa- rameters. Performances keep improving with more training data and larger size. It doesn’t seem to matter if the network is trained under the supervision of labels, weakly super- vised with image captions, or even self-supervised without human-provided guidance. Its sheer size also hides another mystery: as its size increases, it can be fine-tuned success- fully to many unseen tasks. *: Equal advising. 1 arXiv:2312.01280v2 [cs.CV] 30 Mar 2024 What can these two massive systems, the brain and deep network, tell about each other? By identifying ‘what’ deep features are most relevant for each brain voxel fMRI predic- tion, we can obtain a picture of deep features mapped onto a brain (literally), as shown by the brain-to-network mapping in Figure 1. The key insight is that deep networks trained with the same architecture, but different objectives and data, pro- duce drastically different computation layouts of interme- diate layers, even if they can produce similar brain en- coding scores and other downstream task scores. For ex- ample, we found intermediate layers of CLIP align hier- archically to the visual brain. However, there are unex- pected non-hierarchical bottom-up and top-down structure in supervised classification and segmentation-trained mod- els. Moreover, for many models, when scaling up in param- eters and training data, they tend to lose hierarchical align- ment to the brain, except CLIP, which improved hierarchical alignment to the brain after scaling up. Suppose the brain’s hierarchical organization is a tem- plate for efficient, modular, and generalizable computation; an ideal computer vision model should align with the brain: the first layer of the deep network matches the early visual cortex, and the last layer best matches high-level regions. Our fine-tuning results show that networks with more hi- erarchy organization tend to (qualitatively) maintain their hidden layers better after fine-tuning on small datasets, thus suffering less (quantitatively) from catastrophic forgetting. We conjecture that better alignment to the brain is one way to find a robust model that adapts to dynamic tasks and scales better with larger models and more data. Our analysis crucially depends on a robust mapping be- tween deep 4D features: spatial, layer, channel, and scale (class token vs local token) to the brain. Our fundamen- tal assumption is that this mapping should be: a) brain- topology constrained, and b) factorized in feature dimen- sions of space, layer, channel, and scale. This is important because independent 4D image features to brain mapping are highly unconstrained, and learning a shared mapping across images, with brain-topology constraint and factor- ized representation, is statistically more stable. Our contribution is summarized as the following: 1. We introduce a factorized, brain-topological smooth se- lection that produces an explicit mapping between deep features: space, layer, channel, and scale (class token vs local token) to the brain. 2. We pioneer a new network visualization by coloring the brain using layer-selectors, exposing the inner workings of the network. 3. We found that brain-like hierarchically organized net- works suffer less from catastrophic forgetting after fine- tuning. 1: The Algonauts 2023 competition: http://algonauts.csail.mit.edu/ 2. Background and Related Work Hierarchy of the Visual Brain In Figure 2, visual brain is organized into regions, each region has specialized func- tions. Image processing in visual brain is organized in a hierarchical and feed-forward fashion. Starting from re- gion V1 to V4, neurons were found to have increasing receptive field size and represent more abstract concepts [14, 15, 65], the late visual brain has semantic regions such as face (FFA), body (EBA), and place (OPA, PPA). INPUT V1 V4 FFA high-level Figure 2. Image features (selected channels) for brain ROIs. V1 is orientation filtering, V4 segmentation, FFA face-selective. Brain Encoding Benchmarks Open challenge and com- petitions on brain encoding model have generated broad interests [9, 10, 19, 51, 52, 59, 62]. Large-scale open- source datasets are growing rapidly in both quantity and quality [2, 7, 18, 24, 31]. The Algonauts1 2023 competition [19] is the first to use a massive high-quality 7-Tesla fMRI dataset [2]. The high-quality and large-scale of this datasets enabled models that can recover brain-to-space mapping from naturalistic image stimuli [47], which was only pos- sible with synthetic stimuli [15]. Our brain encoding model methods is a direct extension of the Algonauts 2023 com- petition winning method Memory Encoding Model [67]. In this work, we added a scale axis for feature selection. Explain Brain by Deep Networks After fitting brain en- coding models to predict brain response, gradient-based methods have been used to explain how brain works: orientation-selective neurons in V1 [16, 37, 47], category- selective regions in late visual brain [28, 34, 35, 42, 45, 50]. Gradient-based methods can also generate maximum- excited images [4, 20, 30, 55, 66]. Meanwhile, studies try to find the best performance pre-trained model for each brain ROI [12, 40, 49, 61, 70] from a zoo of supervised [27, 29, 44, 54], self-supervised [8, 21, 23, 32, 41], image generation [46], and 3D [38, 48, 58] models. Features can be efficiently cached and are plug-in-and-play [22, 57, 60]. Different from the main-stream study that use deep net- works to explain the brain’s functionality. In this work, we use existing knowledge of the brain’s functionality to ex- plain feature computation in deep networks. 2 Figure 3. Topological Constrained, Factorized, Brain-to-Network Selectors for CLIP. Top: factorized-selectors trained with topolog- ical constraints improved confidence of the mapping (color brightness) and mapping smoothness (colored as Section 3.2). Bottom left: individual layer-selector weight ˆωlayer, note layer 4 is mostly aligned with V1, and the last two are aligned with the body (EBA) and face (FFA) region. Bottom right: space-selector ˆuspace: 3D voxels, dots, are mapped to the image space with color dots indicating the layers. For later layers, only center image regions are selected. 3. Methods: Brain Encoding Model Figure 1 presents an overview of our methods. In the brain encoding task, one needs to predict a large number vox- els (vertices), of visual cortex’s fMRI responses as a func- tion of the observed image. This encoding task is under- constrained: since each subject has her/his unique men- tal process, a successful brain encoding model needs to be highly individualized, thus significantly reducing the train- ing example per voxel. Most of the current approaches treat each brain voxel independently. This leads to a major reduc- tion in signal-to-noise ratio, particularly for our analysis. Our fundamental innovations are two-fold. First, we en- force brain-and-network topology-constrained prediction. Brain voxels are not independent but are organized locally into similar “tasks”, and globally into diverse functional re- gions. Similarly, Neural networks show local feature simi- larity across adjacent layers while ensuring diversity for far- away ones. The local smoothness constraints significantly reduce uncertainties in network-to-brain mapping. Second, we propose a factorized feature selection across three independent dimensions of space, layers, and scales (local vs global token). This factorized representation leads to a more robust estimation because feature selection in each dimension is more straightforward, and learning can be more efficient across training samples. For example, the spatial feature selection only needs to find the center of the pixel region for each brain voxel, similar to retinotopy. The layer or scale selection estimates the size of the pixel re- gion: the early layer typically has a smaller receptive field size. Note that the factorized feature selection is soft: multi- ple layers or spatial locations can be selected, as determined by the brain prediction training target. 3.1. Factorized, Topological Smooth, Brain-to- network Selection (FactorTopy) We used a pre-trained image backbone model (ViT) to pro- cess input image X into features V . The entire feature V is organized along four dimensions: space, layer, scale (class t