Clifford-Steerable Convolutional Neural Networks

In classical physics, space and time are treated as separate entities: space is described using Euclidean geometry, and time is considered a continuous one-dimensional flow. This separation leads to different notions of distance in space and time. In space, to measure how far apart two points are from each other, we use the Euclidean distance: \[ d^2 = \Delta x^2 + \Delta y^2 + \Delta z^2 \] Time, on the other hand, is thought to be absolute: if we were to place two clocks at different locations, they would read the same time, regardless of their spatial separation or relative motion. In other words, the timing of events is assumed to be the same for all observers.

In relativistic physics, the situation is different. There, space and time are unified into a single entity - spacetime. It is described by Minkowski geometry, where the distance between two points is given by: \[ d^2 = c^2 \Delta t^2 - (\Delta x^2 + \Delta y^2 + \Delta z^2) \] where \(c\) is the speed of light. To better build an intuition, let us imagine four clocks moving at different velocities along a rod, and velocities are close to the speed of light. As time is compressed in the direction of motion at high velocities, the clocks have to travel longer distances to show the same time. If we now draw the trajectories of the clocks in spacetime, we would see that they end up on the same hyperbola. This hyperbola represents the set of all spacetime points that are separated from the origin by a fixed proper time interval.

We can now look at the transformations that preserve the distance in pseudo-Euclidean spaces - isometries. For Euclidean spaces, these are rotations, translations, and reflections which constitute the Euclidean group \(E(n)\). For pseudo-Euclidean spaces, the set of isometries includes translations, spatial rotations, reflections, and also boosts between inertial frames. They constitute the pseudo-Euclidean group \(E(p,q)\), which as special cases includes the Poincaré group for Minkowski spacetime \(E(1,3)\) and the Euclidean group \(E(0,n) \equiv E(n) \).

Mathematically, the pseudo-Euclidean group is defined as a semidirect product of the pseudo-orthogonal group \(O(p,q)\) (origin-preserving transformations) and the translation group \( (\mathbb{R}^{p+q}, +) \): \[ E(p,q) := O(p,q) \ltimes (\mathbb{R}^{p+q}, +). \] This means that any isometry can be represented as a composition of a linear transformation from the pseudo-orthogonal group \(O(p,q)\) and a translation.

If we now want to learn a map \(F: f_{\text{in}} \to f_{\text{out}}\), we must respect the geometry of the space. Consequently, the map should respect the transformation laws of the feature fields, which imposes the equivariance constraint: \[ F \circ \rho_\text{in}(g) = \rho_{\text{out}}(g) \circ F. \] This is nothing but a consistency requirement: the function should work the same way regardless of the symmetry transformation applied to the input.

Standard CNNs are already translation equivariant by design, as the same kernel is applied at each location in the input. A lot of work has been done to extend this to the full set of isometries of the Euclidean space, leading to the framework of Steerable CNNs. The core idea here is to use the convolutional operator, as it guarantees translation equivariance, and constrain the space of kernels to be \(O(p,q)\)-equivariant. More formally, a kernel \(k\) should satisfy the following property, called steerability constraint: \[ k(gx) = \rho_{\text{out}}(g) k(x) \rho_{\text{in}}(g)^{-1}. \] The procedure then would be to solve the constraint by hand, obtain the kernel basis, and simply use it in the standard CNN. This approach is general, although quite elaborate, as one needs to solve the constraint for every group of interest \(G\).

We employ Clifford-algebra based neural networks as they allow very easily to build \(O(p,q)\)-equivariant MLPs (see the CEGNN paper for more details). The main advantage is that the framework trivially generalizes to pseudo-orthogonal groups \(O(p,q)\) regardless of the dimension or metric signature of the space. Furthermore, it is possible to have a single implementation of an \(O(p,q)\)-equivariant MLP, which ultimately allows us to build \(E(p,q)\)-equivariant CNNs within a single framework.

The core idea of CEGNNs is to build the nonlinear map by first applying a linear transformation of each grade of the input multivector, and taking the geometric product with the original input. As geometric product is equivariant, the resulting map is \(O(p,q)\)-equivariant. Nonlinearities are efficiently implemented through a gating mechanism: a standard non-linearity (e.g., ReLU) is applied to the scalar part of the multivector, and the result is then multiplied with the original multivector.

The way the implicit kernels works is given below:

1. Given kernel size, discretize \([-1,1]^{p+q}\) to compute the grid as relative positions.
2. Embed the relative positions as vector part of multivectors \(\to\) batch.
3. Compute multivector matrices (vectorized) for each element of the batch.
4. Partially evaluate the geometric product by multiplying the batch of matrices with Cayley table of the algebra.
5. Reshape the output to \((C_{\text{out}} \cdot 2^{p+q},C_{\text{in}} \cdot 2^{p+q}, X_1, ..., X_{p+q})\).

The resulting tensor is ready to be used in any CNN implementation, e.g. torch.nn.ConvNd or jax.lax.conv:

1. Given: input multivector field of shape \((B, C_{\text{in}}, X_1, ..., X_{p+q}, 2^{p+q})\).
2. Compute the kernel as described above.
3. Reshape the input field to \((B, C_{\text{in}} \cdot 2^{p+q}, X_1, ..., X_{p+q})\).
4. Perform the convolution operation \(\to\) output multivector field of shape \((B, C_{\text{out}} \cdot 2^{p+q}, X_1, ..., X_{p+q})\).
5. Reshape the output field to \((B, C_{\text{out}}, X_1, ..., X_{p+q}, 2^{p+q})\).

Theoretically, there is a gap between CS-CNNs and Steerable CNNs in terms of the completeness of the kernel basis. Specifically, CS-CNNs do not cover the full space, and there are two main reasons for that:

1. Multivector representations: multivectors are not able to represent irreps with frequencies higher than 1. This can be illustrated using the (O(3)) group: the vector/bivector part of a multivector is 3-dimensional, whereas irreps with angular frequency 2 are 5-dimensional. It follows that it is not possible neither to represent such irreps (e.g. a tensor) nor to learn a map that involves them. The latter is precisely the reason why the kernel basis of CS-CNNs is not complete - vector to vector mapping requires frequency 2 kernels, which are not present in the basis. This issue is, however, alleviated when applying multiple convolutional layers, which recover the missing degrees of freedom (see Appendix B in the paper for proof).
2. Kernel's input: the kernel's input is limited to the relative position and its norm (shell), which form the 0- and 1-grade components of the input multivector. With this input, the CEGNN backbone is theoretically not able to learn multivectors with non-zero 2-grade. Precisely, the reason for that is that 2-grade appears from the interaction of two unique 1-grades, which is not possible as only one 1-grade is present in the input. There are, of course, ways to overcome this limitation, which we did not explore in this work, but I would be happy to chat about it.

Overall, a lot of work can be done to improve the model, both at the high level (improving expressivity of implicit kernels) and at the low level (developing more expressive basis based on multivectors). Let me know if you are interested in collaboration or just want to chat about the paper :)