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: The paper proposes a self-supervised capsule architecture for 3D point clouds. In order to circumvent the customary need to train on pre-aligned datasets, the authors introduce the Canonical Capsules to compute the K-part decomposition of a point cloud that allows unsupervised "object-centric" learning of 3D representations without requiring a semantically aligned dataset. The main idea of the canonical capsules is to enforce invariance/equivariance by relating the decomposition result of two random rigid transformations.  The authors compare the proposed method against state-of-the-art methods on 3D point cloud reconstruction, canonicalization, and unsupervised classification.