Masked BRep Autoencoder via Hierarchical Graph Transformer

Yifei Li; Kang Wu; Wenming Wu; Xiao-Ming Fu

cs.GRMar 16, 2026

Masked BRep Autoencoder via Hierarchical Graph Transformer

Authors: Yifei Li, Kang Wu, Wenming Wu, Xiao-Ming Fu

Organizations: University of Science and Technology of China · Hefei University of Technology

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.

Figures & tables

Figure 1: Given a CAD model, our method first applies a BRep encoder to construct a gAAG from the extracted BRep information. A Graph encoder is then implemented to update the graph features with both global and local information. During pre-training, we train an MAE by reconstructing the randomly masked faces and edges for BRep representation learning (the upper part of the figure). In the fine-tuning stage, we train a new network, formed by connecting a task-specific head behind the encoder, using a small amount of labeled data with different loss functions (lower part of the figure).

Figure 2: Our BRep encoder extracts and fuses geometric and attribute features into a unified gAAG representation via parallel CNN and MLP branches.

Panel of Figure 3 — Figure 3: Left: The overall architecture of the Graph encoder, processing multi-resolution face features $F_{\text{low}}$ , $F_{\text{high}}$ , and edge features $E$ . Right: The internal structure of the $k$ -th CSMA block. CA and SA denote cross-attention and self-attention layers, respectively, while Q, K, and V represent queries, keys, and values.

Figure 4: Four reconstruction examples. For each example, the sequence from left to right shows: (1) the input BRep model, (2) the full surface point cloud, (3) the masked input where black denotes masked faces and blue denotes unmasked faces, and (4) the reconstructed point cloud.

		Method	0.1%	0.5%	1%	1.5%	2%	3%	100%
MFInstSeg	Acc	AAGNet	46.49	76.62	85.19	90.30	92.31	92.37	99.13
		BRepFormer	33.43	86.14	91.23	95.77	96.51	97.64	99.62
		BRepMAE	88.81	94.75	96.99	97.44	97.66	98.20	99.57
		Ours	88.75	96.19	97.31	97.67	97.92	98.26	99.53
	mIoU	AAGNet	34.01	67.51	77.50	84.64	86.76	87.32	98.36
		BRepFormer	7.47	62.66	73.89	86.17	88.25	91.85	98.75

Table 1: Comparison across MFInstSeg, MFCAD++, and CADSynth under varying supervision ratios. Accuracy (Acc, %) and mIoU (%) are reported.

Model	10-shot		20-shot		Full
	Acc	mIoU	Acc	mIoU	Acc	mIoU
BRep-BERT	59.26	19.02	62.30	24.20	95.14	82.88
BRT	54.98	13.54	56.69	14.99	94.48	79.23
BRepFormer	50.24	12.31	51.75	13.38	94.02	78.73
Ours	72.33	34.63	74.31	38.22	97.02	86.75

Table 2: Performance comparison on Fusion 360 Gallery Segmentation under different levels of supervision. The

k

-shot setting denotes using

k

labeled training samples in total. Accuracy (Acc, %) and mIoU (%) are reported.

Figure 5: Machining feature recognition under the full supervised setting. We accurately predict various machining features, such as passages, pockets, slots, and steps.

Figure 6: Examples of modeling segmentation results on the Fusion 360 Gallery segmentation dataset under the full supervised setting. Our model accurately predicts per-face labels across 8 operation types, with distinct colors for each.

Panel of Figure 7 — Figure 7: Performance comparison with BRT under the 10-shot setting. Left: The segmentation task, where our method predicts face labels more accurately. Right: The part classification task, where our method correctly identifies all shapes as n (class 13), while BRT misclassifies the last two.

Model	10-shot Acc	20-shot Acc	Full Acc
BRep-BERT	68.71	75.92	97.76
BRT	54.84	62.69	97.35
BRepFormer	76.53	84.25	97.59
Ours*	81.83	88.47	97.78
Ours	82.51	89.58	97.91

Table 3: Comparison on SolidLetter classification under different levels of supervision. The

k

-shot setting denotes using

k

labeled training samples for each class. “Ours*” denotes our model pre-trained without SolidLetter data. Accuracy (Acc, %) is reported.

