Topographic VAEs learn Equivariant Capsules

1 Introduction . Many parts of the brain are organized topographically . Famous examples are the ocular dominance maps and the orientation maps in V1 . What is the advantage of such organization and what can we learn from it to develop better inductive biases for deep neural network architectures ? 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . One potential explanation for the emergence of topographic organization is provided by the principle of redundancy reduction [ 1 ] . In the language of Information Theory , redundancy wastes channel capacity , and thus to represent information as efficiently as possible , the brain may strive to transform the input to a neural code where the activations are statistically maximally independent . In the machine learning literature , this idea resulted in Independent Component Analysis ( ICA ) which linearly transforms the input to a new basis where the activities are independent and sparse [ 2 , 13 , 29 , 44 ] . It was soon realized that there are remaining higher order dependencies ( such as correlation between absolute values ) that can not be transformed away by a linear transformation . For example , along edges of an image , linear-ICA components ( e.g . gabor filters ) still activate in clusters even though the sign of their activity is unpredictable [ 48 , 56 ] . This led to new algorithms that explicitly model these remaining dependencies through a topographic organization of feature activations [ 27 , 45 , 46 , 59 ] . Such topographic features were reminiscent of pinwheel structures observed in V1 , encouraging multiple comparisons with topographic organization in the biological visual system [ 28 , 30 , 42 ] . A second , almost independent body of literature developed the idea of “ equivariance ” of neural network feature maps under symmetry transformations . The idea of equivariance is that symmetry transformations define equivalence classes as the orbits of their transformations , and we wish to maintain this structure in the deeper layers of a neural network . For instance , for images , asserting a rotated image contains the same object for all rotations , the transformation of rotation then defines an orbit where the elements of that orbit can be interpreted as pose or angular orientation . When an image is processed by a neural network , we want features at different orientations to be able to be combined to form new features , but we want to ensure the relative pose information between the features is preserved for all orientations . This has the advantage that the equivalence class of rotations for the complex composite features is guaranteed to be maintained , allowing for the extraction of invariant features , a unified pose , and increased data efficiency . Such ideas are reminiscent of the capsule networks of Hinton et al . [ 21 , 22 , 51 ] , and indeed formal connections to equivariance have been made [ 39 ] . Interestingly , by explicitly building neural networks to be equivariant , we additionally see geometric organization of activations into these equivalence classes , and further , the elements within an equivalence class are seen to exhibit higher-order non-Gaussian dependencies [ 40 , 41 , 56 , 57 ] . The insight of this connection between topographic organization and equivariance hints at a possibility to encourage approximate equivariance from an induced topology in feature space . To build a model , we need to ask what mechanisms could induce topographic organization of observed transformations specifically ? We have argued that removing dependencies between latent variables is a possible mechanism ; however , to obtain the more structured organisation of equivariant capsule representations , the usual approach is to hard-code this structure into the network , or to encourage it through regularization terms [ 4 , 15 ] . To achieve this same structure unsupervised , we propose to incorporate another key inductive bias : “ temporal coherence ” [ 18 , 24 , 52 , 60 ] . The principle of temporal coherence , or “ slowness ” , asserts than when processing correlated sequences , we wish for our representations to change smoothly and slowly over space and time . Thinking of time sequences as symmetry transformations on the input , we desire features undergoing such transformations to be grouped into equivariant capsules . We therefore suggest that encouraging slow feature transformations to take place within a capsule could induce such grouping from sequences alone . In the following sections we will explain the details of our Topographic Variational Autoencoder which lies at the intersection of topographic organization , equivariance , and temporal coherence , thereby learning approximately equivariant capsules from sequence data completely unsupervised . 2 Related Work . The history of statistical models upon which this work builds is vast , including sparse coding [ 44 ] , Independant Component Analysis ( ICA ) [ 2 , 13 , 29 ] , Slow Feature Analysis ( SFA ) [ 54 , 60 ] , and Gaussian scale mixtures [ 41 , 48 , 56 , 57 ] . Most related to this work are topographic generative models including Generative Topographic Maps [ 6 ] , Bubbles [ 25 ] , Topographic ICA [ 27 ] , and the Topographic Product of Student ’ s-t [ 46 , 59 ] . Prior work on learning equivariant and invariant representations is similarly vast and also has a deep relationship with these generative models . Specifically , Independant Subspace Analysis [ 26 , 53 ] , models involving temporal coherence [ 18 , 24 , 52 , 60 ] , and Adaptive Subspace Self Organizing Maps [ 35 ] have all demonstrated the ability to learn invariant feature subspaces and even ‘ disentangle ’ space and time [ 19 , 53 ] . Our work assumes a similar generative model to these works while additionally allowing for efficient estimation of the model through variational inference [ 33 , 50 ] . Although our work is not the first to combine Student ’ s-t distributions and variational inference [ 7 ] , it is the first to provide an efficient method to do so for Topographic Student ’ s-t distributions . Another line of work has focused on constructing neural networks with equivariant representations separate from the framework of generative modeling . Analytically equivariant networks such as Group Equivariant Neural Networks [ 11 ] , and other extensions [ 9 , 16 , 17 , 55 , 49 , 58 , 61 , 62 ] propose to explicitly enforce symmetry to group transformations in neural networks through structured weight sharing . Alternatively , others propose supervised and self-supervised methods for learning equivariance or invariance directly from the data itself [ 4 , 14 , 15 ] . One related example in this category uses a group sparsity regularization term to similarly learn topographic features for the purpose of modeling invariance [ 31 ] . We believe the Topographic Variational Autoencoder presented in this paper is another promising step in the direction of learning approximate equivariance , and may even hint at how such structure could be learned in biological neural networks . Furthermore , the idea of disentangled representations [ 3 ] has also been been connected to equivariance and representation theory in multiple recent papers [ 8 , 12 , 10 , 20 ] . Our work shares a fundamental connection to this distributed operator definition of disentanglement , where the slow roll of capsule activations can be seen as the latent operator . Recently , the authors of [ 34 ] demonstrated that incorporating the principle of ‘ slowness ’ in a variational autoencoder ( VAE ) yields the ability to learn disentangled representations from natural sequences . While similar in motivation , the generative model proposed in [ 34 ] is unrelated to topographic organization and equivariance , and is more aligned with traditional notions of disentanglement . Finally , and importantly , in the neuroscience literature , another popular explanation for topographic organization arises as the solution to the ‘ wiring length ’ minimization problem [ 36 ] . Recently , models which attempt to incorporate wiring length constraints have been shown to yield topographic organization of higher level features , ultimately resembling the ‘ face patches ’ found in primates [ 32 , 38 ] . Interestingly , the model presented in this paper organizes activity based on the same statistical property ( local correlation ) as the wiring length proxies developed in [ 38 ] , but from a generative modeling perspective , demonstrating a computationally principled explanation for the same phenomenon . 3 Background . The model in this paper is a first attempt at bridging two yet disjoint classes of models : Topographic Generative Models , and Equivariant Neural Networks . In this section , we will provide a brief background on these two frameworks . 3.1 Topographic Generative models . Inspired by Topographic ICA , the class of Topographic Generative models can be understood as generative models where the joint distribution over latent variables does not factorize into entirely independent factors , as is commonly done in ICA or VAEs , but instead has a more complex ‘ local ’ correlation structure . The locality is defined by arranging the latent variables into an n-dimensional lattice or grid , and organizing variables such that those which are closer together on this grid have greater correlation of activities than those which are further apart . In the related literature , activations which are nearby in this grid are defined to have higher-order correlation , e.g . correlations of squared activations ( aka ‘ energy ’ ) , asserting that all first order correlations are removed by the initial ICA de-mixing matrix . Such generative models can be seen as hierarchical generative models where there exist higher level independent ‘ variance generating ’ variables V which are combined locally to generate the variances σ = ϕ ( WV ) of the lower level topographic variables T ∼ N ( 0 , σ2I ) , for an appropriate non-linearity ϕ . The variables T are thus independent conditioned on σ . Other related models which can be described under this umbrella include Independent Subspace Analysis ( ISA ) [ 26 ] where all variables within a predefined subspace ( or ‘ capsule ’ ) share a common variance , and ‘ temporally coherent ’ models [ 24 ] where the energy of a given variable between time steps is correlated by extending the topographic neighborhoods over the time dimension [ 25 ] . The topographic latent variable T can additionally be described as an instance of a Gaussian scale mixture ( GSM ) . GSMs have previously been used to model the observed non-Gaussian dependencies between coefficients of steerable wavelet pyramids ( interestingly also equivariant to translation & rotation ) [ 48 , 56 , 57 ] . 3.2 Group Equivariant Neural Networks . Equivariance is the mathematical notion of symmetry for functions . A function is said to be an equivariant map if the the result of transforming the input and then computing the function is the same as first computing the function and then transforming the output . In other words , the function and the transformation commute . Formally , f ( τρ [ x ] ) = Γρ [ f ( x ) ] , where τ and Γ denote the ( potentially different ) operators on the domain and co-domain respectively , but are indexed by the same element ρ . It is well known that convolutional maps in neural networks are translation equivariant , i.e. , given a translation Γρ ( applied to each feature map separately ) and a convolutional map f ( · ) , we have f ( Γρ [ x ] ) = Γρ [ f ( x ) ] . This can be extended to other transformations ( e.g . rotation or mirroring ) using Group convolutions ( G-convolutions ) [ 11 ] . As a result of the design of G-convolutions , feature maps that are related to each other by a rotation of the filter/input are grouped together . Moreover , a rotation of the input results in a transformation ( i.e . a permutation and rotation ) on the activations of each of these groups in the output . Hence , we can think of these equivalence class groups as capsules where transformations of the input only cause structured transformations within a capsule . As we will demonstrate later , this is indeed analogous to the structure of the representation learned by the Topographic VAE with temporal coherence – a transformation of the input yields a cyclic permutation of activations within each capsule . However , due to the approximate learned nature of the equivariant representation , the Topographic VAE does not require the transformations τρ to constitute a group .

Summary: The paper introduces a VAE that can model the spatial and temporal “coherence” of latent variables. The idea is to construct student’s-t distributions from initial Gaussian latents. These student’s-t distributions are made coherent by sharing some of the Gaussian latents that are used in their construction. These variables are then divided into blocks called capsules where the variables in different blocks do not share any Gaussian latent vars. Temporal coherence is achieved by applying a similar “sharing” of Gaussians in a moving time window. However, the paper also suggests “rolling” the variables inside a capsule. The paper’s claim is that this rolling of the variables leads to an equivariant representation. Both qualitative and quantitative experiments on variations of MNIST and dSprites demonstrate the ability of this model in learning representations that are spatially and temporally coherent.