A Recipe for Geometry-Aware 3D Mesh Transformers

TL;DR

This study meticulously examines the various components of a geometry-aware 3D mesh transformer, from tokenization to structural encoding, assessing the contribution of each, and introduces a spectral-preserving tokenization rooted in algebraic multi-grid methods.

Abstract

Utilizing patch-based transformers for unstructured geometric data such as polygon meshes presents significant challenges, primarily due to the absence of a canonical ordering and variations in input sizes. Prior approaches to handling 3D meshes and point clouds have either relied on computationally intensive node-level tokens for large objects or resorted to resampling to standardize patch size. Moreover, these methods generally lack a geometry-aware, stable Structural Embedding (SE), often depending on simplistic absolute SEs such as 3D coordinates, which compromise isometry invariance essential for tasks like semantic segmentation. In our study, we meticulously examine the various components of a geometry-aware 3D mesh transformer, from tokenization to structural encoding, assessing the contribution of each. Initially, we introduce a spectral-preserving tokenization rooted in algebraic multigrid methods. Subsequently, we detail an approach for embedding features at the patch level, accommodating patches with variable node counts. Through comparative analyses against a baseline model employing simple point-wise Multi-Layer Perceptrons (MLP), our research highlights critical insights: 1) the importance of structural and positional embeddings facilitated by heat diffusion in general 3D mesh transformers; 2) the effectiveness of novel components such as geodesic masking and feature interaction via cross-attention in enhancing learning; and 3) the superior performance and efficiency of our proposed methods in challenging segmentation and classification tasks.

Figures (4)

• Figure 1: Our framework at a glance. The two-level approach can be extended to include multiple layers of varying resolutions, each with different patch sizes, which is a common practice for segmentation.
• Figure 2: Comparison between two tokenization methods, Root-Node Selection (RNS) and METIS, reveals significant differences. On the left (A-C), the interpolated approximation of LBO eigenfunctions using RNS demonstrates superior performance compared to METIS. Additionally, on the bottom right of the bunny, the Heat Kernel Signature (HKS) used for patches outperforms the METIS partitioning method. To visually demonstrate the supernodes and partitions based on RNS, Figure $D$ showcases the results on a sample object alongside its counterpart method, METIS ($E$).
• Figure 3: Segmentation results on human body segmentation.
• Figure 4: Segmentation results for RNA mesh segmentation. $A$ and $D$ are the ground truth. $C$ and $F$ are the results.