4D Panoptic Segmentation as Invariant and Equivariant Field Prediction

The task of panoptic segmentation is first proposed in the image domain [23] and later extended to LiDAR point clouds with the release of a large-scale outdoor LiDAR point cloud dataset with panoptic labels, SemanticKITTI [3]. Similar to the semantic segmentation [8, 4, 21, 5, 6] and panoptic segmentation techniques in the image domain [31, 49, 7], their 3D counterparts can be classified into proposal-based and proposal-free methods. Proposal-based methods [3, 38] require a detection module to locate the objects first and then predict the instance mask for each bounding box and conduct semantic segmentation on the background pixels. This strategy needs to deal with potential conflicts among the segmentations. On the other hand, proposal-free methods conduct semantic segmentation first and then cluster the points belonging to different instances. The clustering strategy impacts overall efficiency and performance. Offset prediction and centerness prediction are two major approaches for this. Offset prediction [19, 48] means that each point predicts the offset vector to the instance center, and the clustering is conducted on the predicted centers. Centerness prediction [1] is to regress a heatmap of the closeness to the instance center at any spatial location. Then the local maximums on the heatmap are used to cluster the points nearby. These two strategies can also be combined [52, 26]. Other clustering methods also exist. For example, [18] proposes an end-to-end clustering layer. [34] uses a graph network to cluster over-segmented points into instances.

The 4D task is to provide temporally consistent object IDs on top of 3D panoptic segmentation. MOPT [22] is an early attempt to provide tracking ID for panoptic outputs. 4D-PLS [1] proposes evaluation metrics and a strong baseline method for this task. It accumulates point clouds in sequential timestamps to a common frame, applies segmentation on the aggregated point clouds, and clusters the instances using the centerness prediction. 4D-DS-Net [20] and 4D-StOP [24] follow this pipeline but use offset prediction to cluster the instances. CA-Net [28] uses an off-the-shelf 3D panoptic segmentation network and learns the temporal instances association through contrastive learning.

The equivariance to translations enables CNNs to generalize over translations of image content with much fewer parameters compared with fully connected networks. Equivariant networks extend the symmetries to rotations, reflections, permutations, etc. G-CNN [11] enables equivariance to 90-degree rotations of images. Steerable CNNs extend the symmetry to continuous rotations [43]. The input is also extended from 2D images to spherical [12, 15] and 3D data [42]. To deal with infinite groups (e.g., continuous rotations), generalized Fourier transforms and irreducible representations are adopted to formulate convolutions in the frequency domain [45, 41]. Equivariant graph networks [36] and transformers [17] are also proposed as non-convolutional equivariant layers. Equivariant models have applications in various areas such as physics, chemistry, and bioimaging [41, 17], where symmetries play an important role. They also attract research interest in robotic applications. For example, equivariant networks with SO(3)- and SE(3)-equivariance are developed to process 3D data such as point clouds [14, 10], meshes [13], and voxels [42]. However, due to the added complexity, most equivariant networks for 3D perception are restricted to relatively simple tasks with small-scale inputs, such as object-wise classification, registration, part segmentation, and reconstruction [35, 54, 9]. In the following, we will review recent progress in extending equivariance to large-scale outdoor 3D perception tasks.

LiDAR perception demands two main considerations in equivariant models. Firstly, outdoor scene perception typically needs a sophisticated network design, and incorporating equivariant layers that are compatible with conventional network layers can build on existing successful designs. Secondly, due to the large sizes of LiDAR point clouds, it’s essential to create expressive equivariant models that fit within the memory constraints of standard GPUs.

Existing work mainly follows two strategies. With the first strategy, inputs are projected to rotation invariant features using local reference frames [27, 47]. In this way, the changes are mainly at the first layer of the network, causing limited memory overhead. The main drawback is that the invariant feature could cause information loss and limit performance. The second strategy adopts group convolution, by augmenting the domain of feature maps to include rotations [50, 46]. While achieving improved performance with the help of equivariance, they consume twice [46] or four times [50] of memory as their non-equivariant counterparts due to the augmented dimension of the feature maps and convolutions. With the two strategies, equivariant networks have been applied in 3D object detection [50, 46, 47] and semantic segmentation [27].

A feature map $f_{0}$ is a field, assigning a value to each point in some space. In the context of point cloud perception, we have $f_{0}:\mathbb{R}^{3}\rightarrow V$, where $V$ is some vector space. This map can represent the geometry of the point cloud, i.e., $f_{0}(x)=1$ for every point $x$ in the point cloud and $f_{0}(x)=0$ otherwise. It can also represent point properties or arbitrary learned point features. For example, the feature map $f_{0}:x\mapsto SC(x)$, where $SC(x)$ is the semantic class label of the point $x$, represents the semantic segmentation of a point cloud. We denote the space of all such feature maps as $\mathcal{F}_{0}$.

