Masked BRep Autoencoder via Hierarchical Graph Transformer
这篇论文介绍的模型架构名为Masked BRep Autoencoder (MBRE),它是一种专为 CAD 模型(边界表示,BRep)设计的自监督学习框架。其核心是一个分层图 Transformer (Hierarchical Graph Transformer)架构。
该架构的设计初衷是为了解决 CAD 模型中存在的极端尺度差异(例如巨大的平面与微小的圆角共存)以及数据冗余问题。
模型整体架构
该模型由编码器 (Encoder)和解码器 (Decoder)两大部分组成,采用了类似于 MAE (Masked Autoencoder) 的自监督学习逻辑,通过重建被遮掩(Masked)的几何信息来学习特征。
1. 编码器 (Encoder)
编码器分为两个阶段,将原始的 CAD 几何数据逐步转化为高层的语义表示:
- BRep 编码器 (BRep Encoder):
- 将 BRep 模型转化为gAAG (Geometric Attributed Adjacency Graph),即带有几何属性的邻接图。
- 双分辨率策略:针对每个面(Face),同时采样3 × 3 3 \times 33×3(全局语义)和13 × 13 13 \times 1313×13(局部细节)两种分辨率的 UV 网格。
- 使用并行的2D CNN处理面几何,MLP处理面属性;使用1D CNN处理边(Edge)几何。
- 图编码器 (Graph Encoder):这是核心部分,结合了全局与局部学习:
- 跨尺度互注意力模块 (CSMA Block):使用低分辨率特征作为查询(Query),从高分辨率特征中提取关键几何变化,捕捉全局长程依赖。
- 局部拓扑消息传递 (MPNN):在 Transformer 处理完全局信息后,使用消息传递神经网络聚合邻域的拓扑连接关系。
2. 解码器 (Decoder)
为了降低重建难度,解码器采用了两阶段设计:
- 图解码器 (Graph Decoder):首先在图空间中重建中间节点和边的特征(使用 MPNN 结构)。
- BRep 解码器 (BRep Decoder):
- 包含四个并行分支,将图特征映射回显式的几何和属性。
- 几何分支:采用FoldingNet结构(一种专门用于变形点云的解码器),将潜变量恢复为 3D 坐标、法向量等。
- 属性分支:使用 MLP 回归预测面的类型、面积、质心等离散属性。
核心技术特点
| 组件 | 功能描述 |
|---|---|
| Masking 策略 | 在输入层直接遮掩 70% 的原始几何和属性,强制模型根据上下文推断缺失结构。 |
| CSMA 模块 | 解决 CAD 冗余问题,防止模型被大面积的平坦表面“淹没”,专注关键特征。 |
| 双分辨率采样 | 显式分离全局形状信息与局部精细结构。 |
| 两阶段重建 | 通过中间特征监督和显式几何重建(5 项 Loss 联合优化)增强学习稳定性。 |
“我们的层级图 Transformer 编码器通过跨尺度互注意力(CSMA)瓶颈无缝集成了全局几何上下文,并通过显式的局部消息传递保留了关键的局部细节。” Conclusion
要将这个预训练模型扩展到零件识别(Part Classification),实际上是完成从“理解局部零件”到“总结整体身份”的转变。
按照论文中的做法,可以分为以下四个具体步骤:
1. 提取全局特征(Global Feature Aggregation)
编码器最初输出的是每个面(Face)的特征。对于零件识别这种“整体性”任务,需要把所有面的特征“揉”在一起,形成一个代表整个零件的向量:
- 做法:将编码器输出的所有面特征进行全局池化(Global Pooling),例如取平均值或最大值。
- 意义:这个全局向量就像零件的“指纹”,包含了模型整体的几何分布和拓扑结构信息。
2. 挂载分类头(Task-specific Head)
在预训练好的编码器后面,接一个专门负责分类的子网络:
- 结构:通常是一个简单的多层感知机(MLP)。
- 输出:最后一层的大小对应你想要识别的零件种类数量(比如 26 类字母零件或各种工业标准件)。
3. 监督微调(Fine-tuning)
利用带标签的零件数据集(如论文中提到的SolidLetter)对模型进行训练:
- 两阶段更新:
- 分类头:使用较大的学习率,从头开始学习如何根据特征下结论。
- 预训练编码器:使用极小的学习率进行微调,或者在数据极少时直接冻结(Frozen)。
- 损失函数:使用标准的分类型损失函数(如 Cross-Entropy Loss)。
4. 实验证明的效果
根据论文的实验数据,这种扩展方式非常高效:
- 高精度:在 SolidLetter 26 类字母零件分类任务中,准确率达到了97.91%。
- 泛化性:即使在预训练时完全没见过某种类型的零件,仅靠少量的标注样本(如 10-shot 或 20-shot),模型也能快速学会识别。
性能表现参考
您可以参考下表了解其在不同标注数据量下的表现(以 SolidLetter 数据集为例):
| 模型 | 10-shot 准确率 | 20-shot 准确率 | 全量数据准确率 |
|---|---|---|---|
| BRep-BERT | 68.71% | 75.92% | 97.76% |
| BRepFormer | 76.53% | 84.25% | 97.59% |
| Ours (本模型) | 82.51% | 89.58% | 97.91% |
“在微调阶段,我们训练一个新的网络,该网络是通过在编码器后面连接一个特定任务的头部形成的,并使用带有不同损失函数的少量标签数据进行训练。” Method