BRep-BERT: Pre-training Boundary Representation BERT with Sub-graph Node Contrastive Learning

Obtaining effective entity feature representations is crucial in the field of Boundary Representation (B-Rep), a key parametric representation method in Computer-Aided Design (CAD). However, the lack of labeled large-scale database and the scarcity of task-specific label sets pose significant challenges. To address these problems, we propose an innovative unsupervised neural network approach called BRep-BERT, which extends the concept of BERT to the B-Rep domain. Specifically, we utilize Graph Neural Network (GNN) Tokenizer to generate discrete entity labels with geometric and structural semantic information. We construct new entity representation sequences based on the structural relationships and pre-train the model through the Masked Entity Modeling (MEM) task. To address the attention sparsity issue in large-scale geometric models, we incorporate graph structure information and learnable relative position encoding into the attention module to optimize feature updates. Additionally, we employ geometric sub-graphs and multi-level contrastive learning techniques to enhance the model's ability to learn regional features. Comparisons with previous methods demonstrate that BRep-BERT achieves the state-of-the-art performance on both full-data training and few-shot learning tasks across multiple B-Rep datasets. Particularly, BRep-BERT outperforms previous methods significantly in the few-shot learning scenarios. Comprehensive experiments demonstrate the substantial advantages and potential of BRep-BERT in handling B-Rep data representation. Code will be released at https://github.com/louyz1026/Brep_Bert.