For the group of transformations $G$, we use the rotation group $\mathrm{SO}(3)$ for the formulation, for which $\mathrm{SO}(2)$ is also valid.

A rotation $R\in\mathrm{SO}(3)$ can be applied to a point and to a feature map. A rotated feature map $[Rf_{0}]$ is simply the feature map of the rotated point cloud. A point at $x$ goes to $Rx$ after the rotation, thus $f_{0}(x)$ and $[Rf_{0}](Rx)$ are the features of the same point before and after rotating the point cloud. The relation between them depends on the property of $V$. For instance, if $V=SC$, then we know

$$[Rf_{0}](Rx)=f_{0}(x),\quad\forall R\in\mathrm{SO}(3),$$

(1)

i.e., the rotation does not change the semantic class.

As another example, if $V$ represents the surface normal vector, then we have

$$[Rf_{0}](Rx)=Rf_{0}(x),\quad\forall R\in\mathrm{SO}(3),$$

(2)

where on the right-hand side, $R$ is applied to $f_{0}(x)\in\mathbb{R}^{3}$, meaning that the normal vector of a given point rotates along with the point cloud.

In both cases, the feature map of a rotated point cloud, $[Rf_{0}]$, can be generated from $f_{0}$, the feature map of the original point cloud. This property is called equivariance. In the case of Eq. 1, we call $f_{0}$ an invariant scalar field. In the case of Eq. 2, we call $f_{0}$ an equivariant vector field.

A dense prediction task on a point cloud is to reproduce a target field (e.g., $x\mapsto SC(x)$) using a feature map realized by a neural network. Naturally, it would be helpful to equip the network with the same invariant and/or equivariant properties as the target field. However, a general feature map learned by a network is typically neither invariant nor equivariant to rotations. To fix this, we can augment the space of feature maps to $\mathcal{F}=\{f:\mathbb{R}^{3}\times\mathrm{SO}(3)\rightarrow V\}$, defined by

$$f(x,R):=[R^{-1}f_{0}](R^{-1}x)$$

(3)

for some $f_{0}\in\mathcal{F}_{0}$, i.e., the augmented feature map at rotation $R$ equals the original feature map rotated by $R^{-1}$. In this way, we have, $\forall R,R^{\prime}\in\mathrm{SO}(3)$,

$$[Rf](Rx,R^{\prime})\!=\\!=\!f(x,R^{-1}\!R^{\prime}),$$

(4)

which means that the augmented feature map of a rotated point cloud, $[Rf]$, can be generated by the augmented feature map of the original point cloud $f$, indicating that $f$ is equivariant. Equivariant feature maps satisfying Eq. 3 can be constructed using group convolutions [11, 10, 53].

Assuming that we can fit the target field of a point cloud without rotation: $f(x,I)=f_{gt}(x)$, where $I$ is identity rotation. We want to show that the fitting generalizes to the rotated point clouds.

If $f_{gt}$ is an invariant scalar field, from Eq. 4, we have

$$[Rf](Rx,R)=f(x,I)=f_{gt}(x)=[Rf_{gt}](Rx),$$

(5)

meaning that the feature map $[Rf]$ at the rotational coordinate $R$, $[Rf](\cdot,R)$, fits the target $[Rf_{gt}]$ for the rotated point cloud.

If $f_{gt}$ is an equivariant vector field, we have

$$[Rf](Rx,R)\!=\!f(x,I)\!=\!f_{gt}(x)\!=\!R^{-1}[Rf_{gt}](Rx),$$

(6)

which means that we need to apply a rotational matrix multiplication on features in $[Rf](\cdot,R)$ to fit $[Rf_{gt}]$, i.e., $[Rf_{gt}](Rx)=R[Rf](Rx,R),\forall x\in\mathbb{R}^{3}$.

The analysis above shows that the fitting of $f\in\mathcal{F}$ to $f_{gt}\in\mathcal{F}_{0}$ generalizes over rotations, if the rotational coordinate (i.e., the second argument in $f$) is the same as the actual rotation of the point cloud. However, the actual rotation $R$ is usually unknown during inference; thus needs to be learned. In practice, this is formulated as a rotation classification task to select the best rotational channel in a feature map $f$. We will discuss how we instantiate this in the architecture in Section 4.3.

When the target field $f_{gt}$ is invariant, there is another fitting strategy. For equivariant feature map $f$ that satisfies Eq. 3, we can build a rotation-invariant layer by marginalizing over the second argument of $f$, i.e.,

$$f_{inv}(x)=\prod_{R}f(x,R),$$

(7)

