Clebsch-Gordan Nets: a Fully Fourier Space Spherical Convolutional Neural Network

Introduction

Despite the many recent breakthroughs in deep learning, we still do not have a satisfactory understanding of how deep neural networks are able to achieve such spectacular perfomance on a wide range of learning problems. One thing that is clear, however, is that certain architectures pick up on natural invariances in data, and this is a key component to their success. The classic example is of course Convolutional Neural Networks (CNNs) for image classification . Recall that, fundamentally, each layer of a CNN realizes two simple operations: a linear one consisting of convolving the previous layer’s activations with a (typically small) learnable filter, and a nonlinear but pointwise one, such as a ReLU operatorReal CNNs typically of course have multiple channels, and correspondingly multiple filters per layer, but this does not fundamentally change the network’s invariance properties.. This architecture is sufficient to guarantee translation equivariance, meaning that if the input image is translated by some vector $t$, then the activation pattern in each higher layer of the network will translate by the same amount. Equivariance is crucial to image recognition for two closely related reasons: (a) It guarantees that exactly the same filters are applied to each part the input image regardless of position. (b) Assuming that finally, at the very top of the network, we add some layer that is translation invariant, the entire network will be invariant, ensuring that it can detect any given object equally well regardless of its location.

Recently, a number of papers have appeared that examine equivariance from the theoretical point of view, motivated by the understanding that the natural way to generalize convolutional networks to other types of data will likely lead through generalizing the notion of equivariance itself to other transformation groups . Letting $\smash{f^{s}}$ denote the activations of the neurons in layer $s$ of a hypothetical generalized convolution-like neural network, mathematically, equivariance to a group $G$ means that if the inputs to the network are transformed by some transformation $g\in G$, then $\smash{f^{s}}$ transforms to $T^{s}_{g}(f^{s})$ for some fixed set of linear transformations $\smash{\{T^{s}_{g}\}_{g\in G}}$. (Note that in some contexts this is called “covariance”, the difference between the two words being only one of emphasis.)

In the present paper we propose a spherical CNN that differs from in two fundamental ways:

Rather than a pointwise nonlinearity in real space, our network takes the tensor (Kronecker) product of the activations in each layer followed by decomposing the result into irreducible fragments using the so-called Clebsch–Gordan decomposition. This way, we get a “fully Fourier space” neural network that avoids repeated forward and backward Fourier transforms.

The resulting architecture is not only more flexible and easier to implement than , but our experiments show that it can also perform better on some standard datasets.

Convolutions on the sphere

and then applying a nonlinearity $\sigma$, such as the Re-LU operator:

Defining $\smash{T_{x}(h^{s})(y)=h^{s}(y-x)}$, which is nothing but $h^{s}$ translated by $x$, allows us to equivalently write (1) as

where the inner product is $\langle f^{s-1}\!,T_{x}(h^{s})\rangle=\sum_{y}f^{s-1}(y)\,T_{x}(h^{s})(y).$ What this formula tells us is that fundamentally each layer of the CNN just does pattern matching: $f^{s}(x)$ is an indication of how well the part of $\smash{f^{s-1}}$ around $\smash{x}$ matches the filter $\smash{h^{s}}$.

Equation 3 is the natural starting point for generalizing convolution to the unit sphere, $S^{2}$. An immediate complication that we face, however, is that unlike the plane, $S^{2}$ cannot be discretized by any regular (by which we mean rotation invariant) arrangement of points. A number of authors have addressed this problem in different ways . Instead of following one of these approaches, similarly to recent work on manifold CNNs , in the following we simply treat each $f^{s}$ and the corresponding filter $h^{s}$ as continuous functions on the sphere, $f^{s}(\theta,\phi)$ and $h^{s}(\theta,\phi)$, where $\theta$ and $\phi$ are the polar and azimuthal angles. We allow both these functions to be complex valued, the reason for which will become clear later.

The inner product of two complex valued functions on the surface of the sphere is given by the formula

where ∗ denotes complex conjugation. Further, $h$ (dropping the layer index for clarity) can be moved to any point $(\theta_{0},\phi_{0})$ on $S^{2}$ by taking $\smash{h^{\prime}(\theta,\phi)=h(\theta-\theta_{0},\phi-\phi_{0})}$. This suggests that the generalization of 3 to the sphere should be

Unfortunately, this generalization would be wrong, because it does not take into account that $h$ can also be rotated around a third axis. The correct way to generalize cross-correlations to the sphere is to define $h\star f$ as a function on the rotation group itself, i.e., to set

where $h_{R}$ is $h$ rotated by $R$, expressible as

with $x$ being the point on the sphere at position $(\theta,\phi)$ (c.f. ).

Cohen et al. observe that the double integral in (6) would be extremely inconvenient to compute in a neural network. As mentioned, in the case of the sphere, just finding the right discretizations to represent $f$ and $h$ is already problematic. As an alternative, it is natural to represent both these functions in terms of their spherical harmonic expansions

