I recently worked on a project for representation learning on 3D shapes using their point clouds and meshes – specifically, 3D scans of teeth. The problem was straitforward: given a 3D point cloud from a scan of a single tooth in a manufacturing facility that makes crowns, determine which CAD design mesh, among tens of thousands of possible designs, it belongs to. The scan and the designs are not in the same orientation/pose, and do not have the same number of points.

There is a geometric solution to this problem - it involves specific steps in 3D-vision for testing each candidate mesh against the 3d scan cloud: (a) global registration using something like ICP (iterative closest point) to align the scan point cloud to each design mesh as best as possible, (d) measuring global deviation of the scan points with respect to the design mesh. Not only is this is extrmely slow, it does not work very well, i.e. it does produce false positives. Further it is sensitive to the initial orientation of the scan with respect to the orientation of the design mesh.

TLDR; SE3-equivariant architectures do not necesarily beat simpler architectures with rotational augmentation

The fact that the scan and the design are not in the same orientation but must still map to the same embedding gives the problem a flavor of SE3-invariance. I built a self-supervised learning model for these 3D objects using two different architectures, both trained using a InfoNCE-style SupCon loss, where for each anchor sample, there are multiple positives and many negatives.

Contrastive Learning
We want the learnt embedding to be pose invariant
  • One, using an inherently SE3-equivariant architecture called DeepSphere that builds on the Spherical CNN work of Cohen, Welling et al, by replacing the high computational complexity operations of convolutaions with spherical harmonics that yield exact SO3 equivariance, with an approximation: the 3D point cloud or mesh of the object (the tooth inthis case) is first projected onto a sphere to create a spherical depth map. This is then discretized onto a grid that lives on the sphere (some discrete projections are: Healpix, equi-angular, and Fibonacci). I used Healpix using the healpy library. We now have a regular grid, kind of like a 2D image, but instead of a flat plane, it lives on a sphere. DeepSphere then treats these pixels on the sphere as a graph with each pixel being a node with a fixed number of neighbors, and applies a graph neural network whose filters are radial, i.e. the weight of the filters are learnt but are radially symmetric – this is called Chebyshev convolution. If the spherical signal rotates, so does the output of the graph neural network. The extent to which this architecture is equivariant depends upon the number of pixels on the grid – therefore there is a tradeoff between the amount of compute and the level of equivariance achieved. SO3 invariance is achieved using a pooling layer at the end. I liked the DeepSpere paper and recommend reading it.
3D mesh of the bust of Nefertiti
As an example, here is 3D mesh of the bust of Nefertiti.

The spherical depthmap, created by casting a ray through the point cloud of the above mesh's vertices. In practice one would cast a ray to every pixel on the healpix grid of a sphere, and intersect with the mesh.

A second channel made using the cosine of the angle between each ray and the normal at the point of intersection.
  • Second, I tried to use a PointTransformerV3 architecture, which is not SE3-equivariant, but can be driven to learn equivariance by aggressive augmentation. The PTv3 architecture is surprisingly fast for a transformer based architecture – and it achieves this by using a space filling curve (Hilbert curves or Z-curves) to transform a 3D problem consisting of a point cloud into a 1D problem! On the space filling curves, the ‘neighbors’ of any point in 3D space are likely to be real 3D nearest neighbors of that point, thereby eliminating the need for 3D nearest neighbors when considering which points to attend to. The ‘aggressive augmentation’ is achievd by applying a random SO3 rotation to each 3D scan and its corresponding matched design file as the positve sample.

To my surprise, the PTv3 architecture achieved very strong performance – better than the spherical CNN when using the same batch size, and surpassed the spherical CNN when the batch size of the PTv3 model was increased (contrastive learning benefits from large batches because they provide a large number of negatives for each positive pair sample in the batch). I also used rotational augmentation on the spherical CNN since this form of spherical CNN is not prfectly SE3-equivariant.