SAM-Guided Masked Token Prediction for 3D Scene Understanding

Foundation models have significantly enhanced 2D task performance, and recent works like Bridge3D have successfully applied these models to improve 3D scene understanding through knowledge distillation, marking considerable advancements. Nonetheless, challenges such as the misalignment between 2D and 3D representations and the persistent long-tail distribution in 3D datasets still restrict the effectiveness of knowledge distillation from 2D to 3D using foundation models. To tackle these issues, we introduce a novel SAM-guided tokenization method that seamlessly aligns 3D transformer structures with region-level knowledge distillation, replacing the traditional KNN-based tokenization techniques. Additionally, we implement a group-balanced re-weighting strategy to effectively address the long-tail problem in knowledge distillation. Furthermore, inspired by the recent success of masked feature prediction, our framework incorporates a two-stage masked token prediction process in which the student model predicts both the global embeddings and the token-wise local embeddings derived from the teacher models trained in the first stage. Our methodology has been validated across multiple datasets, including SUN RGB-D, ScanNet, and S3DIS, for tasks like 3D object detection and semantic segmentation. The results demonstrate significant improvements over current State-of-the-art self-supervised methods, establishing new benchmarks in this field.

Paper

Similar papers

Peer review

Reviewer icGN5/10 · confidence 5/52024-07-01

Summary

This paper proposes a two-stage SAM-guided pre-training method for 3D scene understanding. The authors present a group-balanced re-weighting method for long-tail representation distillation and a SAM-guided tokenization method to seamlessly align 2D and 3D region-level features. Extensive experiments on various downstream tasks show the effectiveness of the proposed pre-training method.

Strengths

1. Adopt SAM to guide pre-training makes sense, since the semantic information can help model to learn high-quality patterns. 2. The results are good, surpassing previous state-of-the-art methods significantly.

Weaknesses

1. For 3D object detection results, even with the proposed pre-training method, the model's overall mAP is lower than SOTA 3D detector like CAGroup3D and VDETR. Can this work be applied on these more advanced architectures? If so, supporting experimental results are required. 2. Using SAM to guide 3D pretrained has already been explored in SEAL [1], more comparison or discussion is required. 3. This paper should not only compare with transformer-based pretraining methods, but also other state-of-the-art pretrianing techniques like PPT [2]. [1] Segment Any Point Cloud Sequences by Distilling Vision Foundation Models, NeurIPS 2023 [2] Towards Large-scale 3D Representation Learning with Multi-dataset Point Prompt Training, CVPR 2024

Questions

It seems the models are pretrained on independent RGB-D frames, and then finetuned on reconstructed room-level point clouds. There may be some domain gap between the single-view point clouds and the complete room-level ones. I wonder is this work applicable to the online 3D perception setting [1, 2], which is more relevant to RGB-D perception and is a more valuable setting. [1] Fusion-aware point convolution for online semantic 3d scene segmentation, CVPR 2020 [2] Memory-based Adapters for Online 3D Scene Perception, CVPR 2024

Rating

5

Confidence

5

Soundness

3

Presentation

3

Contribution

2

Limitations

Limitations is discussed.

Reviewer DTv45/10 · confidence 4/52024-07-12

Summary

This paper introduces a novel method to enhance 3D scene understanding by addressing misalignment between 2D and 3D representations and the long-tail distribution in 3D datasets. The proposed approach involves a SAM-guided tokenization method for seamless alignment and a group-balanced re-weighting strategy to handle representation imbalance. It also incorporates a two-stage masked token prediction process for effective knowledge distillation from 2D foundation models to 3D networks.

Strengths

The SAM-guided tokenization method effectively aligns 2D and 3D representations, overcoming the limitations of traditional KNN-based tokenization techniques. The introduction of a group-balanced re-weighting strategy addresses the long-tail distribution problem in 3D datasets, improving the representation of under-represented samples.

Weaknesses

Mask generation is an important step of the proposed method, probably a visual comparison with current SOTA methods would be more helpful to justify its effectiveness. There are typos in this manuscript, and the reference style is very confused (e.g., Bridge3D (10)).

Questions

What is training time for fine-tuning? How about memory consumption compared to Bridge3D? Bridge3D already mentioned in their paper that one of their limitations is that the current work primarily focuses on indoor 3D scene understanding. Is the same limitation applied for this paper?

Rating

5

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

Can this method apply to outdoor scene understanding?

Reviewer EQZo7/10 · confidence 3/52024-07-13

Summary

The paper proposes a 3D transformer tokenization technique to align 3D representations with 2D ones, distilling from 2D pre-trained knowledge from SAM. The method achieves favorable performance on 3D object detection and semantic segmentation compared to prior self-supervised learning methods.

