A Recipe for Geometry-Aware 3D Mesh Transformers

TL;DR

This paper develops a geometry-aware 3D mesh transformer that uses spectral-preserving tokenization and advanced structural embeddings, improving segmentation and classification performance on unstructured 3D mesh data.

Contribution

It introduces a spectral-preserving tokenization method and novel structural embeddings for 3D meshes, addressing invariance and efficiency issues in transformer-based models.

Findings

Spectral-preserving tokenization enhances mesh representation.

Heat diffusion-based embeddings improve structural awareness.

Proposed methods outperform baseline models in segmentation and classification.

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.