In this work, we study the influence of group actions on data and how they may impact the architecture and performances of neural networks, especially convolutional neural networks (CNN). CNNs rely on assumed translational symmetries in data and have shown very robust performance in imaging tasks, especially medical imaging ones, and they are highly parameter-efficient due to their weight-sharing property. When data offer more structure than simply translation, this can be used to build generalized CNNs. This is especially the case for the task at hand—classification and segmentation of Diffusion Weighted Imaging (DWI) data. These Group and Geometric CNNs (GCNN) have been studied intensively and applied in many situations in the few past years (Masci et al., 2015; Cohen and Welling, 2016a; Boscaini et al., 2016; Bekkers et al., 2018; Cohen et al., 2020 to cite a few).

DWI is a non-invasive image modality that provides local information about water diffusion in tissues by means of measuring spin displacement (Tuchs, 2004). It provides three-dimensional diffusion information at each location x that can be encoded as a function f_x on the two-dimensional sphere S². A field of these functions, on a given domain, can be represented as a function f:ℝ³×S² → ℝ. If a sample is rotated and translated, the acquired signal should reflect, up to the limitations of acquisition protocol, this transformation. The group in question is the group of 3D rigid motions, SE(3), and the space ℝ³×S² is a homogeneous space under the action of SE(3): A point in ℝ³×S² can be transformed in any other point by a rigid transformation. This notion of homogeneous space is at the heart of the extension of CNNs to GCNNs (Cohen et al., 2020; Bekkers, 2019).

Our task at hand is the classification/segmentation of diffusion data. The inductive bias provided by the knowledge of these transformations may prove important for our task, especially when the amount of annotated data is limited. The problem boils down to how to incorporate this knowledge. The most classical approach is to use data augmentation, reflecting the expected symmetries in the data, in the hope that the network will be able to learn it during the training phase, learning symmetry-aware kernels.

Incorporating, on the other hand, some information about the symmetries of the data in the model has been shown to boost the performances of these networks (Bekkers et al., 2018). But how much of this information is needed for a given task? To provide an answer, for the DWI segmentation task, we propose several networks, which gradually incorporate these symmetries in their architecture and study their performances. In addition, instead of performing convolution on non-Euclidean data in a spectral fashion using Fourier-type transformations, we implement convolution in all our experiments in a direct way, as is usually done in the image analysis community. In other words, we use regular representations of groups to encode the group actions in the models, instead of irreducible representations. Our experiments, in some sense, perform a group action ablation study. We start with a “naive” CNN and then incorporate spherical symmetries, resulting in a SO(3)-GCNN, discarding the spatial aspect of the data. The spatial aspect is then added in the form of a standard CNN coupled with spherical symmetries, and then, we build a network where roto-translational transformations are used in almost all steps. This work demonstrates empirically the improvement in performance. The results are, however, not always clear-cut. Previous works associated with group convolutions have addressed the capabilities of their models in comparison with data augmentation but, to the best of our knowledge, have not touched the comparison between models tested under randomly transformed test sets. This is what our ablation study is providing. It not only shows the impact of embedding transformations in the model but also gives a systematic analysis on the comparison among different group actions and the corresponding elements of network architecture with respect to the interplay between rotations and translations - the physically justified roto-translation group and the simpler direct product of translations and rotations - imposed in the models, and their relation to data augmentation, both in the training and test set. In the study we provide, the GCNN built from 3D-translations on one hand and rotations, on the other hand, seems to perform better than a SE(3)-GCNN. However, the SE(3)-network generalizes better to unseen rotated data than the previous one. The reason may lie in the particular type of data used - our DWI scans come from the Human Connectome Project (HCP) (Van Essen et al., 2013) are highly preprocessed, including a form of alignment – and this may impact the results. Nevertheless, for every model we propose, we also experiment training them with data augmentation to compare with our equivariant networks. We show that the more equivariance we incorporate into the model, the better the model resists the inconsistency of distributions between training and testing data.

The contribution of this work is as follows.

In the rest of this paper, we review related work, both around CNN and DWI classification problem. Then, we introduce the theoretical setup of GCNN and build several networks. Thereafter, we study and discuss their performances. Our implementation and experiments are publicly available at https://github.com/rliu-p/se3gcnn.

Deep Learning (DL) for non-flat data, or using more complex group actions than just translations, is currently getting more attention from the research field. When it comes to non-flat data, such as the point-wise spherical signals in DWI, particularly relevant related works are the following. A non-rotationally invariant modification was proposed by Boscaini et al. (2016). Schnell et al. (2009) developed an Support Vector Machine (SVM) using rotation-invariant features extracted from Spherical Harmonic decomposition of the HARDI signals, while Skibbe and Reisert (2017) introduced a toolkit for 3D image processing based on Spherical Tensor Algebra (STA), which is particularly well-suited for tasks requiring rotational invariance, such as image enhancement, reconstruction, and feature detection.