Strengths

* The paper provides a good explanation contrasting the proposed method with prior works that develop 3D representations distilling from 2D foundation models, highlighting the drawbacks of KNN-based point tokenization which well motivates the proposed solution. * Several techniques, such as group-balanced re-weighting and two-stage teacher-forcing training, are adopted in the framework and ablations show their effectiveness. * The performance is strong, with advantages compared to prior self-supervised methods.

Weaknesses

* It is not discussed why SAM is chosen as the distillation source for tokenization, instead of other 2D foundation models such as DINO, or even generative ones such as Stable Diffusion. * The representation is evaluated on the task of 3D object detection and 3D semantic segmentation, which intuitively seem to correspond well to the strength of SAM as SAM may have implicitly learned the notion of objects and object semantics during its training phase. Are there tasks where the proposed representation could be useful?

Questions

* What's the input to $F_{2D, i}$ and $F_{3D, i}$ in Eq (1)? Where is $i$ sampled from? * $1/O_i$ in Eq (1) should probably be written as $1/|O_i|$. * What is $n_{min}$ and $n_{max}$ in line 209?

Rating

7

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

Limitations are discussed.

Area Chair Bw8K2024-08-12

Discussion

Thank you for being a reviewer for NeurIPS2024, your service is invaluable to the community! The authors have submitted their feedback. Could you check the rebuttal and other reviewers' comments and start a discussion with the authors and other reviewers? Regards, Your AC

Reviewer KyqU5/10 · confidence 5/52024-07-13

Summary

The paper proposes a self-supervised method for understanding 3D scenes by predicting the 3D mask of the point cloud. The masks are initialized from Segment Anything (SAM), followed by two stages of the knowledge distillation framework to train the 3D teacher and student networks. The method is evaluated on SUN RGB-D, ScanNet, and S3DIS datasets.

Strengths

The paper is generally well-written and easy to follow.

Weaknesses

- The first concern is the method's innovation. Other methods, like Seal [P1], have already proposed the idea of distilling the knowledge from 2D foundation models into the 3D network for mask prediction. However, it is not discussed or compared. - The insight into method design is not elaborated. It is not clear why two-stage training is necessary for self-supervised learning. How about three-stage, four-stage, or momentum updating proposed in MeanTeacher? It is better to present more insight and intuitive explanations. Besides, it is difficult to understand why group-balanced re-weighting will work since the pseudo labels are also extremely noisy, i.e., it also suffers from long-tail issues. - More recent methods should be comparied to verify the efficientness of the proposed method. For example, CLIP2Scene is disscussed but not comparied. Seal [P1], also utilize 2D models mask prediction to regularize 3D network. P1. Segment Any Point Cloud Sequences by Distilling Vision Foundation Models, NeurIPS 2023.

Questions

refer to the weaknesses

Rating

5

Confidence

5

Soundness

2

Presentation

3

Contribution

2

Limitations

yes

Reviewer icGN2024-08-08

The authors' rebuttal has solved most of my concerns. I suggest them to add more comparison and apply the proposed method to more architecture rather than only transformers. Also, the authors can include some visualization and demo on the online 3D perception setting, which shows great application potential.

Authorsrebuttal2024-08-09

We sincerely appreciate your positive feedback. Your suggestions are invaluable in helping us further refine and enhance the quality of our paper. Regarding your suggestion to include more comparisons and apply the proposed method to a broader range of architectures beyond transformers, we conducted a related experiment in Table 1 of the rebuttal. Specifically, we applied our method to the CAGroup3D architecture, which has a different backbone design. However, it is important to note that a primary motivation behind our method is to address the misalignment between 2D and 3D representations caused by the traditional KNN tokenization method specifically within transformer structures. In future work, we plan to explore a more universal approach that better aligns 2D and 3D representations across all architectures. Additionally, we agree that the inclusion of more visualization results would enhance the demonstration of our method's potential in online perception settings. We will incorporate additional visualization results in the revised version to showcase these applications better.

Reviewer DTv42024-08-11

Thanks for the rebuttal, most of my concerns are properly addressed.

Authorsrebuttal2024-08-11

We sincerely appreciate your positive feedback. Your suggestions are invaluable in helping us further refine and enhance the quality of our paper.

Reviewer KyqU2024-08-12

reply to authors

Thanks for the rebuttal, most of my concerns are addressed.

Authorsrebuttal2024-08-12

We sincerely appreciate your positive feedback. Your suggestions are invaluable in helping us further refine and enhance the quality of our paper.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC