Point-BERT: Pre-training 3D Point Cloud Transformers with Masked Point Modeling

Point-BERT : Pre-trainning 3D Point Cloud Transformers with Masked Point Modeling

Xumin Yu*,1, Lulu Tang∗,1,2, Yongming Rao∗,1, Tiejun Huang2,3, Jie Zhou1, Jiwen Lu†,1,2

 

https://arxiv.org/abs/2111.14819

Point-BERT: Pre-training 3D Point Cloud Transformers with Masked Point Modeling

2023 CVPR

Abstract

 

• BERT의 개념을 3D point cloud로 확장한 새로운 패러다임인 Point-BERT
• 포인트 클라우드 transformer를 사전 학습하기 위한 Masked Point Modelling(MPM) 작업을 고안
  • 포인트 클라우드를 여러 개의 local point patch로 나눈 다음, 의미 있는 local 정보를 포함하는 토큰을 생성하기 위해 dVAE가 포함된 토크나이저 설계
  • 입력 포인트 클라우드의 일부 패치를 무작위로 마스킹하고 이를 backbone 트랜스포머에 입력
• BERT 스타일의 사전 학습 전략이 포인트 클라우드 트랜스포머의 성능을 크게 향상시킴

 

Introduction

 

• 트랜스포머 구조가 NLP와 이미지 처리에서 좋은 결과를 보여 주고 있지만, 3D point cloud의 도메인에서는 아직 충분히 연구되지 않았다. 기존의 트랜스포머 기반 포인트 클라우드 모델들은 local feature aggregation과 neighbor embedding에서 특정한 inductive bias를 도입하곤 한다. (이게 표준이 아님)
• 트랜스포머의 직접적인 적용은 포인트 클라우드 작업에서 만족스러운 성능을 달성하지 못 한다. 
  • 보통 limited annotated 3D data 때문에
  • 귀납적 편향이 없는 순수한 트랜스포머는 대량의 훈련 데이터가 필요하다. 
    • 이미지는 그런 대규모 데이터셋이 있지만, 포인트 클라우드는 그런 대량 데이터셋이 상대적으로 부족하다. 
  • 이러한 어려움은, 레이블이 없는 3D 데이터로부터 학습하는 연구를 촉진시켰다. 
• BERT를 포인트 클라우드에 적용할 때 어려운 점은, 포인트 클라우드를 위한 사전 정의된 vocabulary가 없다는 것이다. 
  • 문장의 단어는 기본적인 문맥적 의미 정보를 포함하지만, 포인트 클라우드의 단일 포인트는 거의 의미를 포함하지 않는다. 

Point-BERT

 

• 본 논문에서 하고자 하는 것은 BERT 스타일의 pre-training strategy를 point cloud에서 효과적으로 사용할 수 있도록 하는 전략을 수립하는 것이다. 
  • * 여기서 BERT 스타일 pre-training strategy가 의미하는 것 : 
    •  bidirectional 
      • 기존 모델들이 텍스트를 왼쪽 -> 오른쪽으로 읽었던 것과 달리 BERT는 문장의 앞뒤 문맥을 동시에 파악할 수 있다. 
    • Masked Language Modeling
    • Next Sentence Prediction
    • Transfer Learning
      • 대량의 텍스트로 사전 학습 후 특정 작업에 맞게 fine-tuning하여 사용
  • 따라서 MLM strategy를 point cloud에 맞게 모방한 Masked Point Modeling을 사용했다. 