The above provide methods for DL-based processing of data on arbitrary manifolds. When the manifold, however, is a homogeneous space, i.e., there is a group action by which any two points on the manifolds can be reached, theory simplifies via a natural generalization of classical convolutions in group convolution neural networks (GCNNs), as was presented in Cohen et al. (2018); Bekkers et al. (2018); and Kondor and Trivedi (2018). GCNNs guarantee global equivariance. However, global equivariance can be complicated and elusive when the underlying geometry is non-trivial, which was discussed in Cohen et al. (2019). An elementary construction on a general manifold is proposed by Schonsheck et al. (2018) via a fixed choice of geodesic paths used to transport filters between points on the manifold, ignoring the effects of path dependency, i.e., holonomy when paths are geodesics. The removal of this path dependency can be obtained by summarizing local responses over local orientations, which is what was done by Masci et al. (2015). To explicitly deal with holonomy, Sommer and Bronstein (2020) proposed a theoretical breakthrough using convolution construction on manifolds based on stochastic processes via the frame bundle.

On the other hand, Cohen et al. (2018) lifted spherical functions to the 3D-rotation group SO(3) and used a generalization of Fourier transform on it to perform convolution. Elaldi et al. (2021) proposed an equivariant spherical deconvolution method to learn the orientation distribution function (ODF). Bouza et al. (2021) generalized convolution to manifold-valued convolutions using Volterra Series, preserving its equivariance. With the generalization of convolution to more complex group actions than translation, several authors (Gens and Domingos, 2014; Cohen and Welling, 2016a; Weiler et al., 2018b,a; Worrall et al., 2017; Kondor and Trivedi, 2018; Bekkers et al., 2018; Andrearczyk et al., 2020; Chakraborty et al., 2018a,b, 2020; Graham et al., 2020) explored the group convolution path for Lie groups and the homogeneous spaces of these groups. Knigge et al. (2022) proposed a separable convolution setup on Lie groups. The relation between group actions, principal bundles and related vector bundles, and convolutional architectures is currently explored (Cohen et al., 2019, 2020; Aronsson, 2022). The latter elucidates important relations between differential geometry of bundles and Reproducible Kernel Hilbert Spaces. Links between partial differential equations, symmetries, and GCNN are studied in Smets et al. (2021). A unifying framework for equivariant DL on manifolds, connecting both the bundle and homogeneous space viewpoint, is given in Weiler et al. (2021) through a notion of coordinate indepencent convolutions.

Most CNNs approach for the processing of DWI signals discards its specific structure. For instance, Golkov et al. (2016) built multi-layer perceptrons in q-space for kurtosis and NODDI mappings. However, the importance of spherical equivariant or invariant structure has been acknowledged for some years now. The importance of the extraction of rotationally invariant features beyond Fractional Anisotropy (Basser et al., 1994) has been recognized in series of DWI works. For instance, Caruyer and Verma (2015) developed invariant polynomials of spherical harmonic (SH) expansion coefficients and discussed their application in population studies. Schwab et al. (2013) proposed a related construction using eigenvalue decomposition of SH operators. Novikov et al. (2018) and Zucchelli et al. (2020) argued their usefulness for understanding microstructures in relation to DWI.

Chakraborty et al. (2018a) proposed a rotation equivariant construction inspired by Cohen et al. (2018) for disease classification. The same authors (Banerjee et al., 2019) used a S²×ℝ⁺ CNN using SHORE function representation for classification in Parkinson's Disease. Sedlar et al. (2020) used a spherical U-Net for f-ODF estimation. The same authors (Sedlar et al., 2021) used a spherical CNN for microstructure parameter estimation, using spherical harmonics representations. Müller et al. (2021) proposed a sixth-D, 3D space and q-space NNs with roto-translation/rotation equivalence properties, targeted at DWI data. Poulenard et al. (2022) reviewed several implementations of SE(3) neural networks and showcased a comparison among these networks. In their work, steerable CNNs generalize better than group CNNs while dealing with inconsistent distributions between training and testing data for 3D images.

While most equivariant methods use spectral representation of groups, we propose an SE(3) network for DWI data that uses regular representation of groups such that the whole model is light-weight, and the implementation for convolution is not only direct but also separable, improving efficiency. A similar idea was used in Chen et al. (2021), for 3D point cloud feature extraction, with, however, important architectural differences due to the nature of input data. Both our method and Chen et al. (2021) implemented regular representations of groups in a separable fashion; however, their separable kernels are only over the spatial and rotation interactions while we additionally split the rotation interactions over 2 axes, making use of the factorization of the icosahedron group into 12 × 5 rotations. We do this to further boost efficiency. Furthermore, Chen et al. (2021) include an attention mechanism in the interaction layers, while instead, we use non-linearities between the separate interaction steps. As our operations are strictly local, including attention mechanism would introduce unnecessary computational overhead, whereas in Chen et al. (2021) the attention mechanism could be critical as a selection mechanism among the global interactions between many points within the point cloud. In addition, we compared our method to Müller et al. (2021) which uses steerable filter bases (spectral representation of groups) for the SE(3) group. In our experiments, in comparison with Müller et al. (2021), we found out, however, that our direct convolution implementation of SE(3) GCNN does not perform inferior to its steerable alternative, and our method is a lot more light-weight.

