PMT-MAE: Dual-Branch Self-Supervised Learning with Distillation for Efficient Point Cloud Classification

📄 arXiv: 2409.02007v2 📥 PDF

作者: Qiang Zheng, Chao Zhang, Jian Sun

分类: cs.CV

发布日期: 2024-09-03 (更新: 2024-09-16)

💡 一句话要点

PMT-MAE：双分支自监督学习与蒸馏，高效点云分类。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 点云分类 自监督学习 Transformer MLP 知识蒸馏 双分支网络 特征提取

📋 核心要点

  1. 现有方法在点云特征提取和理解方面存在不足，需要更有效的自监督学习方法。
  2. PMT-MAE采用双分支结构，结合Transformer和MLP的优势，并通过蒸馏学习更强的特征表示。
  3. 实验表明，PMT-MAE在ModelNet40上取得了93.6%的准确率，优于Point-MAE和Point-M2AE，且训练效率高。

📝 摘要（中文）

本文提出了一种名为PMT-MAE（Point MLP-Transformer Masked Autoencoder）的新型自监督学习框架，用于点云分类。PMT-MAE采用双分支架构，集成了Transformer和MLP组件以捕获丰富的特征。Transformer分支利用全局自注意力进行复杂的特征交互，而并行的MLP分支通过共享的全连接层处理tokens，提供互补的特征转换路径。然后，融合机制结合这些特征，增强模型学习全面3D表示的能力。在更复杂的教师模型Point-M2AE的指导下，PMT-MAE采用了一种蒸馏策略，包括预训练期间的特征蒸馏和微调期间的logit蒸馏，确保有效的知识转移。在ModelNet40分类任务上，PMT-MAE在不采用投票策略的情况下，达到了93.6%的准确率，超过了基线Point-MAE（93.2%）和教师Point-M2AE（93.4%），突显了其学习判别性3D点云表示的能力。此外，该框架展示了高效率，预训练和微调仅需40个epochs。PMT-MAE的有效性和效率使其非常适合计算资源有限的场景，使其成为实际点云分析的有前景的解决方案。

🔬 方法详解

问题定义：论文旨在解决点云分类任务中，如何更有效地利用自监督学习方法提取判别性特征的问题。现有方法，如Point-MAE，在特征提取能力和训练效率上仍有提升空间，尤其是在计算资源受限的情况下。

核心思路：论文的核心思路是设计一个双分支的自监督学习框架，结合Transformer的全局建模能力和MLP的局部特征提取能力，并通过知识蒸馏，将更复杂的教师模型的知识迁移到学生模型，从而提高模型的性能和效率。

技术框架：PMT-MAE的整体架构包含预训练和微调两个阶段。在预训练阶段，输入点云经过masking后，分别进入Transformer分支和MLP分支。Transformer分支利用自注意力机制学习全局特征，MLP分支通过共享全连接层学习局部特征。然后，通过融合机制将两个分支的特征进行融合。在微调阶段，利用预训练好的模型进行点云分类任务。整个框架采用Point-M2AE作为教师模型，在预训练阶段进行特征蒸馏，在微调阶段进行logit蒸馏。

关键创新：PMT-MAE的关键创新在于双分支架构和蒸馏策略的结合。双分支架构能够同时利用Transformer和MLP的优势，从而提取更丰富的特征。蒸馏策略能够有效地将教师模型的知识迁移到学生模型，从而提高模型的性能和效率。

关键设计：在网络结构方面，Transformer分支采用标准的Transformer encoder结构，MLP分支采用共享的全连接层。在损失函数方面，预训练阶段采用masked autoencoder的重建损失和特征蒸馏损失，微调阶段采用交叉熵损失和logit蒸馏损失。预训练和微调均采用40个epochs。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

PMT-MAE在ModelNet40点云分类任务上取得了显著的成果，在不使用投票策略的情况下，达到了93.6%的准确率，超过了Point-MAE (93.2%) 和 Point-M2AE (93.4%)。此外，该模型仅需40个epochs进行预训练和微调，展示了其高效性。

🎯 应用场景

PMT-MAE适用于各种需要高效点云分析的场景，例如自动驾驶中的环境感知、机器人导航、三维重建、以及工业检测等。其高效性和高性能使其在计算资源有限的边缘设备上部署成为可能，具有广阔的应用前景。

📄 摘要（原文）

Advances in self-supervised learning are essential for enhancing feature extraction and understanding in point cloud processing. This paper introduces PMT-MAE (Point MLP-Transformer Masked Autoencoder), a novel self-supervised learning framework for point cloud classification. PMT-MAE features a dual-branch architecture that integrates Transformer and MLP components to capture rich features. The Transformer branch leverages global self-attention for intricate feature interactions, while the parallel MLP branch processes tokens through shared fully connected layers, offering a complementary feature transformation pathway. A fusion mechanism then combines these features, enhancing the model's capacity to learn comprehensive 3D representations. Guided by the sophisticated teacher model Point-M2AE, PMT-MAE employs a distillation strategy that includes feature distillation during pre-training and logit distillation during fine-tuning, ensuring effective knowledge transfer. On the ModelNet40 classification task, achieving an accuracy of 93.6\% without employing voting strategy, PMT-MAE surpasses the baseline Point-MAE (93.2\%) and the teacher Point-M2AE (93.4\%), underscoring its ability to learn discriminative 3D point cloud representations. Additionally, this framework demonstrates high efficiency, requiring only 40 epochs for both pre-training and fine-tuning. PMT-MAE's effectiveness and efficiency render it well-suited for scenarios with limited computational resources, positioning it as a promising solution for practical point cloud analysis.