and similarly for $h$. Similarly to usual Fourier series, in practical scenarios spherical harmonic expansions are computed up to some limiting “frequency” $L$, which depends on the desired resolution.

Remarkably, the above notions of harmonic analysis on the sphere and the rotation group are closely related. In particular, it is possible to show that each Fourier component of the spherical cross correlation (6) that we are interested in computing is given simply by the outer product

Generalized spherical CNNs

Equation (12) describes the behavior of spherical harmonic vectors under rotations, while (15) describes the behavior of Fourier matrices. However, the latter is equivalent to saying that each column of the matrices separately transforms according to (12). One of the key ideas of the present paper is to take this property as the basis for the definition of covariance to rotations in neural nets. Thus we have the following definition.

To fully define our neural network, we need to describe three things: 1. The form of the linear transformations in each layer involving learnable weights, 2. The form of the nonlinearity in each layer, 3. The way that the final output of the network can be reduced to a vector that is rotation invariant, since that is our ultimate goal. The following subsections describe each of these components in turn.

In a covariant neural network architecture, the linear operation of each layer must be covariant. As described in the Introduction, in classical CNNs, convolution automatically satisfies this criterion. In the more general setting of covariance to the action of compact groups, the problem was studied in . The specialization of their result to our case is the following.

2 Covariant nonlinearities: the Clebsch–Gordan transform

Differentiable nonlinearities are essential for the operation of multi-layer neural networks. Formulating covariant nonlinearities in Fourier space, however, is more challenging than formulating the linear operation. This is the reason that most existing group equivariant neural networks perform this operation in “real space”. However, as discussed above, moving back and forth between real space and the Fourier domain comes at a signficiant cost and leads to a range of complications involving quadrature on the transformation group and numerical errors.

Exploiting Lemma 3, the nonlinearity used in our generalized Spherical CNNs consists of computing (17) between all pairs of fragments. In matrix notation,

where $\sqcup$ denotes merging matrices horizontally. Note that this operation increases the size of the activation substantially: the total number of fragments is squared, which can potentially be problematic, and is addressed in the following subsection.

A more unusual feature of the CG nonlinearity is that its essentially quadratic nature. Quadratic nonlinearities are not commonly used in deep neural networks. Nonetheless, our experiments indicate that the CG nonlinearity is effective in the context of learning spherical images. It is also possible to use higher CG powers, although the computational cost will obviously increase.

3 Limiting the number of channels

4 Final invariant layer

Note that in contrast to other architectures such as that involve repeated forward and backward transforms, thanks to their fully Fourier nature, for Clebsch–Gordan nets, in both training and testing, the elements of $\smash{F^{S}_{0}}$ are guaranteed to be invariant to rotations of arbitrary magnitude not just approximately, but in the exact sense, up to limitations of finite precision arithmetic. This is a major advantage of Clebsch–Gordan networks compared to other covariant architectures.

5 Summary of algorithm

In summary, our Spherical Clebsch–Gordan network is an $S+1$ layer feed-forward neural network in which apart from the initial spherical harmonic transform, every other operation is a simple matrix operation. In the forward pass, the network operates as follows.

Therefore, the type of $\widehat{f}^{0}$ is $\tau^{0}=(n_{\textrm{in}},n_{\textrm{in}},\ldots,n_{\textrm{in}})$, and $\widehat{f}^{0}$ is stored as a collection of $L+1$ matrices $\{F^{0}_{0},F^{0}_{1},\ldots,F^{0}_{L}\}$ of sizes $1\times n_{\textrm{in}}$, $3\times n_{\textrm{in}}$, $5\times n_{\textrm{in}}$,…

For layers $s=1,2,\ldots,S-1$, the Fourier space activation $\widehat{f}^{s}=(F^{s}_{0},F^{s}_{1},\ldots,F^{s}_{L})$ is computed as follows:

Experiments

In this section we describe experiments that give a direct comparison with those reported by Cohen et al. . We choose these experiments as the Spherical CNN proposed in is the only direct competition to our method. Besides, the comparison is also instructive for two different reasons: Firstly, while the procedure used in is exactly equivariant in the discrete case, for the continuous case they use a discretization which causes their network to partially lose equivariance with changing bandwidth and depth, whereas our method is always equivariant in the exact sense. Secondly, owing to the nature of their architecture and discretization, use a more traditional non-linearity i.e. the ReLU, which is also quite powerful. In our case, to maintain full covariance and to avoid the quadrature, we use an unconventional quadratic non-linearity in fourier space. Because of these two differences, the experiments will hopefully demonstrate the advantages of avoiding the quadrature and maintaining full equivariance despite using a purportedly weaker nonlinearity.

We use a version of MNIST in which the images are painted onto a sphere and use two instances as in : One in which the digits are projected onto the northern hemisphere and another in which the digits are projected on the sphere and are also randomly rotated.