The networks we present will be built from the principle of expanding CNNs to groups and their homogeneous spaces, on which they act by extending convolution operations to functions on groups and their homogeneous spaces. For the rotation group SO(3) and the sphere S² as SO(3)-homogeneous space, the common path for implementing convolutions/correlations is to use irreducible representations (Cohen and Welling, 2016b). This approach can be computationally very intensive, unless one restricts to very low-order irreducible representations, with a resolution trade-off worse than the approximation of SO(3) by the icosahedral rotation group. So we do not follow that path here.

In the next section, we provide the theoretical background for extending convolutions to functions on groups. For the reader's convenience, standard concepts from group theory and group actions that are used to build our new convolution layers are presented in Appendix 1.

In this section, we first list all the detailed network setups, after which we present the results of the experiments. We evaluate our method on the DWI brain dataset from the human connectome project (HCP) (Van Essen et al., 2013). We classify the human brains into four regions - cerebrospinal fluid (CSF), subcortical, white matter (WM), and gray matter (GM). An illustration of the task can be found in Figure 3.

We use the preprocessed DWI data (Van Essen et al., 2013) and normalize each DWI scan for the b-1000 images with the voxel-wise average of the b₀. We use the brain masks provided in the dataset to obtain the voxels of interest, while background is ignored. The labels provided with the T1-image are transformed to the DWI using nearest neighbor interpolation (Figure 3). The resolution of the DWI images is 145 × 174 × 145, and the resolution of the T1-images is 260 × 311 × 260. Focal Loss (Lin et al., 2018) is used to counter the class imbalance of the four brain regions. For Focal Loss, all experiments use γ = 2 and use α = (0.35, 0.35, 0.15, 0.15) for CSF, subcortical, WM, and GM, respectively. For the Watson Kernel, all experiments that used this interpolation (Type 2 discretization) have κ = 10. Batch size for all experiments is 100, and the learning rate for all experiments is 0.001.

The resistance to data variation that has been shown by our fully equivariant network was demonstrated on synthetically augmented data - with 90-degree rotations. Even though this synthetic augmentation did not cost any loss of signals or any interpolation-caused inaccuracy, it is desirable to verify the robustness of more complex group actions in CNNs using data with real-world variations (e.g., subjects scanned in different positions, affine variations in shapes). Acquiring this type of data is another challenge. On the other hand, data augmentation seems to be very robust against the variations in the rotated test set. However, this is because the augmentations applied in the training set and the test set are identical, and they modeled exactly the same distribution in the data. Our proposed equivariant methods deal with inconsistent distributions between the training set and the test set much better, which is usually the case in real world. In addition, our method outperforms (Müller et al., 2021) with the same amount of information given to the models. Even though both methods show similar resistance to variations in the distributions of the training and test set, our model has a more light-weight implementation using regular group representation with separable kernels. Furthermore, the experiments we conducted using Schnell et al. (2009) have shown the power of equivariant learning in non-Euclidean spaces. Using rotation-invariant features as in Schnell et al. (2009) is beneficial in terms of getting consistent features from spherical functions, regardless of the orientation. However, extracting invariant features from the very beginning also discards potentially valuable orientational information that is implicitly embedded in the data, and discarding spatial information completely severely weakens the capability of the model. This is easily shown by the fact that our SO(3) models that also discard spatial relationships outperform (Schnell et al., 2009).

In conclusion, we presented a systematic study of GCNNs of various group actions with the application to DWI segmentation. We interpreted images of DWI scans (I:ℝ³×S² → ℝ) as functions in the homogeneous spaces of groups with different complexities of symmetries and provided a detailed analysis of how different levels of complexities of these symmetries impact the performance of the network. It is shown from the models OursDecoupled and OursFull that whether or not more complex transformations should be imposed in the model is not always a clear-cut, since while tested on the original test set, OursDecoupled has a slightly better performance. OursDecoupled incorporates a mathematically well-defined, but physically impossible group action, yet it is computed more cheaply, while OursFull incorporates the SE(3) action, which corresponds to the expected physical transformations of the data. And, under any physically realistic turbulence in the test data resulting in unseen distributions, adding to the model possible transformations of the data (to the limit of their discretizations) provides a more stable performance. Therefore, we emphasize the importance of imposing the full roto-translation transformations in models as it is the kind that appears in the data. From the experiments, we conclude that (1) exploiting the spatial-directional interactions in the data is crucial for efficient learning of the features; (2) incorporating complex group actions of 3D rigid motions—SE(3)—might not be essential for highly aligned and preprocessed data such as the human connectome project (HCP) (Van Essen et al., 2013), but it shows significantly higher resistance to variations in data. For real-world raw data in which the positions of subjects are not perfectly aligned as in Van Essen et al. (2013), our proposal shows significant potential.