Towards robust vision by multi-task learning on monkey visual cortex

1 Introduction . Although machine learning algorithms have witnessed enormous progress thanks to the recent success of deep learning methods [ 1 ] , current state-of-the-art deep models [ 2–4 ] still fall behind the generalization abilities of biological brains [ 5 , 6 ] . This includes a lack of robustness to image 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . corruptions as pointed out by Hendrycks and Dietterich [ 7 ] , who measured a network ’ s performance on 15 different image corruptions applied to the ImageNet [ 8 ] test-set . Studies with similar image corruptions have proven to severely decrease performance in classification networks while having a smaller impact on human perception [ 9 ] , suggesting that the ability to extrapolate is weak in these networks compared to the mammalian visual system . This gap in extrapolation has previously been attributed to differences in feature representations [ 10 , 11 ] and internal strategies for decision making [ 12 ] between humans and CNNs . Historically , neuroscience has inspired many innovations in artificial intelligence [ 13 , 14 ] , and most of this transfer between neuroscience and machine learning happens on the implementational level [ 15 , 13 ] . However , we currently know too little about the structure of the brain at the level of detail needed to transfer functional generalization properties from biological to artificial systems [ 5 ] . To transfer functional inductive biases from the brain to deep neural networks ( DNNs ) , it may thus be better to consider the representational level by capturing biological feature representations in the responses of biological neurons to visual input – abstracting away from the implementational level . In fact , deep neural networks have set new standards in capturing brain activity across multiple areas in the visual cortex [ 16 ] , becoming the state-of-the-art for neural response prediction of the primary visual cortex ( area V1 , [ 17 ] , but also see Marques et al . [ 18 ] for biologically inspired models that perform similarly well ) . Recent work has shown that CNNs , which were fitted to V1 neural data , can generalize well to other neurons , animals , and stimuli [ 19 , 20 ] . Vice versa , prior work suggests that enforcing brain-like representations in CNNs via neural data from humans [ 21 ] , mice [ 22 ] , or monkeys [ 23 ] can have beneficial effects on the generalization abilities of these networks to stimuli outside their training distribution for object recognition . Our work expands on this line of research by exploring the extrapolation capabilities of multi-task learning models ( MTL ; [ 24 ] ) trained on image classification and prediction of neural responses from monkey V1 . This approach was proposed as the neural co-training hypothesis by Sinz et al . [ 5 ] , but remained untested to the best of our knowledge . We implement MTL via a shared representation between image classification and neural response prediction ( Fig . 1 ) . Our hypothesis is that MTL with neural data regularizes the shared representation to inherit good functional inductive biases from neural data , and improves the network ’ s generalization abilities to out-of-distribution images , thus rendering it more robust . We empirically test this idea using common corruptions on tiny ImageNet ( TIN ) 1 . We show that MTL on monkey V1 has a positive effect on generalization as it increases the model ’ s robustness to image distortions , even though it is trained on undistorted images only . We compare our model with a robust Oracle model to quantify what performance improvement can be expected given that only parts of the network are shared during MTL . Subsequently , we develop a constrained reconstruction based method to analyze learned sensitivities and invariances in the feature representation of the different models . We find that the robust models qualitatively exhibit different feature sensitivities than a standard classification model . Finally , we show that the feature sensitivity of the monkey V1 co-trained model is related to salient image features , consistent with existing theories about the role of V1 [ 25 ] . Overall , our results support the neural co-training hypothesis and further expand the scope of prior results , by exploring the relationship between brain-like representations and robustness . To the best of our knowledge , we are the first to analyze learned feature representations of neurally co-trained models and thus take a step towards a semantic understanding of what makes biological vision robust . 2 Neural multi-task learning . Data The images for the classification task and the neurophysiological experiments were selected from the ImageNet dataset [ 8 ] . For the classification task , we used a grayscale version of TIN1 . The tiny ImageNet dataset ( TIN ) is a subset of ImageNet containing 100000 training images of 200 classes ( 500 for each class ) downsized to images of size 64x64 . In addition , each class has 50 validation images and 50 test images . For neural prediction , we used neurophysiological recordings of 458 neurons from the primary visual cortices ( area V1 ) of two fixating awake macaque monkeys , recorded with a 32-channel depth electrode during 15 ( monkey 1 ) and 17 ( monkey 2 ) sessions . In each session , approximately 1000 trials of 15 images were presented – each image for 120ms . We extracted the spike count from 40ms to 160ms after image onset . The image set presented to the monkey consists of 24075 images from 964 categories – 25 images per category . Of those , 24000 were designated to model training and 75 to testing . For each training trial , a new subset of 15 images was randomly sampled from the training set . Test images were displayed in fixed order in 5 test trials consisting of 15 images each . Each of the test trials were randomly interleaved among training trials and repeated 40-50 times per session . All images were converted to gray-scale and presented at 420×420 px , covering 6.7° visual angle for the monkey , resulting in 63 pixels per degree ( ppd ) . For model training , images were downscaled and cropped to 64×64 pixels , corresponding to 14.0 ppd . Architecture All our experiments were based on a variant of the VGG-19 architecture [ 26 ] with batch normalization layers [ 27 ] after every convolutional layer ( Figure 1 ) . To allow for arbitrary image sizes , we made the network fully convolutional by replacing the fully connected readout by three convolutional layers with dropout of 0.5 after the first two , and a final max-pooling operation and softmax [ 28 ] . We predicted neural responses by feeding the output of the convolutional layer conv-3-1 , shown by Cadena et al . [ 17 ] to be optimal for predicting V1 responses , into a Gaussian readout [ 19 ] yielding a spike count prediction per neuron and image . Models We use a VGG-19 , like we described it above , trained on grayscale TIN to serve as the Baseline for image classification in our experiments . To prepare our neural co-training , similar to Li et al . [ 22 ] , we first trained a Monkey Predictor model on the image-response pairs of our recorded neural data . We then used that model to predict neural responses for all input images of the TIN classification dataset . These predicted responses served as the basis neural dataset we used in our MTL approach . This allowed us to balance the amount of data we have for each task and it removed trial-to-trial variability in the neural data . Since co-training only affects the shared representation up to layer conv-3-1 , we can not expect the network to be as robust as a network where all layers are trained on data augmented with the image distortions . To explore the limits on robustness resulting from sharing lower layers only , we trained a classification model with a 1:1 mixture of clean and distorted images drawn from the pool of 14 ImageNet-C [ 7 ] corruptions ( cf . Figure 9 for examples ) . To push the robustness to the frozen part , we added a second loss that penalizes the Euclidean distance between the outputs of layer conv-3-1 for the same image augmented with different corruptions – similar to Chen et al . [ 29 ] . We then froze all layers up to conv-3-1 , re-initialized the rest , and re-trained the remaining network on clean data only . We refer to this model as the Oracle since it has access to the image distortions during training – unlike our MTL models . To demonstrate that MTL can in principle transfer robustness properties without showing distorted images in training , we generated neural responses from our Oracle model for all images of the clean TIN dataset by freezing the Oracle model and training a Gaussian readout on top of layer conv-3-1 for 10 epochs to predict V1 data . Then , we trained a model on the resulting neural responses alongside clean image classification using MTL . We call this model MTL-Oracle . This model also gives us a realistic “ upper bound ” for our MTL experiments . For our main experiment , we trained MTL with the neural responses generated from the Monkey Predictor model , and refer to it as MTL-Monkey . This model has never seen distorted images at any point during training . To demonstrate that MTL-Monkey has an effect beyond introducing noise into the training , we perform a control experiment which we refer to as MTL-Shuffled . For this , we train a model on the same neural data but with shuffled responses across images for all neurons . An overview of all models used in this study can be found in Table 1 . 1https : //www.kaggle.com/c/tiny-imagenet/overview Training We used cross-entropy loss for single task image classification and Poisson loss for single task neural prediction . For multi-task training , the challenge was to find the optimal balance between the two objectives that achieves a reasonable performance on each task individually , and allows both tasks to benefit from each other by learning common representations . To put both objectives on the same scale , we used their corresponding negative log-likelihood and learned their balance through trainable observation noise parameters σ [ 30 ] . This yields a combined loss of 1 σ2c LCE ( θ , ψc ) + 12σ2nLMSE ( θ , ψn ) + log σc + log σn where θ are the shared parameters and ψc , σc and ψn , σn are the task-specific parameters for classification and neural prediction , respectively . The classification objective LCE was the standard cross-entropy , analogous to the single-task case . For MTL on neural data , we used mean-squared error LMSE because the targets are predictions from the network trained on neural data and not the original noisy neural responses . For optimization , we accumulated the gradients over the different losses to optimize the shared parameters θ in a single combined gradient step . By definition , the two loss components would contribute equally to the learning process . We standardized all pixel values with the mean and standard deviation of the training set , and augmented the images by random cropping , horizontal flipping , and rotations in a range of 15◦ for classification . Optimization was performed using stochastic gradient descent with momentum in all classification-related cases , and Adam for single task neural prediction [ 31 ] . We used a batchsize of 128 and weight decay with a factor of 5 · 10−4 throughout all our experiments . The initial learning rate was determined for each task individually and reduced by a ( task-specific ) factor via an adaptive learning rate schedule . The schedule reduces the learning rate depending on the validation performance – classification performance in the case of MTL – when the rate of improvement is not above 10−4 for 5 consecutive epochs . The training was stopped when we reach either five learning rate reduction steps or a maximum number of epochs , that we defined for each task . This training setup was determined via prior hyper-parameter searches on the validation-set . We repeated every experiment with five different random initializations . Error bars were obtained by bootstrapping ( 250 repetitions ) .

Summary: The paper introduces a new method of regularizing deep networks with neural data. Specifically, the network is co-trained on a main task and regression of neural recordings (or at least deep-network predictions of these recordings for new images). The analysis shows that a network trained on neural responses is sensitive to different features, and encodes gaze-salient image regions more strongly, than control networks.