Masked BRep Autoencoder via Hierarchical Graph Transformer

TL;DR

This paper presents a self-supervised learning framework using a hierarchical graph transformer for CAD BRep models, enabling effective downstream task performance with limited labeled data.

Contribution

It introduces a masked graph autoencoder with a hierarchical graph transformer architecture for representation learning on CAD BRep models, improving downstream task accuracy with minimal labeled data.

Findings

High performance on downstream tasks with limited labeled data

Significant improvement over existing methods in low-data regimes

Effective modeling of geometric dependencies in CAD models

Abstract

We introduce a novel self-supervised learning framework that automatically learns representations from input computer-aided design (CAD) models for downstream tasks, including part classification, modeling segmentation, and machining feature recognition. To train our network, we construct a large-scale, unlabeled dataset of boundary representation (BRep) models. The success of our algorithm relies on two keycomponents. The first is a masked graph autoencoder that reconstructs randomly masked geometries and attributes of BReps for representation learning to enhance the generalization. The second is a hierarchical graph Transformer architecture that elegantly fuses global and local learning by a cross-scale mutual attention block to model long-range geometric dependencies and a graph neural network block to aggregate local topological information. After training the autoencoder, we replace its decoder with a task-specific network trained on a small amount of labeled data for downstream tasks. We conduct experiments on various tasks and achieve high performance, even with a small amount of labeled data, demonstrating the practicality and generalizability of our model. Compared to other methods, our model performs significantly better on downstream tasks with the same amount of training data, particularly when the training data is very limited.