XEdgeConv: Leveraging graph convolutions for efficient, permutation- and rotation-invariant dense 3D medical image segmentation

Christian Weihsbach, Lasse Hansen, Mattias P Heinrich

Keywords: Group convolutional neural networks, robust segmentation, rotation and permutation equivariance

Abstract: Deep learning-based 3D anatomical segmentation models that employ convolution kernels have become ubiquitous in medical imaging. Currently, there exist trade-offs between model capacity, the complexity of inference and accuracy. To cope with geometric invariances, reflections (axes flips) of input data in training and test-time augmentations are often used, but cause redundancies in computations. Group equivariance is one solution to enforce invariance w.r.t. rotation and reflection, but it comes at the cost of complicated inference. To address those issues, we first explore a simple yet effective method that directly learns symmetric kernels. To further boost performance and achieve full rotational and reflection equivariance, we propose a novel concept that extends the idea of EdgeConvs, that have so far been used in geometric point cloud learning, from graphs into voxelised grids and integrate this into the state-of-the-art framework for medical 3D segmentation, the nnUNet. Our XEdgeConv kernel reduces the parameter count by 93% and computational operations 20-fold while maintaining very high segmentation accuracies on two challenging 3D multi-organ segmentation tasks and it clearly outperforms alternative parameter reduction strategies. https://github.com/multimodallearning/XEdgeConv