The baseline model is a classical CNN with 5 $\times$ 5 filters and 32, 64, 10 channels with a stride of 3 in each layer (roughly 68K parameters). This CNN is trained by mapping the digits from the sphere back onto the plane, resulting in nonlinear distortions. The second model we use to compare to is the Spherical CNN proposed in . For this method, we use the same architecture as reported by the authors i.e. having layers $S^{2}$ convolution – ReLU – $SO(3)$ convolution – ReLU – Fully connected layer with bandwidths 30, 10 and 6, and the number of channels being 20, 40 and 10 (resulting in a total of 58K parameters).

For our method we use the following architecture: We set the bandlimit $L_{max}=10$, and keep $\tau_{l}=\lceil\frac{12}{\sqrt{2l+1}}\rceil$, using a total of 5 layers as described in section 3.5, followed by a fully connected layer of size 256 by 10. We use a variant of batch normalization that preserves covariance in the Fourier layers. This method takes a expanding average of the standard deviation for a particular fragment for all examples seen during training till then and divide the fragment by it (in testing, use the average from training); the parameter corresponding to the mean in usual batch normalization is kept to be zero as anything else will break covariance. Finally, we concatenate the output of each $F_{0}^{s}$ in each internal layer (length 24 each, as each is $\tau_{0}=12$ complex numbers), as well as the original coefficient at $l=0$ (length 2), into a $SO(3)$ invariant vector of length 122. (We observed that having these skip connections was crucial to facilitate smooth training.) After that, we use the usual batch normalization on the concatenated results before feeding it into fully connected layers of length 256, a dropout layer with dropout probability 0.5, and finally a linear layer to to 10 output nodes. The total number of parameters was 285772, the network was trained by using the ADAM optimization procedure with a batch size of 100 and a learning rate of $5\times 10^{-4}$. We also used L2 weight decay of $1\times 10^{-5}$ on the trainable parameters.

We report three sets of experiments: For the first set both the training and test sets were not rotated (denoted NR/NR), for the second, the training set was not rotated while the test was randomly rotated (NR/R) and finally when both the training and test sets were rotated (denoted R/R).

We observe that the baseline model’s performance deteriorates in the three cases, more or less reducing to random chance in the R/R case. While our results are better than those reported in , they also have another characteristic: they remain roughly the same in the three regimes, while those of slightly worsen. We think this might be a result of the loss of equivariance in their method.

Atomization Energy Prediction

Our spherical CNN architecture has the same parameters and hyperparameters as in the previous subsection except that $\tau_{l}=15$ for all layers, increasing the number of parameters to 1.1 M. Following , we share weights amongst atoms and each molecule is represented as a $N\times F$ tensor where $F$ represents $\smash{F^{s}_{0}}$ scalars concatenated together. Finally, we use the approach proposed in to ensure permutation invariance. The feature vector for each atom is projected onto 150 dimensions using a MLP. These embeddings are summed over atoms, and then the regression target is trained using another MLP having 50 hidden units. Both of these MLPs are jointly trained. The final results are presented below, which show that our method outperforms the Spherical CNN of Cohen et al.. The only method that delivers better performance is a MLP trained on randomly permuted Coulomb matrices , and as point out, this method is unlikely to scale to large molecules as it needs a large sample of random permutations, which grows rapidly with $N$.

D Shape Recognition

Finally, we report results for shape classification using the SHREC17 dataset , which is a subset of the larger ShapeNet dataset having roughly 51300 3D models spread over 55 categories. It is divided into a 70/10/20 split for train/validation/test. Two versions of this dataset are available: A regular version in which the objects are consistently aligned and another where the 3D models are perturbed by random rotations. Following we focus on the latter version, as well as represent each 3D mesh as a spherical signal by using a ray casting scheme. For each point on the sphere, a ray towards the origin is sent which collects the ray length, cosine and sine of the surface angle. In addition to this, ray casting for the convex hull of the mesh gives additional information, resulting in 6 channels. The spherical signal is discretized using the Discroll-Healy grid with a bandwidth of 128. We use the code provided by for generating this representation.

We use a ResNet style architecture, but with the difference that the full input is not fed back but rather different frequency parts of it. We consider $L_{max}=14$, and first train a block only till $L=8$ using $\tau_{l}=10$ using 3 layers. The next block consists of concatenating the fragments obtained from the previous block and training for two layers till $L=10$, repeating this process till $L_{max}$ is reached. These later blocks use $\tau_{l}=8$. As earlier, we concatenate the $\smash{F^{s}_{0}}$ scalars from each block to form the final output layer, which is connected to 55 nodes forming a fully connected layer. We use Batch Normalization in the final layer, and the normalization discussed in 4 in the Fourier layers. The model was trained with ADAM using a batch size of 100 and a learning rate of $5\times 10^{-4}$, using L2 weight decay of 0.0005 for regularization. We compare our results to some of the top performing models on SHREC (which use architectures specialized to the task) as well as the model of Cohen et al.. Our method, like the model of Cohen et al. is task agnostic and uses the same representation. Despite this, it is able to consistently come second or third in the competition, showing that it affords an efficient method to learn from spherical signals.

Conclusion

References