SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Congyue Deng1* Jiahui Lei2* Bokui Shen1 Kostas Daniilidis2,3 Leonidas Guibas1
$^{1}$ Stanford University  $^{2}$ University of Pennsylvania  $^{3}$ Archimedes, Athena RC
{congyue, willshen, guibas}@cs.stanford.edu, {leijh, kostas}@cis.upenn.edu
Equivariance has gained strong interest as a desirable network property that inherently ensures robust generalization. However, when dealing with complex systems such as articulated objects or multi-object scenes, effectively capturing interpart transformations poses a challenge, as it becomes entangled with the overall structure and local transformations. The interdependence of part assignment and per-part group action necessitates a novel equivariance formulation that allows for their co-evolution. In this paper, we present *Banana*, a Banach fixed-point network for pointcloud segmentation with inter-part equivariance by construction. Our key insight is to iteratively solve a fixed-point problem, where point-part assignment labels and per-part SE(3)-equivariance co-evolve simultaneously. We provide theoretical derivations of both per-step equivariance and global convergence, which induces an equivariant final convergent state. Our formulation naturally provides a strict definition of inter-part equivariance that generalizes to unseen inter-part configurations. Through experiments conducted on both articulated objects and multi-object scans, we demonstrate the efficacy of our approach in achieving strong generalization under inter-part transformations, even when confronted with substantial changes in pointcloud geometry and topology.
From articulated objects to multi-object scenes, multi-body systems, i.e. shapes composed of multiple parts where each part can be moved separately, are prevalent in various daily-life scenarios. However, modeling such systems presents considerable challenges compared to rigid objects, primarily due to the infinite shape variations resulting from inter-part pose transformations with exponentially growing complexities. Successfully modeling these shapes demands generalization across potential inter-part configurations. While standard data augmentation techniques can potentially alleviate such problems, exhaustive augmentation can be highly expensive especially given the complexity of such systems compared to rigid objects. Meanwhile, recent rigid-shape analysis techniques have made significant progress in building generalization through the important concept of equivariance [39, 14, 79, 2, 81, 71, 20, 11, 15, 3, 36, 42, 55]. Equivariance dictates that when a transformation is applied to the input data, the network's output should undergo a corresponding transformation. For example, if a is rotated by 90 degrees, the segmentation masks should also rotate accordingly. However, to the best of our knowledge, no existing work has managed to extend the formulation of equivariance to inter-part configurations.
In this paper, we address the challenge of achieving inter-part equivariance in the context of part segmentation, a fundamental task in analysis, which is also the key to generalizing equivariance from single object to multi-body system. Extending equivariance to multi-body systems presents notable challenges in both formulation and realization. When dealing with a shape that allows for inter-part motions, the model must exhibit equivariance to both global and local state changes, which can only be defined by combining part assignment and per-part transformations. For example, to model an
Global equivariance
Inter-part equivariance
Figure 1: "Chicken or the egg" problem with inter-part equivariance. Left: Rigid SE(3)-equivalence. If a rigid transformation is applied to the input, the output segmentation transforms accordingly. Right: Inter-part equivariance. Analogously, we want the input and output to be coherent under per-part transformations. But such a definition of equivariance requires part segmentations, which is also exactly the desired network output, resulting in a "chicken or the egg" problem.
oven's inter-part states, we first need to differentiate between the parts corresponding to the door and the body and then define the movement of each individual part. An intriguing "chicken or the egg" problem arises when attempting to build inter-part equivariance without access to a provided segmentation, where segmentation is necessary to define equivariance, but segmentation itself is the desired output that we aim to produce.
Our key insight to tackle this seeming dilemma is to model inter-part equivariance as a sequential fixed-point problem by co-evolving part segmentation labels and per-part SE(3)-equivariance. We provide theoretical derivations for both the per-step behavior and global convergence behavior, which are crucial aspects of a fixed-point problem. We show that our formulation establishes per-step equivariance through network construction, which then induces an overall inter-part equivariance upon convergence. Thus, by having equivariant per-step progression and global convergence, our formulation naturally gives rise to a strict definition of inter-part equivariance through iterative inference that generalizes to unseen part configurations. We further bring our formulation to concrete model designs by proposing a novel part-aware equivariant network with a weighted message-passing paradigm. With localized network operators and per-part SE(3)-equivariant features, the network is able to guarantee per-step inter-part equivariance as well as facilitate stable convergence during the iterative inference. We test our framework on articulated objects to generalize from static rest states to novel articulates and multi-object scenes to generalize from clean synthetic scenes to cluttered real scans. Our model shows strong generalization in both scenarios even under significant changes in pointcloud geometry or topology.
To summarize, our key contributions are
- To the best of our knowledge, we are the first to provide a strict definition of inter-part equivalence for pointcloud segmentation and introduce a learning framework with such equivalence by construction.
- We propose a fixed-point framework with one-step training and iterative inference and show that the per-step equivariance induces an overall equivariance upon convergence.
- We design a part-aware equivariant message-passing network with stable convergence.
- Experiments show our strong generalization under inter-part configuration changes even when they cause subsequent changes in pointcloud geometry or topology.
Equivariant pointcloud networks. Existing works in equivariant 3D learning mainly focus on rigid SE(3) transformations. As a well-studied problem, it has developed comprehensive theories [39, 14, 79, 2, 81] and abundant network designs [71, 20, 11, 15, 3, 36, 42, 55], which benefit a variety of 3D vision and robotics tasks ranging from pose estimation [44, 46, 54, 61, 89], shape reconstruction [11, 10], to object interaction and manipulation [19, 27, 60, 63, 80, 82]. A few recent works have extended SE(3) equivariance to part level [86, 41, 50]. [86] employs a local equivariant feature extractor for object bounding box prediction, showing robustness under object-level and scene-level pose changes. [41] learns an SE(3)-equivariant object prior and applies it to object
Training with ground-truth segmentation
Iterative inference with per-step equivariance
Figure 2: Single-state training and novel-state iterative inference. Left: Training with ground-truth segmentation as both network input and output. Right: Equivariant Banach iterations. At each step  $k$ , the network is equivariant to its current segmentation input  $\mathbf{y}^{(k)}$  by construction. At the end of the iterations, it converges to its fixed point with an induced overall equivariance.
detection and segmentation in scenes with robustness to scene configuration changes. [50] learns part-level equivariance and pose canonicalization from a collection of articulated objects with a two-stage coarse-to-fine network. However, all these works are purely heuristics-based, and none has provided a strict definition of inter-part equivariance.
Multi-body systems. From small-scale articulated objects to large-scale multi-object scenes, there is a wide range of 3D observations that involve multiple movable parts or objects, exhibiting various configurations and temporal variations. This has sparked a rich variety of works dedicated to addressing the challenges in multi-body systems, studying their part correspondences [56, 49, 8, 48, 43], segmentations [85, 76, 67, 31, 30, 70, 7, 12], reconstructions [34, 37, 32, 53, 40], or rearrangements [78, 68], to name just a few. However, as we know the system is acted via products of group actions; a more structured network that exploits such motion prior by construction is desired. Only a few works have studied this problem. [86] exploits global and local gauge equivariance in 3D object detection via bottom-up point grouping but without a systematical study of inter-object transformations. [41] further exploits equivariance of object compositions in scene segmentation, but as an EM-based approach, it requires exhaustive and inefficient enumerations and no global convergence is guaranteed. In contrast, our approach offers a theoretically sound framework that achieves multi-body equivariance by construction, which enables robust analysis for various tasks and provides a more structured and reliable solution compared to existing methods.
Iterative inference. From recurrent neural networks [23, 29] to flow [16, 59, 17, 38, 83] or diffusion-based generative models [65, 66, 28, 64, 51, 88], iterative inference has played many diverse yet essential roles in the development of computer vision. Though naturally adapted to sequential or temporal data, iterative inference can also be applied to static learning problems such as image analysis [22, 21, 73, 26] and pointcloud segmentation [84, 47]. Recent studies also show that iterative inference presents stronger generalizability than one-step forward predictions [62, 6], with explanations of their relations to the "working memory" of human minds [4, 5] or human visual systems [45, 35]. In our work, we employ iterative inference on static pointclouds to co-evolve two intertwined attributes: part segmentation and inter-part equivariance, in order to address a seemingly contradictory situation.
3 Inter-Part Equivariance
We start by defining inter-part equivariance on point clouds with given part segmentations (Sec. 3.1). Then, we will introduce how to extend this definition to unsegmented point clouds using a fixed-point framework (Sec. 3.2) with one-step training and iterative inference. The iterations have per-step equivariance by network construction, which is shown to induce an overall equivariance upon convergence (Sec 3.3). Finally, we will also explain how to eliminate part orders for instance segmentation (Sec 3.4).
3.1 Multi-Body Systems
Let  $\mathbf{X} \in \mathbb{R}^{N \times 3}$  be a pointcloud with  $P$  parts. Segmentations on  $\mathbf{X}$  are represented as point-part assignment matrices  $\mathbf{y} \in [0,1]^{N \times P}$  which sum up to 1 along the  $P$  dimension. As the action of  $\mathrm{SE}(3)^P$  on  $\mathbf{X}$  and its resulting equivariance is subject to the part decompositions, we define the action of  $\mathrm{SE}(3)^P$  on (pointcloud, segmentation) pairs  $(\mathbf{X},\mathbf{y})$ . When  $\mathbf{y}$  is a binary mask, the  $P$
parts of  $\mathbf{X}$  are  $P$  disjoint sub-pointclouds and each  $\mathrm{SE}(3)$  component of  $\mathrm{SE}(3)^P$  acts on one part separately. More concretely, for any transformation  $\mathbf{A} = (\mathbf{T}_1,\dots ,\mathbf{T}_P)\in \mathrm{SE}(3)^P$ , we define  $\mathbf{A}(\mathbf{X},\mathbf{y})\coloneqq (\mathbf{X}',\mathbf{y})$  where the  $n$ -th point  $\mathbf{x}_n$  is transformed to  $\mathbf{x}_n'$  by
$$
\mathbf {x} _ {n} ^ {\prime} := \sum_ {p = 1} ^ {P} \mathbf {y} _ {n p} \left(\mathbf {x} _ {n} \mathbf {R} _ {p} + \mathbf {t} _ {p}\right). \tag {1}
$$
where  $\mathbf{R}_p$  and  $\mathbf{t}_p$  are the rotation and translation components of  $\mathbf{T}_p\in \mathrm{SE}(3)$ . For soft segmentation masks, we view each point as probabilistically assigned to one of the  $P$  parts.
Equivalence. Similar to  $\mathrm{SE}(3)$ -equivariance on rigid shapes  $f(\mathbf{T}\mathbf{X}) = \mathbf{T}f(\mathbf{X}),\forall \mathbf{T}\in \mathrm{SE}(3)$ , we define the inter-part equivariance on  $(\mathbf{X},\mathbf{y})$  pairs by saying that a function  $f$  is  $\mathrm{SE}(3)^P$ -equivariant if  $\forall \mathbf{A}\in \mathrm{SE}(3)^P$ ,
$$
f (\mathbf {A} (\mathbf {X}, \mathbf {y})) = \mathbf {A} f (\mathbf {X}, \mathbf {y}). \tag {2}
$$
Note that if the function outputs  $f(\mathbf{X}, \mathbf{y})$  are no longer (pointcloud, segmentation) pairs, the  $\mathrm{SE}(3)^P$ -action on its outputs needs to be specified. A special case is when  $\mathrm{SE}(3)^P$  acts trivially on  $f(\mathbf{X}, \mathbf{y})$  by  $f(\mathbf{A}(\mathbf{X}, \mathbf{y})) = f(\mathbf{X}, \mathbf{y})$ , which defines the common-sense "invariance". For a pointcloud segmentation network with per-point part-label outputs  $\mathbf{y} \in [0,1]^{N \times P}$ , we desire it to be  $\mathrm{SE}(3)^P$ -invariant under the above definitions, that is  $f(\mathbf{A}(\mathbf{X}, \mathbf{y})) \equiv \mathbf{y}$ ,  $\forall \mathbf{A} \in \mathrm{SE}(3)^P$ .
3.2 Banach Fixed-Point Iterations
"Chicken or the egg". The "chicken or the egg" nature of segmentation with inter-part equivariance immediately emerges from the above definition: to enforce  $\mathrm{SE}(3)^P$ -equivariance for a neural network  $f(\cdot ;\Theta)$ , the part segmentation  $\mathbf{y}$  is required as input, which is also exactly the desired output of the network (Fig. 1). That this,
$$
f (\mathbf {X}, \mathbf {y}; \Theta) = \mathbf {y}. \tag {3}
$$
To resolve this dilemmatic problem, we approach Eq. 3 as a fixed-point equation on function  $f(\cdot ,\mathbf{y};\Theta)$  w.r.t.  $\mathbf{y}$ . And instead of using single-step forward predictions, we solve it with Banach fixed-point iterations (Fig. 2).
Training. At training time, we aim to optimize the network weights  $\Theta$  such that the labeled  $s(\mathbf{X}_i,\mathbf{y}_i), i\in \mathcal{I}$  from the dataset become fixed points of the function  $f(\cdot ,\Theta)$ :
$$
\Theta_ {*} = \underset {\Theta} {\arg \min } \frac {1}{| \mathcal {I} |} \sum_ {i \in \mathcal {I}} \| f \left(\mathbf {X} _ {i}, \mathbf {y} _ {i}; \Theta\right) - \mathbf {y} _ {i} \|. \tag {4}
$$
However, there exists a trivial solution which is the identity function  $f(\cdot ,\mathbf{y};\Theta)\equiv \mathbf{y}$ , meaning that any arbitrary segmentation can satisfy the objective for any point clouds. Thus, to avoid such degenerated cases, we have to limit the network expressivity. More specifically, we limit the Lipschitz constant  $L$  of the network w.r.t.  $\mathbf{y}$  which is defined by
$$
\| f (\mathbf {X}, \mathbf {y} _ {1}; \Theta) - f (\mathbf {X}, \mathbf {y} _ {2}; \Theta) \| \leq L \| \mathbf {y} _ {1} - \mathbf {y} _ {2} \|, \quad \forall \mathbf {X}, \forall \mathbf {y} _ {1}, \mathbf {y} _ {2}. \tag {5}
$$
If  $f$  degenerates to the identical mapping, its Lipschitz  $L = 1$ , which violates the above constraint. When  $L < 1$ ,  $f$  is a contractive function and has a unique fixed point based on the Banach fixed-point theorem. Our training objective is to align this fixed point with the ground-truth segmentation.
Iterative inference. At inference time, given an input pointcloud  $\mathbf{X}_0$ , the learnable parameters  $\Theta_*$  fixed in  $f$  and we look for a segmentation  $\mathbf{y}$  that satisfies
$$
f _ {\mathbf {X} _ {0}, \Theta_ {*}} (\mathbf {y}) := f \left(\mathbf {X} _ {0}, \mathbf {y}; \Theta_ {*}\right) = \mathbf {y}. \tag {6}
$$
We solve this equation with Banach fixed-point iterations [1]
$$
\mathbf {y} ^ {*} = \lim  _ {k \rightarrow \infty} \mathbf {y} ^ {(k)} = \lim  _ {k \rightarrow \infty} f _ {\mathbf {X} _ {0}, \Theta_ {*}} ^ {(k)} \left(\mathbf {y} ^ {(0)}\right) \tag {7}
$$
where  $\mathbf{y}^{(0)}$  is a random initialization of the segmentation on  $\mathbf{X}_0$ . For a contractive function  $f$ , the uniqueness of  $\mathbf{y}^*$  induces a well-defined mapping from pointclouds to segmentations  $f_{\Theta}^{*}: \mathbf{X}_{0} \to \mathbf{y}^{*}$ .
3.3 Equivariance
For an input pairs  $(\mathbf{X},\mathbf{y})$  with known segmentation  $\mathbf{y}$ , we employ an  $\mathrm{SE}(3)$ -equivariant backbone [15] on each individual part and construct an  $\mathrm{SE}(3)^P$ -equivariant network, which will be further explained in Section 4. This guarantees the inter-part equivariance at training time and at each timestep during inference, but most importantly we desire the overall  $\mathrm{SE}(3)^P$ -equivariance at the convergence point of the Banach iteration. Here we prove two properties:
Self-coherence of convergent state. At each inference timestep  $k$ , we have  $\mathrm{SE}(3)^P$ -equivariance subject to  $\mathbf{y}^{(k)}$ , that is,  $\forall \mathbf{A} \in \mathrm{SE}(3)^P$
$$
f \left(\mathbf {A} \left(\mathbf {X} _ {0}, \mathbf {y} ^ {(k)}\right); \Theta\right) = f \left(\mathbf {X} _ {0}, \mathbf {y} ^ {(k)}; \Theta\right), \quad \forall k \in \mathbb {N}, \tag {8}
$$
Because of the continuity of  $f$  and the compactness of  $[0,1]^{N\times P}$ , we can take limits w.r.t.  $k$  on both sides, which gives
$$
f \left(\mathbf {A} \left(\mathbf {X} _ {0}, \mathbf {y} ^ {*}\right); \Theta\right) = f \left(\mathbf {X} _ {0}, \mathbf {y} ^ {*}; \Theta\right) = f _ {\Theta} ^ {*} \left(\mathbf {X} _ {0}\right). \tag {9}
$$
This shows the self-coherence of the convergent point between  $f_{\Theta}^{*}$  and  $\mathbf{y}^{*}$ .
Generalization to novel inter-part states. Suppose the network has seen  $(\mathbf{X},\mathbf{y}_{\mathrm{gt}})$  in the training set and learned  $f_{\Theta}(\mathbf{X},\mathbf{y}_{\mathrm{gt}};\Theta_{*}) = \mathbf{y}_{\mathrm{gt}}$  with the loss function in Eq. 4 minimized to zero. Now apply an inter-part transformation  $\mathbf{A}\in \operatorname {SE}(3)^P$  and test the network on  $(\mathbf{X}^{\prime},\mathbf{y}_{\mathrm{gt}})\coloneqq \mathbf{A}(\mathbf{X},\mathbf{y}_{\mathrm{gt}})$ . We would like to show  $f_{\Theta}^{*}(\mathbf{X}^{\prime}) = f_{\Theta}^{*}(\mathbf{X}) = \mathbf{y}_{\mathrm{gt}}$ . First of all, we know that  $\mathbf{y}_{\mathrm{gt}}$  itself is a fixed point of  $f(\mathbf{X}^{\prime},\cdot ;\Theta)$ , which is due to the equivariance w.r.t.  $\mathbf{y}_{\mathrm{gt}}$  at training time:
$$
f \left(\mathbf {X} ^ {\prime}, \mathbf {y} _ {\mathrm {g t}}; \Theta\right) = f (\mathbf {A} \left(\mathbf {X}, \mathbf {y} _ {\mathrm {g t}}\right); \Theta) = f \left(\mathbf {X}, \mathbf {y} _ {\mathrm {g t}}; \Theta\right). \tag {10}
$$
Now if  $f$  is a contractive function on  $\mathbf{y}$ , we have the uniqueness of Banach fixed-points and thus  $\mathbf{y}_{\mathrm{gt}}$  is the only fixed point for  $f(\mathbf{X}', \cdot; \Theta)$ , implying that the iterations in Eq. 7 must converge to  $\mathbf{y}_{\mathrm{gt}}$ . Equivalence of the per-step iteration and the final convergent state are illustrated in Fig. 2 right.
In less ideal cases when the training loss is not zero but a small error  $\varepsilon$ , we view it as the distance  $\| \mathbf{y}^{(1)} - \mathbf{y}^{(0)} \|$  between the ground-truth  $\mathbf{y}^{(0)} = \mathbf{y}_{\mathrm{gt}}$  and the one-step iteration output  $\mathbf{y}^{(1)}$ . The distance between the actual fixed-point  $\mathbf{y}^*$  and the ground-truth  $\mathbf{y}_{\mathrm{gt}}$  is then bounded by
$$
\left\| \mathbf {y} ^ {*} - \mathbf {y} _ {\mathrm {g t}} \right\| = \sum_ {k = 0} ^ {\infty} \left\| \mathbf {y} ^ {(k + 1)} - \mathbf {y} ^ {(k)} \right\| \leq \sum_ {k = 0} ^ {\infty} L ^ {k} \left\| \mathbf {y} ^ {(1)} - \mathbf {y} ^ {(0)} \right\| = \frac {1}{1 - L} \varepsilon \tag {11}
$$
where  $L$  is the Lipschitz constant of the network.
3.4 Part Permutations
A pointcloud segmentation represented by  $\mathbf{y} \in [0,1]^{N \times P}$  inherently carries an order between the  $P$  parts, which is the assumption in semantic segmentation problems. But in various real-world scenarios, it is common to encounter situations where multiple disjoint parts are associated with the same semantic label, without any clear or coherent part orderings. In order to tackle this challenge, we also extend our method to encompass the instance segmentation problem without part orderings. For simplicity, here we assume that all parts can be permuted together by  $\mathfrak{S}_P$ . In practice, permutations only occur among the parts within each semantic label, for which we can easily substitute the  $\mathfrak{S}_P$  below with its subgroup and the conclusion still holds. We define an equivalence relation on  $[0,1]^{N \times P}$  by
$$
\mathbf {y} _ {1} \sim \mathbf {y} _ {2} \iff \exists \sigma \in \mathfrak {S} _ {P}, s. t. \mathbf {y} _ {1} \sigma = \mathbf {y} _ {2}, \tag {12}
$$
which gives us a quotient space  $[0,1]^{N\times P} / \mathfrak{S}_P$  and each equivalent class  $\hat{\mathbf{y}}$  in this quotient space represents an instance segmentation without part ordering. We can define a metric on  $\hat{\mathbf{y}}$  by
$$
d \left(\hat {\mathbf {y}} _ {1}, \hat {\mathbf {y}} _ {2}\right) := \min  _ {\sigma \in \mathfrak {S} _ {P}} \| \mathbf {y} _ {1} \sigma - \mathbf {y} _ {2} \|, \tag {13}
$$
which makes  $\left([0,1]^{N\times P} / \mathfrak{S}_P,d\right)$  a complete metric space, on which the Banach fixed-point theorem holds and so are the iterations and convergence properties. Proofs for the well-definedness of  $d$  and the completeness of the quotient space are shown in the supplementary material.
Figure 3: Part-aware equivariant network. Left: Segmentation-weighted message passing. Right: Overall architecture for segmentation label update. Color gradients on the tensors indicate the usage of segmentation labels,
4 Part-Aware Equivariant Network
In this section, we present our part-aware equivariant network for segmentation label updates (Fig.3). We utilize a SE(3)-equivariant framework called Vector Neurons (VN) [15] as the backbone to encode per-point features  $\mathbf{V} \in \mathbb{R}^{N \times C \times 3}$ , ensuring equivariance within each part. Subsequently, these features are transformed into invariant representations, eliminating relative poses to facilitate inter-part information propagation and global aggregations.
Segmentation-weighted message passing. Key to our part-aware equivariant network is a message-passing module weighted by the input segmentation  $\mathbf{y}$  (Fig. 3 left). Intuitively, given a point  $\mathbf{x}_n \in \mathbf{X}$  with latent feature  $\mathbf{V}_n$ , for its neighborhood points  $\mathbf{x}_m \in \mathcal{N}$ , we only allow information propagation between  $\mathbf{x}_n$  and  $\mathbf{x}_m$  if they belong to the same part. When the part segmentation  $\mathbf{y}$  is a soft mask, we compute the probability  $p_{nm}$  of  $\mathbf{x}_n$  and  $\mathbf{x}_m$  belonging to the same part by  $p_{nm} = \mathbf{y}_n\mathbf{y}_m^t$ ,  $\mathbf{y}_n, \mathbf{y}_m \in [0,1]^{1 \times P}$  and weight the local message passing with
$$
\mathbf {V} _ {n} ^ {\prime} = \sum_ {m \in \mathcal {N}} p _ {n m} \varphi \left(\mathbf {V} _ {m} - \mathbf {V} _ {n}, \mathbf {V} _ {n}\right) / \sum_ {m \in \mathcal {N}} p _ {n m}. \tag {14}
$$
Proof of inter-part equivariance for this module can be found in the supplementary material.