[论文解读] Can We Solve 3D Vision Tasks Starting from A 2D Vision Transformer?

Yi Wang, Zhiwen Fan+3 more|arXiv (Cornell University)|Sep 15, 2022

Advanced Neural Network Applications被引用 5

一句话总结

本文提出 Simple3D-Former，一种极简主义的 3D 视觉 Transformer 模型，通过膨胀 patch 和位置嵌入并修改输出头层，将标准 2D ViT 适配至 3D 视觉任务，无需架构重设计，即在 3D 分类、分割和检测任务上实现强劲性能。实验表明，预训练的 2D ViT 权重显著提升 3D 泛化能力，尤其在低数据设置下表现突出，建立了一个强大且可迁移的基线模型。

ABSTRACT

Vision Transformers (ViTs) have proven to be effective, in solving 2D image understanding tasks by training over large-scale image datasets; and meanwhile as a somehow separate track, in modeling the 3D visual world too such as voxels or point clouds. However, with the growing hope that transformers can become the "universal" modeling tool for heterogeneous data, ViTs for 2D and 3D tasks have so far adopted vastly different architecture designs that are hardly transferable. That invites an (over-)ambitious question: can we close the gap between the 2D and 3D ViT architectures? As a piloting study, this paper demonstrates the appealing promise to understand the 3D visual world, using a standard 2D ViT architecture, with only minimal customization at the input and output levels without redesigning the pipeline. To build a 3D ViT from its 2D sibling, we "inflate" the patch embedding and token sequence, accompanied with new positional encoding mechanisms designed to match the 3D data geometry. The resultant "minimalist" 3D ViT, named Simple3D-Former, performs surprisingly robustly on popular 3D tasks such as object classification, point cloud segmentation and indoor scene detection, compared to highly customized 3D-specific designs. It can hence act as a strong baseline for new 3D ViTs. Moreover, we note that pursing a unified 2D-3D ViT design has practical relevance besides just scientific curiosity. Specifically, we demonstrate that Simple3D-Former naturally enables to exploit the wealth of pre-trained weights from large-scale realistic 2D images (e.g., ImageNet), which can be plugged in to enhancing the 3D task performance "for free".

研究动机与目标

• 探究标准 2D 视觉 Transformer 是否可通过最小化架构修改适配至 3D 视觉任务。
• 评估将预训练的 2D ViT 权重迁移至 3D 任务以提升泛化能力的有效性。
• 建立统一的 2D-3D ViT 框架，实现跨领域的知识迁移。
• 证明简单、非定制化的 Transformer 可与高度工程化的 3D 专用模型性能相当或更优。
• 探索使用单一 Transformer 主干网络同时支持 2D 和 3D 视觉任务，并通过共享预训练实现的可能性。

提出的方法

• 将 2D patch 嵌入层扩展以处理 3D 体素或点云数据，将 3D patch 视为类似 2D 的 token 处理。
• 设计一种 3D 感知的位置编码，以保持 3D 空间维度中的几何结构。
• 将 2D 分类头替换为针对对象分类、分割和检测等任务定制的 3D 特异性头。
• 利用预训练的 2D ViT ImageNet 权重作为初始化，将先验知识注入 3D 模型。
• 通过知识蒸馏从教师 ViT 中保留 2D 泛化能力，同时在 3D 数据上微调。
• 端到端训练模型，结合交叉熵损失与蒸馏损失，以在低数据条件下维持性能。

实验结果

研究问题

• RQ1 标准 2D 视觉 Transformer 是否可通过最小架构修改有效适配至 3D 视觉任务？
• RQ2 在大规模 2D 数据（如 ImageNet）上预训练在多大程度上能提升 3D 视觉性能？
• RQ3 从 2D 教师 ViT 中进行知识蒸馏是否能增强 3D 泛化能力，尤其是在低数据设置下？
• RQ4 统一的 Transformer 主干网络是否可通过共享预训练同时服务于 2D 和 3D 视觉任务？
• RQ5 极简主义 3D ViT 的性能与高度定制化的 3D 专用架构相比如何？

主要发现

• Simple3D-Former 在 3D 对象分类任务上表现优异，在 ModelNet40 上达到 84.5% 的准确率，在 ShapeNet Part Segmentation 上达到 83.1%，优于未使用 2D 预训练的基线模型。
• 在 SUN RGB-D 3D 检测基准上，Simple3D-Former 实现 57.6% mAP@25 和 32.0% mAP@50，尽管架构改动极少，性能与 3DETR 和 H3DNet 相当。
• 使用 2D 预训练使 ModelNet40 上性能提升 0.7%，ShapeNet Part Segmentation 上提升 1.5%，证明了 2D 先验知识的价值。
• 当结合来自 2D 教师 ViT 的知识蒸馏时，ModelNet40 上性能进一步提升至 88.0%，表明遗忘问题得到缓解，泛化能力增强。
• 在低数据设置（训练数据仅 25%）下，2D 预训练使 ShapeNet 上 3D 部分分割的 mIoU 最高提升 1.6%，表明其具备强大数据效率。
• 消融实验确认，预训练与蒸馏均对最优性能至关重要，当两者同时使用时，ModelNet40 上达到最佳结果（88.0%）