• Point Tokenization
  • Point Embeddings
    • naive한 접근 방법은 하나의 포인트를 하나의 토큰으로 치환하는 것인데, 이러면 computational cost가 너무 크다. 
    • Vision Transformer의 patch embedding strategy에서 영감을 받아, point cloud를 sub-cloud로 바꿔서 그룹화하는 방법을 채택했다. 
    • 입력 포인트 클라우드가 주어지면, FPS(Farthest Point Sampling)을 통해 전체 포인트 클라우드 p에서 g개의 중심점을 샘플링한다. 
    • 이후 K-NN알고리즘을 사용하여 각 중심점에 대해 n개의 가장 가까운 점을 선택하여, g개의 sub cloud를 그룹화한다. 
    • 지역 패치들에서 중심 좌표를 빼서 편향되지 않도록 하여, 지역 패치의 구조적인 패턴과 공간 좌표를 분리한다. 
    • 이러한 sub cloud는 NLP의 단어나 비전 도메인의 이미지 패치처럼 취급될 수 있다. 
    • 우리는 이런 sub cloud를 포인트 임베딩으로 투영하기 위해 mini-pointnet을 추가로 채택했다. 
  • Point Tokenizer
    • 포인트 토크나이저는 포인트 임베딩을 입력으로 받아 이산 포인트 토큰으로 변환한다. 
    • 토크나이저는 포인트 임베딩을 이산 포인트 토큰으로 매핑하고, 여기서 토큰들은 학습된 어휘에 해당한디ㅏ. 
    • 본 논문에서는 DGCNN이 토크나이저 네트워크로 사용된다. 
  • Point Cloud Reconstruction
    • dVAE의 디코더는 포인트 토큰을 입력으로 받아 해당하는 서브 클라우드를 재구성하는 것을 학습한다. 
    • 지역적인 기하학 구조만 가지고는 reconstruction을 하기가 어렵기 때문에, 우리는 이웃하는 포인트 토큰들과의 관계를 구축하기 위해 DGCNN을 채택한다. 
    • 이후 FoldingNet이 서브 클라우드를 재구성하는 데 사용된다. 
  • Transformer Backbone
    • 표준 트랜스포머
  • 따라서 Point-BERT에서는 크게 두 단계의 학습이 진행된다. 
    • 토크나이저 (dVAE) 학습
      • 입력 : 포인트 패치들
      • 학습 과정
        • 인코더 : 패치 -> 잠재 벡터
        • 코드북 : 가장 가까운 코드북 벡터 선택
        • 디코더 : 선택된 토큰으로부터 원본 패치 재구성
      • 손실 함수 
        • reconstruction loss : 원본과 재구성된 패치 간의 차이
    • 트랜스포머 사전 학습
      • MPM 

 

*dVAE (discrete Variational AutoEncoder)

• 일반적인 VAE
  • 입력 -> 연속적인 잠재 벡터 -> 출력
• dVAE
  • 입력 -> 이산적인 잠재 벡터 -> 출력
• Point-BERT에서의 dVAE 구조
  • 입력 : 포인트 패치
  • 인코더 (DGCNN)
  • 이산화 과정 (Gumbel-Softmax)
    • 이산화 과정을 미분가능하게 만듦
  • 이산 토큰 (코드북에서 가장 가까운 인덱스 선택)
  • 디코더 (FoldingNet)
  • 출력 : 재구성된 포인트 패치

 

Experiments

 

• Pre-training Setups
  • Data Setups
    • Pre-traning 데이터셋으로 ShapeNet을 사용했음
    • 각 3D 모델에서 1024개의 점을 샘플링하고 이를 64개의 포인트 패치(sub cloud)로 나누었다. 
    • 각 패치는 PointNet을 이용해 64개의 포인트 임베딩으로 바뀌어 이는 dVAE와 트랜스포머 모두의 입력으로 사용된다. 
  • dVAE Setups
    • 입력 포인트 클라우드의 내부 구조를 모델링하고 패치 간의 관계를 학습하기 위해 4-layer DGCNN을 사용한다. 
    • Vocabulary size는 8192
    • dVAE의 성능은 하이퍼파라미터에 민감하기 때문에, 이미지 기반 dVAE의 구성을 직접 사용할 수 없다. 
  • MPM
    • 입력 포인트 임베딩의 25%~45%가 무작위로 마스킹된다. 
    • MoCo
      • 메모리 뱅크, temperature, momentum 설정 ~

 

Downstream Tasks

 

• Object Classification
  • ModelNet40
• Few-shot 학습
  • K-way N-shot 
    • K개의 클래스가 먼저 무작위로 선택되고, 각 클래스에 대해 (N+20) 개의 객체가 샘플링된다. 
• Part Segmentation
  • 모든 점에 대해 더 세분화된 클래스 레이블을 예측하는 태스크

 

Conclusion and Discussions

 

• BERT 스타일의 사전 학습을 통해 저수준의 구조적인 정보와 고수준의 의미적 특징을 모두 학습하는 3D 포인트 클라우드 트랜스포머를 잘 학습할 수 있다. 
• 다만 다른 트랜스포머 pre-training 방법들처럼 전체 'pretrain + fine tuning' 절차는 상당한 시간이 소요된다.