Configuration	0.1%	0.5%	1%	1.5%	2%	3%	100%
Single Res. ( $13\times 13$ only)	85.83	93.14	95.19	95.88	97.16	97.42	99.27
Triple Res. (w/ $7\times 7$ )	86.41	94.46	96.40	97.04	97.38	97.78	99.42
w/o Masking (0% Mask)	66.41	84.19	92.99	94.20	95.33	95.96	99.43
MPNN first	76.39	92.10	94.95	95.62	96.12	96.41	99.06
Mask gAAG	78.63	92.98	94.36	95.11	95.75	96.48	99.16
Frozen encoder	82.01	92.71	94.70	95.74	96.15	96.56	98.64

Table 4: Ablation study on MFInstSeg under various supervision ratios. The core architectural and pre-training designs are evaluated.

Appendix figures & tables17 assets

Supplementary material from the paper’s appendix.

Appendix

Operator	Input Shape	Output Shape
Conv2D, $k=3,p=1$	$7\times N\times N$	$64\times N\times N$
GroupNorm, $G=8$	$64\times N\times N$	$64\times N\times N$
ReLU	$64\times N\times N$	$64\times N\times N$
Conv2D, $k=3,p=1$	$64\times N\times N$	$128\times N\times N$
GroupNorm, $G=16$	$128\times N\times N$	$128\times N\times N$
ReLU	$128\times N\times N$	$128\times N\times N$

Appendix

Table 5: Architecture of the 2DCNN for face geometry. The same architecture is applied to process two face sampling resolutions. The input grid size is generally denoted as

N\times N

, where

N\in\{3,13\}

. The convolution kernel size is

k

, padding is

p

, and groups for GroupNorm are

G

Operator	Input Shape	Output Shape
Conv1D, $k=3,p=1$	$12\times 13$	$64\times 13$
GroupNorm, $G=8$	$64\times 13$	$64\times 13$
ReLU	$64\times 13$	$64\times 13$
Conv1D, $k=3,p=1$	$64\times 13$	$256\times 13$
GroupNorm, $G=32$	$256\times 13$	$256\times 13$
ReLU	$256\times 13$	$256\times 13$

Appendix

Table 6: Architecture of the 1DCNN for edge geometry. The input is a sequence of 13 points sampled along the edge, each with a 12D vector. The convolution kernel size is

k

, the padding is

p

, and the groups for GroupNorm are

G

Operator	Input Shape	Output Shape
Linear	16	128
LayerNorm	128	128
GELU	128	128
Linear	128	128

Appendix

Table 7: Architecture of the MLP for face attribute.

Operator	Input Shape	Output Shape
Linear	9	128
LayerNorm	128	128
GELU	128	128
Linear	128	128

Appendix

Table 8: Architecture of the MLP for edge attribute.

Operator	Input Shape	Output Shape
Linear	384	256
LayerNorm	256	256
GELU	256	256
Linear	256	256

Appendix

