A ROTATION - EQUIVARIANT CONVOLUTIONAL NEURAL NETWORK MODEL OF PRIMARY VISUAL CORTEX
A BSTRACT
Classical models describe primary visual cortex (V1) as a ﬁlter bank of orientationselective linear-nonlinear (LN) or energy models, but these models fail to predict neural responses to natural stimuli accurately.
Recent work shows that convolutional neural networks (CNNs) can be trained to predict V1 activity more accurately, but it remains unclear which features are extracted by V1 neurons beyond orientation selectivity and phase invariance.
Here we work towards systematically studying V1 computations by categorizing neurons into groups that perform similar computations.
We present a framework for identifying common features independent of individual neurons’ orientation selectivity by using a rotation-equivariant convolutional neural network, which automatically extracts every feature at multiple different orientations.
We ﬁt this rotation-equivariant CNN to responses of a population of 6000 neurons to natural images recorded in mouse primary visual cortex using two-photon imaging.
We show that our rotation-equivariant network outperforms a regular CNN with the same number of feature maps and reveals a number of common features, which are shared by many V1 neurons and are pooled sparsely to predict neural activity.
Our ﬁndings are a ﬁrst step towards a powerful new tool to study the nonlinear functional organization of visual cortex.
I NTRODUCTION
The mammalian retina processes image information using a number of distinct parallel channels consisting of functionally, anatomically, and transcriptomically deﬁned distinct cell types.
In the mouse, there are 14 types of bipolar cells (Euler et al., 2014), which provide input to 30–50 types of ganglion cells (Baden et al., 2016; Sanes & Masland, 2015).
In visual cortex, in contrast, it is currently unknown whether excitatory neurons are similarly organized into functionally distinct cell types.
A functional classiﬁcation of V1 neurons would greatly facilitate understanding its computations just like it has for the retina, because we could focus our efforts on identifying the function of a small number of cell types instead of characterizing thousands of anonymous neurons.
Recent work proposed a framework for learning functional cell types from data in an unsupervised fashion while optimizing predictive performance of a model that employs a common feature space shared among many neurons (Klindt et al., 2017).
The key insight in this work is that all neurons that perform the same computation but have their receptive ﬁelds at different locations, can be represented by a feature map in a convolutional network.
Unfortunately, this approach cannot be applied directly to neocortical areas.
Neurons in area V1 extract local oriented features such as edges at different orientations, and most image features can appear at arbitrary orientations – just like they can appear at arbitrary locations.
Thus, to deﬁne functional cell types in V1, we would like to treat orientation as a nuisance parameter (like receptive ﬁeld location) and learn features independent of orientation.
In the present paper, we work towards this goal.
While we do not answer the biological question whether there are indeed well-deﬁned clusters of functional cell types in V1, we provide the technical foundation by extending the work of Klindt and colleagues (Klindt et al., 2017) and introducing a rotation-equivariant convolutional neural network model of V1.
We train this model directly on the responses of 6000 mouse V1 neurons to learn a shared feature space, whose features are independent of orientation.
We show that this model outperforms state-of-the-art CNNs for system identiﬁcation and allows predicting V1 responses of thousands of neurons with only 16 learned features.
Moreover, for most neurons, pooling from only a small number of features is sufﬁcient for accurate predictions.
R ELATED WORK
Functional classiﬁcation of cell types.
Characterizing neurons according to some identiﬁed, potentially nonlinear response properties has a long history.
For instance, researchers have identiﬁed simple and complex cells (Hubel & Wiesel, 1968), pattern- and component-selective cells (Gizzi et al., 1990), end-stopping cells (Schiller et al., 1976), gain control mechanisms (Carandini et al., 1997) and many more.
However, these properties have been identiﬁed using simple stimuli and the models that have been built for them do not apply to natural stimuli in a straightforward way.
In addition, most of these studies were done using single-cell electrophysiology, which does not sample cells in an unbiased fashion.
Thus, we currently do not know how important certain known features of V1 computation are under natural stimulation conditions and what fraction of neurons express them.
Recent work using two-photon imaging in the monkey (Tang et al., 2018) started closing this gap by systematically investigating response properties of V1 neurons to an array of complex stimuli and found, for instance, that about half of the neurons in V1 are selective to complex patterns such as curvature, corners and junctions rather than just oriented line segments.
Recent work in the retina showed that different types of Ganglion cells (Baden et al., 2016) – and to some extend also bipolar cells (Franke et al., 2017) – can be identiﬁed based on functional properties.
Thus, in the retina there is a relatively clear correspondence of anatomical and genetic cell types and their functional output, and we are getting closer to understanding each retinal cell type’s function.
Learning shared feature spaces for neural populations.
Antolík and colleagues pioneered the idea of learning a shared nonlinear feature representation for a large population of neurons (Antolík et al., 2016), which others have also used in both retina and V1 (Batty et al., 2017; Kindel et al., 2017; Cadena et al., 2019).
Klindt et al.
(2017) proposed a framework to learn functional cell types in an unsupervised fashion as a by-product of performing system identiﬁcation.
They propose a structurally constrained readout layer, which provides a compact characterization of each neuron’s computation.
By enforcing this representation to be sparse, they suggest that the learned features may correspond to distinct functional cell types.
Our present work builds upon this idea and extends it to be applicable to V1.
Rotation equivariance in CNNs.
There is a large body of work on equivariant representations (Nordberg & Granlund, 1996).
Here we review only the most closely related approches.
Cohen &
Welling (2016) introduce group-equivariant CNNs, which include rotation equivariance as a special case and form the basis of our approach.
Weiler et al.
(2017) use a steerable ﬁlter basis to learn rotation-equivariant CNNs.
We essentially use the same approach, but with a different set of steerable ﬁlters (2d Hermite functions instead of circular harmonics).
RotEqNet (Marcos et al., 2016) is related to our work in the sense that it also applies multiple rotated versions of each ﬁlter to the input.
However, instead of maintaining all resulting feature maps as inputs for the next layer, they apply an orientation pooling operation, which reduces each set of feature maps to a two-dimensional vector ﬁeld.
Harmonic networks (Worrall et al., 2016) are an alternative approach that achieves full 360 ◦ rotation equivariance by limiting the structure of the convolutional ﬁlters.
Finally, Dieleman et al.
(2016) achieve rotation equivariance by feeding multiple rotated versions of the input image into parallel CNN streams whose weights are tied together. [Figure] Rank: ... [Figure]
u t p u t f e a t u r e # 1 : 8 x
[Figure]
R OTATION - EQUIVARIANT CNN
Network architecture.
Our architecture follows that of Klindt et al.
The image is ﬁrst processed by a multi-layer convolutional core shared by all neurons (Fig.
1A, blue), which we describe in detail below.
The resulting feature representation is then turned into a response prediction for each neuron by applying a sparse linear combination across space, followed by a sparse linear combination across features and a pointwise, soft-thresholding nonlinearity (Fig.
1A, red).
The entire model is trained end-to-end directly on predicting neural responses, without employing transfer learning or pre-training on auxiliary tasks (Yamins et al., 2014; Cadena et al., 2019).
This model has a number of desirable properties in terms of data efﬁciency and interpretability.
The separation into convolutional core and readout, together with the strong structural constraints on the readout pushes all the ‘heavy lifting’ into the core while the readout weights (spatial mask and feature weights) provide a relatively low-dimensional characterization of each neuron’s function.
Because the core is shared among all neurons (here: thousands), many of which implement similar functions, we can learn complex non-linear functions very accurately.
Equivariance.
A function f : X → Y is called equivariant with respect to a group of transformations Π if such transformations of the input lead to predictable transformations of the output.
Formally, for every π ∈ Π there is a transformation ψ ∈ Ψ such that for every x ∈ X
[Equation]
CNNs are shift-equivariant by construction, meaning that every translation of the image leads to a matching translation in the feature maps (i. e. π and ψ are both translation by a number of pixels horizontally and vertically).
Shift equivariance is a useful property for neural system identiﬁcation, because it allows us to represent many neurons that perform similar computations but in different locations in space by a single convolutional feature map instead of learning each neurons’ nonlinear input-output function ‘from scratch.’
Rotation-equivariant CNN.
Neurons in V1 do not only perform similar functions at different locations, but also extract similar features with different orientations.
Thus, for modeling populations of V1 neurons, learning a representation that is equivariant to rotation in addition to translation would be desirable.
To achieve rotation equivariance, we use group convolutions (Cohen & Welling, 2016).
Here we use the group of all rotations by multiples of 45 ◦ .
That is, for each convolutional ﬁlter in the ﬁrst layer, we have eight rotated copies, each of which produces one feature map (Fig.
1A, blue).
Thus, if we learn 16 different ﬁlters in the ﬁrst layer, it will have a total of 8 × 16 = 128 feature maps.
Formally, if π is a rotation of the image by, say, 45 ◦ , then ψ is a rotation of the feature maps by also 45 ◦ combined with a cyclic permutation of the feature maps by one rotation step.
For the second (and all subsequent) layers, the procedure becomes a bit more involved.
Sticking with the numbers from above, for every feature in the second layer we now learn 128 ﬁlters – one for each input feature map (Fig.
1B, ﬁrst row).
To preserve rotation equivariance, we need to create all rotated copies of these ﬁlters and cyclically permute the feature maps such that each rotated ﬁlter receives the rotated version of the input (depicted by the colored boxed in Fig.
1B, second and following rows).
To implement weight sharing across ﬁlter orientation without aliasing artifacts, we represent the ﬁlters in a steerable basis (Weiler et al., 2017).
We use the two-dimensional Hermite functions in polar coordinates, which form a steerable, orthonormal basis (Fig.
1C; see also Victor et al., 2006; Hu & Victor, 2016).
For ﬁlters of size k , we use all 2d Hermite functions up to rank k , which means we have k ( k + 1) / 2 basis functions.
We sample each ﬁlter at twice the resolution and then downsample by 2 × 2 average pooling to reduce aliasing.
E XPERIMENTS
Neural data.
We recorded the responses of 6005 excitatory neurons in primary visual cortex (layers 2/3 and 4) from one mouse by taking two consecutive scans with a large-ﬁeld-of-view twophoton mesoscope (Sofroniew et al., 2016).
Activity was measured using the genetically encoded calcium indicator GCaMP6f.
V1 was targeted based on anatomical location as veriﬁed by numerous previous experiments performing retinotopic mapping using intrinsic imaging.
We selected cells based on a classiﬁer for somata on the segmented cell masks and deconvolved their ﬂuorescence traces (Pnevmatikakis et al., 2016).
We did not ﬁlter cells according to visual responsiveness.
The aquisition frame rate was 4.8 Hz in both scans.
We monitored pupil position, dilation, and absolute running speed of the animal.
However, because eye movements were rare and we are primarily interested in the average response of neurons given the visual stimulus, we did not further take into account eye position and running speed.
Visual stimuli.
Stimuli consisted of 5500 images taken from ImageNet (Russakovsky et al., 2015), cropped to ﬁt a 16:9 monitor, and converted to gray-scale.
The screen was 55 × 31 cm at a distance of 15 cm, covering roughly 120 ◦ × 90 ◦ .
In each scan, we showed 5400 of these images once (training and validation set) and the remaining 100 images 20 times each (test set).
Each image was presented for 500ms followed by a blank screen lasting between 500ms and 1s.
For each neuron, we extract the accumulated activity between 50ms and 550ms after stimulus onset using a Hamming window.
Tools.
We performed our model ﬁtting and analyses using DataJoint (Yatsenko et al., 2015),
Numpy/Scipy (Walt et al., 2011), Matplotlib (Hunter, 2007), Seaborn (Waskom et al., 2017), Jupyter (Kluyver et al., 2016), Tensorﬂow (Abadi et al., 2015), and Docker (Merkel, 2014).
Preprocessing.
We rescaled the images to 64 × 36 pixels and standardized them by subtracting the mean over all images in the training set and all pixels from each image and dividing by the standard deviation (also taken over all images in the training set and all pixels).
We divided the responses of each neuron by its standard deviation over time.
We did not center the neural responses, because they are non-negative after deconvolution and zero has a clear meaning.
Model ﬁtting and evaluation.
We initialized all weights randomly from a truncated normal distribution with mean zero and standard deviation 0.01.
The biases of the batch normalization layers were initially set to zero.
In contrast, we set the biases in each neuron’s readout to a non-zero initial value, since the neural responses are not centered on zero.
We initialized these biases such that the initial model prediction was on average half the average response of each neuron.
To ﬁt the models, we used the Adam Optimizer (Kingma & Ba, 2014) with an initial learning rate of 0.002, a single learning rate decay and early stopping.
We monitored the validation loss every 50 iterations and decreased the learning rate once by a factor of 10 when the validation loss had not decreased for ﬁve validation steps in a row.
We then further optimized the model until the same criterion is reached again.
To evaluate the models, we use the following procedure.
For each neuron we compute the Pearson correlation coefﬁcient between the model prediction and the average response over the 20 repetitions of the 100 test images.
We then average the correlations over all neurons.
This approach tells us how well the model predicts the average response of neurons to a given stimulus, ignoring trial-to-trial variability (which is interesting in itself, but not the focus of the present work).
Architecture details.
Our architecture consists of three convolutional layers with ﬁlter sizes of 13, 5 and 5 pixels, respectively.
Thus, the receptive ﬁelds of the CNN’s last layer’s units were 21 px, corresponding to ∼ 60 ◦ and covering both classical and extra-classical receptive ﬁeld.
We use zero padding such that the feature maps maintain the same size across layers.
We use 16 ﬁlter sets (i. e.feature maps) in the ﬁrst two layers and 8 to 48 ﬁlter sets in the third layer (number cross-validated, see below).
After each layer, we apply batch normalization followed by a learned scale and bias term for each feature map.
1 After the ﬁrst and second layer, but not after the third, we apply a soft-thresholding nonlinearity f ( x ) = log(
+ exp( x )) .
The feature maps of the third layer then provide the input for each neuron’s readout, which consists of a linear combination ﬁrst over space and then over features, followed by an added bias term and a ﬁnal soft-thresholding nonlinearity.
Thus, each neuron implements a cascade of three LN operations.
Regularization.
We use the same three regularizers as (Klindt et al., 2017).
For smoothness of convolution kernels we set the relative weight of the ﬁrst layer to twice as strong as in the second and third layer to account for the larger kernels.
We apply group sparsity to the convolution kernels of the second and third layer.
We regularize the spatial masks and feature weights in the readout to be sparse by applying the L1 penalty on the 3d tensor that results from taking their outer tensor product.
The weights of all three regularizers are cross-validated as described in the next paragraph.
Model selection.
We cross-validated over the number of ﬁlter sets in the third layer, ranging from 8 to 48 (i. e. 64 to 384 feature maps 2 ) and the strength of the regularizers.
For each architecture, we ﬁt 32 models with different initial conditions and randomly drawn hyperparameters (smoothness of ﬁlters: 0.001–0.03, group sparsity of ﬁlters: 0.001–0.1, sparsity of readout: 0.005–0.03) and chose the best one according to its loss on the validation set (i. e. not using the test set).
For all baseline and control models (see below), we also cross-validated over 32 randomly sampled sets of hyperparameters drawn from the same range of values.
Baseline and control experiments.
As a baseline, we ﬁt a number of regular CNNs without rotation equivariance.
These models are completely identical in terms of number of layers, ﬁlter sizes, readout layer and ﬁtting procedure, except for the weight sharing constraint across orientation.
We ﬁt models with the same number of feature maps as their rotation-equivariant counterparts as well as smaller models with fewer feature maps (see Table 1).
Previous work has enforced the feature weights in the readout to be positive (Klindt et al., 2017).
Because we do not enforce such constraint in the present work, we ran a control experiment, in which we enforce the readout weights (both masks and feature weights) to be positive.
Sparsity of the spatial masks is well justiﬁed, because we know that receptive ﬁelds are localized.
However, it is unclear whether sparsity on the feature weights is desirable, since the function each L [Figure] ss o n v
[Figure]
l i d a t i o n s
[Figure]
t
[Figure]
[Figure]
[Table]
neuron implements may not be well aligned with the coordinate axes spanned by the small number of features we learn.
Thus, we also ﬁt a model without the sparsity regularizer for the feature weights.
R ESULTS
Architecture search and model selection.
We start by evaluating the predictive performance of our rotation-equivariant CNN on a dataset of 6005 neurons from mouse primary visual cortex.
First, we performed an initial architecture search, drawing inspiration from earlier work (Klindt et al., 2017) and exploring different numbers of layers and feature maps and different ﬁlter sizes.
We settled on an architecture with three convolutional layers with ﬁlter sizes of 13, 5 and 5 pixels, respectively, and used 16 ﬁlter sets (i. e. 128 feature maps) in the ﬁrst two layers.
For the third layer, which provides the nonlinear feature space from which all neurons pool linearly, we cross-validated over the number of features, ranging from 8 to 48 (i. e. 64 to 384 feature maps).
Our model achieved an average correlation on the test set of 0 . 47 .
The performance improved with the number of features in the third layer, but the improvement was small beyond 16 features (Fig. 2).
For the following analyses we therefore focus on this model with 16 features as a compromise between model complexity and performance.
It is encouraging that a model with only 16 features is sufﬁcient to accurately model a population as large as 6000 neurons.
Earlier work modeling V1 responses used a shared feature space of 10– (Antolík et al., 2016) or 48 features (Klindt et al., 2017) for populations of 50–100 neurons, which reduced the dimensionality of the feature space by a factor of 2–5 compared to the number of neurons.
Here, we reduce the dimensionality by a factor of 375 while achieving a similar level of performance.
Rotation-equivariant CNN outperforms regular CNN.
To compare our new model to existing approaches, we also ﬁt a regular CNN with identical architecture except for the weight sharing constraint across orientations.
In addition, we ﬁt a number of smaller CNNs with fewer feature maps in each layer, which are more similar in terms of number of parameters, but potentially have less expressive power.
Our rotationequivariant CNN outperforms all baselines (Table 1) and generally requires less data (Fig. 3), showing that enforcing weight sharing across orientations is not only potentially useful for interpreting the model (as we show below), but also serves as a good regularizer to ﬁt a larger, more expressive model. A [Figure]
u
[Figure]
Feature space generalizes to unseen neurons.
To show that our network learns common features of V1 neurons, we performed cross-validation on a held-out set of neurons.
Speciﬁcally, we excluded 20% of the neurons when ﬁtting the network.
We then ﬁxed the rotation-equivariant convolutional core and trained only the readout (spatial mask and feature weights) for those remaining 20% of the neurons.
The resulting test correlation for these neurons was indistinguishable (0.46 vs. 0.47) from the full model, showing that the learned features transfer to neurons not used to train these features.
Feature weights are sparse.
The intuition behind the sparse, factorized readout layer is that the spatial mask encodes the receptive ﬁeld location of each neuron while the feature weights parameterize the neuron’s nonlinear computation.
We now ask how the neurons are organized within this 16dimensional function space.
On the one extreme of the spectrum, each neuron could pick a random direction, in which case sparsity would not be the right prior and the feature weights should be dense.
On the other end of the spectrum, there could be a discrete number of functional cell types.
In this case, each cell type would be represented by a single feature and the feature weights should be maximally sparse, i. e. one-hot encodings of the cell type.
[Figure]
majority of weights is zero; the central bin has been excluded for clarity.
To analyze the feature weights, we ﬁrst marginalize them over orientation.
To do so, we take the sum of squares over allorientations for each of the 16 features and then normalize them such that the energy of the 16 features sums to one.
We ﬁnd that the feature weights are indeed quite sparse (Fig.
5A): most of the neurons use only 1– features (Fig. 5B) and the strongest feature captures more than 50% of the energy for 63% of the neurons.
To ensure that the sparsity of the feature weights is a property of the data and not a trivial result of our L 1 penalty, we performed an ablation study, where we ﬁt a 40° model that applies sparsity only on the spatial masks, but uses L 2 regularization for the feature weights.
This model performed worse than the original model (Table 1) and produced a signiﬁcantly denser weight distribution (Fig. 5B), suggesting that sparsity is indeed a justiﬁed assumption.
Feature weights have consistent sign.
There is a second line of evidence that the features learned by our model are meaningful beyond just providing an arbitrary basis of a 16-dimensional feature space in which neurons are placed at random: the weights that different neurons assign to any given feature have remarkably consistent signs (Fig. 4).
For 11 of 16 features, more than 90% of the non-zero weights have the same sign.
Thus, the negation of one feature that drives a large group of neurons is not a feature that drives many neurons.
Visualization of the different groups of neurons.
Having presented evidence that the features identiﬁed by our model do indeed represent a simple and compact, yet accurate description of a large population of neurons, we now ask what these features are.
As a ﬁrst step, we group all neurons into one of 16 groups based on their strongest feature weight.
This approach is obviously to some extent arbitrary.
By no means to we want to argue that there are exactly 16 types of excitatory neurons in V1 or that excitatory neurons in V1 can even be classiﬁed into distinct functional types.
Nevertheless, we believe that the 16 groups deﬁned in this way are useful in a practical sense, because they represent features that cover more than 50% of the variance of a large fraction of neurons and therefore yield a very compact description of the most important features of V1 computation across the entire excitatory neuronal population.
To visualize what each of these features compute, we pick the 16 most representative examplesof each group and use the model to approximate their linear receptive ﬁeld (RF).
To this end, we compute the model gradient at a gray image (Fig. 6).
The crucial point of this plot is that it shows that the premise of rotation equivariance holds: cells are clustered according to similar spatial patterns, but independent of their preferred orientation.
As expected, most linear RFs resemble Gabor ﬁlters, which show differences in symmetry (odd: #2, #11, #15 vs even: #3, #6, #12) as well polarity (#3 vs.
#6, #12), while some groups exhibit center-surround structure (#4, #14, #16).
D ISCUSSION
We developed a rotation-equivariant convolutional neural network model of V1 that allows us to characterize and study V1 computation independent of orientation preference.
Although the visual system is not equivariant to rotation – there are known biases in the distribution of preferred orientations –, enforcing weight sharing across orientations allowed us to ﬁt larger, more expressive models given a limited dataset.
While our work lays out the technical foundation, we only scratched the surface of the many biological questions that can now be addressed.
Future work will have to investigate the learned features in much more detail, test to what extent they generalize across recording sessions and animals, whether they are consistent across changes in the architecture and – most importantly – whether neurons in V1 indeed cluster into distinct, well-deﬁned functional types and this organization ﬁnds any resemblance in anatomical or genetic properties (Tasic et al., 2017) of the neurons recorded.