where $\prod$ denotes any summarizing operator symmetric to all its arguments, e.g., sum, average, and max.

If we can fit $f_{inv}(x)=f_{gt}(x)$, then based on Eq. 4, we have

$$\begin{split}[Rf_{inv}](Rx)&=\prod_{R^{\prime}}[Rf](Rx,R^{\prime})=\prod_{R^{\prime}}f(x,R^{-1}R^{\prime})\\ &=f_{inv}(x)=f_{gt}(x)=[Rf_{gt}](Rx),\end{split}$$

(8)

which indicates that generalization over rotations holds. In Section 4.3, we discuss different options of implementing the invariant layer in the network.

The model’s group equivariance should match the data’s actual transformations. An overly large group may lead to high computational costs with minimal performance improvement. For the outdoor driving scenario, the $\mathrm{SO}(2)$ group is chosen to represent planar rotations around the gravity axis.

We discretize $\mathrm{SO}(2)$ into a finite group and create a group CNN that is equivariant to the discretized rotations rather than the continuous ones. It allows us to leverage simpler structures akin to conventional deep learning models, allowing integration with existing state-of-the-art networks. Despite discretization, the network is expected to interpolate the equivariance gaps through training with rotational data augmentation.

Specifically, $\mathrm{SO}(2)$ is discretized into cyclic groups $C_{n}$, where $n$ indicates the number of discretized rotations, such as $C_{3}$ for 120-degree rotations. These are referred to as the rotation anchors.

We utilize the point-convolution style equivariant network E2PN [53] as our backbone, which is an equivariant version of KPConv [40]. We describe necessary adaptations to E2PN in Section 4.2, which allow us to build equivariant models on top of SOTA 4D panoptic segmentation network 4D-PLS [1] and 4D-StOP [24], both of which are based on KPConv. We refer to our equivariant models as Eq-4D-PLS and Eq-4D-StOP, respectively.

On a high level, both models first stack the point clouds from several sequential time instances within a common reference frame so that the temporal association becomes part of the instance segmentation. Each network consists of an encoder, a decoder, and prediction heads. The encoders and decoders are very similar for the two models, while the main differences lie in their prediction heads, which formulate the targets as invariant scalar fields and equivariant vector fields, respectively, as discussed in Section 4.3. An overview of our network structure is shown in Fig. 2.

The encoder of the 4D panoptic segmentation networks can be made equivariant by simply swapping the KPConv [40] layers with E2PN [53] convolution layers. However, E2PN is originally designed for $\mathrm{SO}(3)$ equivariance, and uses quotient representations and efficient feature gathering to improve the efficiency. Using it for $\mathrm{SO}(2)$ equivariance requires two adaptations.

First, we use the regular representation instead of the quotient representation in [53]. This is because $\mathrm{SO}(2)$ is abelian, in which case quotient representations cause loss of information (see Appendix D for details).

Second, to apply the efficient feature gathering in E2PN, the spatial position of the convolution kernel needs to be symmetric to the rotation anchors. It allows the rotation of kernel points to be implemented as a permutation of their indices. This implies that different $n$’s in $C_{n}$ impose different constraints on the number and distribution of the kernel points. For example, the default KPConv kernel with 15 points is symmetric to 60-degree rotations and thus can be used to realize $C_{2}$, $C_{3}$, and $C_{6}$ equivariance. However, $C_{4}$ requires a different kernel (for symmetry to 90-degree rotations), for which the 19-point KPConv kernel works.

The original E2PN only provides an encoder to predict a single output for an input point cloud. We need to devise an equivariant decoder for the dense 4D panoptic segmentation. Similar to conventional non-equivariant networks, our decoder adopts upsampling layers and 1-by-1 convolution layers. The upsampling layers simply assign the features of the coarser point cloud to the finer point cloud via nearest neighbor interpolation. The 1-by-1 convolution processes the feature at each point and each rotation independently. It is straightforward to prove the equivariance of such a decoder (see Appendix E).

In Tab. 3, we show the segmentation performance with different invariant pooling layers and the effect of rotational coordinate selection (RCS). In terms of the invariant pooling layers, average pooling is optimal in Eq-4D-StOP, while max pooling is best in Eq-4D-PLS. Eq-4D-PLS favors max pooling, possibly because the point embeddings averaged over all rotational directions may be less discriminative, hindering instance clustering. In contrast, Eq-4D-StOP benefits from average pooling gathering information from all orientations, as the pooled features are used for object property prediction instead of clustering. For the equivariant field (offset) prediction, we compare the rotational coordinate selection with average pooling that wrongly treats the offsets as rotation-invariant targets. The performance drastically decreases when RCS is replaced with average pooling, showing that it is of vital importance to respect the equivariant nature of the targets.