Table 9: Architecture of the MLP for feature fusion. This network is used to fuse the geometric features (256D) and attribute features (128D) into a final 256D vector. The input shape is the concatenated dimension (

256+128=384

Hyperparameter	Value
Embedding Dimension	256
Number of Attention Heads	8
FFN Hidden Dimension	1024
FFN Activation	GELU
Dropout Rate	0.3

Appendix

Table 10: Hyperparameters for our Transformer blocks.

Operator	Input Shape	Output Shape
Linear	$256+2$	512
ReLU	512	512
Linear	512	512
ReLU	512	512
Linear	512	64
Linear	$256+64$	512

Appendix

Table 11: Architecture of the FoldingNet decoder for face geometry.

Operator	Input Shape	Output Shape
Linear	256	256
LayerNorm	256	256
GELU	256	256
Linear	256	16

Appendix

Table 12: Architecture of the MLP decoder for face attribute information.

Operator	Input Shape	Output Shape
Linear	$256+1$	512
ReLU	512	512
Linear	512	512
ReLU	512	512
Linear	512	64
Linear	$256+64$	512

Appendix

Table 13: Architecture of the FoldingNet decoder for edge geometry.

Operator	Input Shape	Output Shape
Linear	256	256
LayerNorm	256	256
GELU	256	256
Linear	256	9

Appendix

Table 14: Architecture of the MLP decoder for edge attribute.

Operator	Input Shape	Output Shape
Linear	256	1024
GELU	1024	1024
Dropout	1024	1024
Linear	1024	256
GELU	256	256
Dropout	256	256

Appendix

Table 15: Architecture of the downstream task head. It consists of three linear layers with GELU activations and dropout (rate = 0.3) applied after each hidden layer.

C

denotes the number of output classes. The input is the encoded feature vector produced by our pre-trained encoder. For machining feature recognition and modeling segmentation, each face’s feature is fed independently. For part classification, a global max pooling is applied over all face features to obtain the input.

Figure 8: The pre-training loss convergence curves. The y-axis denotes the logarithmic loss, showing the total loss and its five individual components over the training epochs.

Stage	Parameters (M)	FLOPs (G)	Latency (ms)	Throughput (models/s)
Pre-training	8.15	20.98	25.99	38.47
Downstream	6.09	5.26	10.42	96.01

Appendix

Table 16: Computational cost profiling for the pre-training and downstream stages. Profiling is based on CAD models with an average of 30 faces.

Configuration	0.1%	0.5%	1%	1.5%	2%	3%	100%
w/o Mask attr	84.48	94.15	95.42	96.34	96.60	97.20	99.41
Feature all	76.77	90.86	93.79	94.67	95.72	96.56	99.33
Geometry masked	78.03	91.14	93.91	95.23	95.81	96.71	99.37
Mask ratio(50%)	86.99	96.02	97.18	97.43	97.60	98.02	99.49
Mask ratio(60%)	87.41	96.12	97.19	97.46	97.73	98.05	99.51
Mask ratio(80%)	87.01	95.66	96.86	97.37	97.62	97.96	99.46

Appendix

Table 17: More Ablation study on MFInstSeg under various supervision ratios.

Figure 9: More reconstruction examples with face coordinates and face normals.

Figure 10: More results of machining feature recognition.

Figure 11: More results of modeling segmentation.

3D Representation Learning

3D representation learning builds geometric features for recognizing, reconstructing, generating, and reasoning about objects and scenes.

See all 191 papers

Masked Topology Modeling for Self-Supervised Learning on Parametric CAD
Jul 22, 2026Heinrich Jiang, Jennifer JangGraph SSL Parametric CAD Modeling
StoryGold AI
BrepCoder: A Unified Multimodal Large Language Model for Multi-task B-rep Reasoning
Feb 25, 2026Mingi Kim, Yongjun Kim, Jungwoo Kang +1Multimodal Large Language Models
Chungnam National University, Daejeon, Republic of Korea
BRepCLIP: Contrastive Multimodal Pretraining on BRep Primitives for CAD Understanding
Jun 3, 2026Muhammad Usama, Didier Stricker, Mohammad Sadil Khan +13D Shape Retrieval Multimodal Pretraining
1DFKI, Germany · 2RPTU Kaiserslautern-Landau, Germany

Jul 22, 2026J/K move · Enter open · S save

Masked Topology Modeling for Self-Supervised Learning on Parametric CAD

Heinrich Jiang, Jennifer Jang

StoryGold AI

Graph SSL Parametric CAD Modeling Boundary Representation Modeling Topological Representation Learning

Open paper AskPDF arXiv

Computer aided design (CAD) is ubiquitous: virtually any modern object was designed using editable CAD tools. However, with the shortage of available CAD datasets in its native editable and parametric format, boundary representation (B-Rep), it is ever more important to develop data-efficient methods for this domain. We present a new self-supervised pretraining task, Masked Topology Modeling (MTM), that leverages the face-adjacency graph, an induced structure unique to B-reps that the encoder can be asked to reconstruct. MTM masks a fraction of edges and trains a small head to predict each masked edge's convexity and curve type from the encoder's post-message-passing face features. We combine MTM with a MoCo-style momentum-queue contrastive learning over B-rep-aware augmentations, a BFS-connected face-region masked-reconstruction objective, and pretraining on the ABC dataset and our new procedurally generated dataset to show strong performance on a number of benchmarks.

Abstract

3D Representation Learning

Explore similar work

Masked Topology Modeling for Self-Supervised Learning on Parametric CAD

BrepCoder: A Unified Multimodal Large Language Model for Multi-task B-rep Reasoning

BRepCLIP: Contrastive Multimodal Pretraining on BRep Primitives for CAD Understanding