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: This paper uses recordings from monkeys to co-train a classifier with a second head from layer conv-3-1 (of VGG-19) that predicts the cortical responses in V1 to the same pictures. The authors test this model on 14 different distortions of images and find that it is more robust than a baseline model trained only to classify the images. They show an upper-bound on robustness via a network trained on distorted images from the same distributions, and show that training on scrambled neural responses provides no advantages.   Finally, they analyze the network in two ways, qualitatively and quantitatively. First they reconstruct several input images corrupted by gaussian noise from conv-3-1 in their model and the various control networks. They find that the monkey-trained network appears to focus more on the object than the noise compared to a network trained only on the images. This is obviously subjective, so they go on to measure how much the monkey-trained network attends to salient regions (based on Deep-Gaze-2) compared to the other models, and find that it is generally more sensitive to the salient regions.