Canonical Capsules: Self-Supervised Capsules in Canonical Pose

1 Introduction . Understanding objects is one of the core problems of computer vision [ 32 , 14 , 38 ] . While this task has traditionally relied on large annotated datasets [ 42 , 22 ] , unsupervised approaches that utilize self-supervision [ 5 ] have emerged to remove the need for labels . Recently , researchers have attempted to extend these methods to work on 3D point clouds [ 59 ] , but the field of unsupervised 3D learning remains relatively uncharted . Conversely , researchers have been extensively investigating 3D deep representations for shape auto-encoding1 [ 61 , 19 , 33 , 16 ] , making one wonder whether these discoveries can now benefit from unsupervised learning for tasks other than auto-encoding . Importantly , these recent methods for 3D deep representation learning are not entirely unsupervised . Whether using point clouds [ 61 ] , meshes [ 19 ] , or implicits [ 33 ] , they owe much of their success to the bias within the dataset that was used for training . Specifically , all 3D models in the popular ShapeNet [ 3 ] dataset are “ object-centric ” – they are pre-canonicalized to a unit bounding box , and , even more importantly , with an orientation that synchronizes object semantics to Euclidean frame axes ( e.g . airplane cockpit is always along +y , car wheels always touch z = 0 ) . Differentiable 3D decoders are heavily affected by the consistent alignment of their output with an Euclidean frame [ 8 , 16 ] as local-to-global transformations can not be easily learnt by fully connected layers . As we will show in Section 4.2 , these methods fail in the absence of pre-alignment , even when data 1Auto-encoding is also at times referred to as “ reconstruction ” or “ shape-space ” learning . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) , virtual . augmentation is used . A concurrent work [ 45 ] also recognizes this problem and proposes a separate learnt canonicalizer , which is shown helpful in downstream classification tasks . In this work , we leverage the modeling paradigm of capsule networks [ 23 ] . In capsule networks , a scene is perceived via its decomposition into part hierarchies , and each part is represented with a ( pose , descriptor ) pair : 1 The capsule pose specifies the frame of reference of a part , and hence should be transformation equivariant ; 2 The capsule descriptor specifies the appearance of a part , and hence should be transformation invariant . Thus , one does not have to worry how the data is oriented or translated , as these changes can be encoded within the capsule representation . We introduce Canonical Capsules , a novel capsule architecture to compute a K-part decomposition of a point cloud . We train our network by feeding pairs of a randomly rotated/translated copies of the same shape ( i.e . siamese training ) hence removing the requirement of pre-aligned training datasets . We then decompose the point cloud by assigning each point into one of the K parts via attention , which we aggregate into K keypoints . Equivariance is then enforced by requiring the two keypoint sets to only differ by the known ( relative ) transformation ( i.e . a form of self-supervision ) . For invariance , we simply ask that the descriptors of each keypoint of the two instances match . With Canonical Capsules , we exploit our decomposition to recover a canonical frame that allows unsupervised “ object-centric ” learning of 3D deep representations without requiring a semantically aligned dataset . We achieve this task by regressing canonical capsule poses from capsule descriptors via a deep network , and computing a canonicalizing transformation by solving a shape-matching problem [ 44 ] . This not only allows more effective shape auto-encoding , but our experiments confirm this results in a latent representation that is more effective in unsupervised classification tasks . Note that , like our decomposition , our canonicalizing transformations are also learnt in a self-supervised fashion , by only training on randomly transformed point clouds . Contributions . In summary , in this paper we : • propose an architecture for 3D self-supervised learning based on capsule decomposition ; • enable object-centric unsupervised learning by introducing a learned canonical frame of reference ; • achieve state-of-the-art performance 3D point cloud auto-encoding/reconstruction , canonicalization , and unsupervised classification . 2 Related works . Convolutional Neural Networks lack equivariance to rigid transformations , despite their pivotal role in describing the structure of the 3D scene behind a 2D image . One promising approach to overcome this shortcoming is to add equivariance under a group action in each layer [ 49 , 9 ] . In our work , we remove the need for a global SE ( 3 ) -equivariant network by canonicalizing the input . Capsule Networks . Capsule Networks [ 23 ] have been proposed to overcome this issue towards a relational and hierarchical understanding of natural images . Of particular relevance to our work , are methods that apply capsule networks to 3D input data [ 64 , 65 , 46 ] , but note these methods are not unsupervised , as they either rely on classification supervision [ 65 ] , or on datasets that present a significant inductive bias in the form of pre-alignment [ 64 ] . In this paper , we take inspiration from the recent Stacked Capsule Auto-Encoders [ 30 ] , which shows how capsule-style reasoning can be effective as long as the primary capsules can be trained in a self-supervised fashion ( i.e . via reconstruction losses ) . The natural question , which we answer in this paper , is “ How can we engineer networks that generate 3D primary capsules in an unsupervised fashion ? ” Deep 3D representations . Reconstructing 3D objects requires effective inductive biases about 3D vision and 3D geometry . When the input is images , the core challenge is how to encode 3D projective geometry concepts into the model . This can be achieved by explicitly modeling multi-view geometry [ 28 ] , by attempting to learn it [ 13 ] , or by hybrid solutions [ 60 ] . But even when input is 3D , there are still significant challenges . It is still not clear which is the 3D representation that is most amenable to deep learning . Researchers proposed the use of meshes [ 53 , 31 ] , voxels [ 54 , 55 ] , surface patches [ 19 , 12 , 10 ] , and implicit functions [ 35 , 33 , 7 ] . Unfortunately , the importance of geometric structures ( i.e . part-to-whole relationships ) is often overlooked . Recent works have tried to close this gap by using part decomposition consisting of oriented boxes [ 50 ] , ellipsoids [ 17 , 16 ] , convex polytopes [ 11 ] , and grids [ 2 ] . However , as previously discussed , most of these still heavily rely on a pre-aligned training dataset ; our paper attempts to bridge this gap , allowing learning of structured 3D representations without requiring pre-aligned data . Canonicalization . One way to circumvent the requirement of pre-aligned datasets is to rely on methods capable of registering a point cloud into a canonical frame . The recently proposed CaSPR [ 39 ] fulfills this premise , but requires ground-truth canonical point clouds in the form of normalized object coordinate spaces [ 52 ] for supervision . Similarly , [ 20 ] regresses each view ’ s pose relative to the canonical pose , but still requires weak annotations in the form of multiple partial views . C3DPO [ 34 ] learns a canonical 3D frame based on 2D input keypoints . In contrast to these methods , our solution is completely self-supervised . The concurrent Compass [ 45 ] also learns to canonicalize in a self-supervised fashion , but as the process is not end-to-end , this results in a worse performance than ours , as it will be shown in Section 4.3 . Registration . Besides canonicalization , many pairwise registration techniques based on deep learning have been proposed ( e.g . [ 56 , 63 ] ) , even using semantic keypoints and symmetry to perform the task [ 15 ] . These methods typically register a pair of instances from the same class , but lack the ability to jointly and consistently register all instances to a shared canonical frame . 3 Method . Our network trains on unaligned point clouds as illustrated in Figure 1 : we train a network that decomposes point clouds into parts , and enforce invariance/equivariance through a Siamese training setup [ 48 ] . We then canonicalize the point cloud to a learnt frame of reference , and perform autoencoding in this coordinate space . The losses employed to train E , K , and D , will be covered in Section 3.1 , while the details of their architecture are in Section 3.2 . Decomposition . In more detail , given a point cloud P 2 RP⇥D of P points in D dimensions ( in our case D=3 ) , we perturb it with two random transformations Ta , Tb 2 SE ( D ) to produce point clouds Pa , Pb . We then use a shared permutation-equivariant capsule encoder E to compute a K-fold attention map A 2 RP⇥K for K capsules , as well as per-point feature map F 2 RP⇥C with C channels : A , F = E ( P ) , ( 1 ) where we drop the superscript indexing the Siamese branch for simplicity . From these attention masks , we then compute , for the k–th capsule its pose ✓k 2 R3 parameterized by its location in 3D space , and the corresponding capsule descriptor k 2 RC : ✓k = P p Ap , kPpP p Ap , k , k = P p Ap , kFpP p Ap , k . ( 2 ) Hence , as long as E is invariant w.r.t . rigid transformations of P , the pose ✓k will be transformation equivariant , and the descriptor k will be transformation invariant . Note that this simplifies the design ( and training ) of the encoder E , which only needs to be invariant , rather than equivariant [ 49 , 46 ] . Canonicalization . Simply enforcing invariance and equivariance with the above framework is not enough to learn 3D representations that are object-centric , as we lack an ( unsupervised ) mechanism to bring information into a shared “ object-centric ” reference frame . Furthermore , the “ right ” canonical frame is nothing but a convention , thus we need a mechanism that allows the network to make a choice – a choice , however , that must then be consistent across all objects . For example , a learnt canonical frame where the cockpit of airplanes is consistently positioned along +z is just as good as a canonical frame where it is positioned along the +y axis . To address this , we propose to link the capsule descriptors to the capsule poses in canonical space , that is , we ask that objects with similar appearance to be located in similar Euclidean neighborhoods in canonical space . We achieve this by regressing canonical capsules poses ( i.e . canonical keypoints ) ✓̄ 2 RK⇥3 using the descriptors 2 RK⇥C via a fully connected deep network K : ✓̄ = K ( ) ( 3 ) Because fully connected layers are biased towards learning low-frequency representations [ 27 ] , this regressor also acts as a regularizer that enforces semantic locality . Auto-encoding . Finally , in the learnt canonical frame of reference , to train the capsule descriptors via auto-encoding , we reconstruct the point clouds with per-capsule decoders Dk : P̃ = [ k Dk ( R̄✓k + t̄ , k ) , ( 4 ) where [ denotes the union operator . The canonicalizing transformation T̄ = ( R̄ , t̄ ) can be readily computed by solving a shape-matching problem [ 44 ] , thanks to the property that our capsule poses and regressed keypoints are in one-to-one correspondence : R̄ , t̄ = argmin R , t 1 K X k k ( R✓k + t ) ✓̄kk22 . ( 5 ) While the reconstruction in ( 4 ) is in canonical frame , note it is trivial to transform the point cloud back to the original coordinate system after reconstruction , as the transformation T̄ 1 is available .

Summary: This paper proposes a way of learning a capsule-based representation of 3D point clouds. The representation decomposes an input point cloud into parts, each corresponding to a capsule. This decomposition is trained to be consistent across SE(3) transformations of the point cloud using loss terms that encourage feature invariance and pose equivariance across two random presentations of the input point cloud. Moreover, the model learns a canonical pose for each part that is tied to the identity of the object (not its presented pose).  The experiments show that 1. The reconstruction quality is high, showing that the model has capacity to fit point clouds from ShapeNet. 1. The canonical pose is indeed consistent (measured by low std deviation in relative angle of rotation). 1. The model can be used for pairwise registration (but DeepGMR does even better using RRI features, but doesn't do any of the other things this model can do.) 1. The learned features lend themselves well to unsupervised clustering/classification. 1. All the proposed loss terms are useful and contribute to better reconstruction and canonicalization.