[CVPR 2024] OneFormer3D: One Transformer for Unified Point Cloud Segmentation

안녕하세요, 스물 여덟 번째 x-review 입니다. 이번 논문은 2024년도 OneFormer3D: One Transformer for Unified Point Cloud Segmentation 입니다. 해당 논문은 3D Segmentation 논문인데, segmentation을 위해 설계된 모델임에도 불구하고 3D detection에서 SOTA에 준하는 성능을 달성하여 segmentation 모듈에서 어떤 부분이 detection에 도움을 줄 수 있었던 건지 궁금하여 읽게 되었습니다. 그럼 바로 리뷰 시작하겠습니다 !

1. Introduction

포인트 클라우드에 대한 3D segmentation은 전체 포인트 중에서 의미있는 영역으로 나누어 포인트를 그룹화하게 됩니다. 의미있다고 함은 동일 semantic한 카테고리에 속하는 포인트이거나 혹은 instance, 즉 같은 하나의 물체에 속하는 포인트로 segmentation 될 수 있겠죠. segmentation은 크게 semantic, instance, 그리고 panoptic 이렇게 3가지 task로 나눌 수 있습니다. 먼저 semantic segmentation은 각각의 포인트가 semantic한 라벨이 할당되도록 semantic 카테고리에 대한 마스크를 출력값으로 가지게 됩니다. 반면 instance segmentation은 각각의 물체에 대한 마스크를 가지는데, 바닥이나 천장과 같은 영역은 특정 물체로 구분할 수 없기 때문에 배경 영역으로 간주됩니다. 즉 물체로 구분할 수 있는 영역의 포인트에만 라벨링이 되는 것이죠. 마지막으로 panoptic segmentation이 앞선 두 segmentation을 합친 개념의 task로 전경에 해당하는 각각의 물체(thing)에 대해서는 마스크를 예측하고 배경에 해당하는 포인트(stuff)에는 semantic한 라벨을 예측합니다. 어떤 기준으로 포인트에 대한 마스크를 예측할지는 조금씩 다르지만 어쨌든 3개의 task 모두 마스크를 예측한다는 동일한 task임에도 불구하고 실질적으로 문제를 해결하는 모델 구조가 완전히 다르다고 합니다. semantic segmentation은 보통 U-Net 구조를 사용하며, instance segmentation은 semantic segmentation 모델을 클러스팅이나, object detection, 혹은 트랜스포머 디코더 구조에 기반한 aggregation 알고리즘을 통해 합치게 됩니다. panoptic segmentation은 2D 이미지에서 우선 segmentation을 수행한 다음에 예측된 마스크를 3차원 공간으로 올려서 포인트 단위로 합친다고 합니다. 저자는 이렇게 각기 달리 수행되는 3가지 segmentation task를 하나의 통합된 방식으로 처리할 수는 없는지에 대한 의문을 가졌습니다. 왜냐하면 2D segmentation에서는 OneFormer라는 방법론이 등장하여 3개의 segmentation task에서 하나의 통합된 구조로 모두 SOTA를 달성할 수 있었기 때문 입니다. OneForemr가 등장하기 이전에도 역시 통합된 구조가 존재하였지만 SOTA에 준하는 성능을 보이지는 못하였으며 여전히 각각의 task에 대해 개별적으로 학습하였을 때 더 좋은 결과를 보였다고 합니다. 하지만 그러한 개별적인 학습은 3배 더 많은 메모리와 시간이 소모되기 때문에 완전히 통합된 구조라고 하기는 어렵습니다. 그렇기 때문에 저자는 3차원 segmentation에서도 통합된 구조의 중요성을 강조하며 2D에서의 흐름을 따라 3D에서도 역시 통합된 구조를 설계하고자 하였고, 그러한 구조가 본 논문에서 제안하는 OneFormer3D 입니다. OneFormer3D는 이전 instance segmentation의 SPFormer 구조의 상당 부분을 그대로 가지고 가기 때문에 필요하다면 method에서 SPFormer 구조에 대해서도 덧붙이면서 설명하려고 합니다. 간단하게, 트랜스포머 디코더에서 instance 쿼리만 사용하던 SPFormer에 추가적으로 semantic 쿼리를 사용하였으며 트랜스포머 기반의 instance segmentation의 불안정한 성능을 새로운 쿼리 선택 알고리즘과 매칭 기법으로 안정화하고자 하였습니다. 이를 통해 3개의 segmentation task에서 모두 SOTA를 달성하는 결과를 얻을 수 있었다고 합니다. 여기서 본 논문의 main contribution을 정리하면 다음과 같습니다.

1. 3D segmentation에서 최초로 하나의 panoptic 데이터셋에서 한번의 학습을 통해 멀티 태스크가 가능한 프레임워크를 설계
2. 새로운 쿼리 선택 알고리즘과 헝가리안 알고리즘을 대신할 수 있는 효율적인 매칭 기법 제안
3. indoor 데이터셋 ScanNet, ScanNet200, 그리고 S3DIS에서 3가지 segmentation task에 대한 SOTA 달성

2. Proposed Method

앞서 이야기한 것 처럼 본 논문의 구조는 SPFormer 구조를 따르며 변형하였는데, Figure 2의 파란색이 기존 구성이며 빨간색이 논문에서 제안하는 부분을 표시하고 있습니다. 전체적인 구조를 살펴보면, 먼저 sparse 3D U-Net 구조로 feature을 추출한 다음 이러한 feature는 flexible pooling을 통과하여 superpoint에서 포인트들의 feature을 단순히 평균내어 feature을 얻게 됩니다. superpoint feature는 semantic과 instance 쿼리를 입력으로 받는 트랜스포머 디코더의 key와 value로 사용되어 디코더는 cross attention을 통해 instance 쿼리에서는 instance identity, semantic 쿼리에서는 semantic 영역의 단일 물체에 대한 마스크를 나타내는 학습된 커널 집합을 출력합니다. 그리고 instance 커널을 end-to-end 방식으로 학습하기 위해서 disentangled 매칭 기법을 사용하였고 결과적으로 3가지의 segmentation을 모두 수행할 수 있게 됩니다.

2.1. Backbone and Pooling

먼저 N개의 포인트 P가 입력으로 들어온다고 가정하면 P \in \mathbb{R}^{N \times 6}라고 표현할 수 있습니다. 그럼 다음 raw한 포인트 클라우드를 복셀 형태로 변환하여 sparse 3D conv로 이루어져 있는 U-Net 백본 네트워크를 통과하여 포인트 feature P’ \in \mathbb{R}^{N \times C}를 얻을 수 있습니다. 여기까지는 SPFormer와 동일하며 그 다음 다른 부분은 superpoint pooling→ Flexible pooling을 적용한 것 입니다. 기존에 superpoint만을 고려한 pooling 레이어를 사용했다면 flexibility를 높이기 위해 복셀까지 고려한 pooling을 적용합니다. superpoint pooling에서는 superpoint feature S \in \mathbb{R}^{N \times C}는 사전에 미리 계산한 superpoint에 대해 백본 네트워크를 타고 나온 포인트 feature P’의 average pooling을 통해 얻을 수 있었습니다. 이러한 구조는 입력 포인트에 M개의 superpoint가 있다고 가정하고 설계하였으며, 본 논문에서도 동일한 가정을 바탕으로 복셀 기반 pooling을 추가하였습니다. 복셀 pooling에서는 복셀 그리드를 기준으로 백본을 통과한 feature을 pooling 하게 됩니다. superpoint를 구하는 클러스터링 과정에서 발생하는 계산량 대비 복셀화에 소모되는 계산량이 더 적기 때문에 복셀 형태가 더 선호될 수 있기 때문에 superpoint 기반과 복셀 기반의 pooling이 모두 가능한 flexible pooling을 제안한 것 입니다. 즉 SPFormer에서는 superpoint의 사용이 절대적이었으나, 본 논문에서는 이미 M개의 superpoint가 입력 포인트에 들어가있다는 가정을 통해 해당 구조에서 superpoint를 사용하지 않고 입력 포인트에 대한 feature을 복셀화하는 것이 문제가 되지 않기에 계산량이 상대적으로 적은 단순 복셀 feature 형태를 고려한 것 입니다. (일단 기본적인 전제는 “superpoint를 사용한다” 이기 때문에 pooling을 거친 feature는 아래에서 계속 superpoint feature라고 지칭합니다.)

2.2. Query Decoder

쿼리 디코더는 K_{ins}+K_{sem}을 쿼리로 입력받아 K_{ins}+K_{sem} 커널로 변홥합니다. 그 다음 superpoint feature는 해당 커널의 컨볼루션을 통과하여 K_{ins} instance와 K_{sem} semantic 마스크를 각각 생성합니다. 쿼리 디코더의 구조는 6개의 sequential한 트랜스포머 디코더 레이어를 통해 쿼리 간의 self attention과 cross attention을 수행합니다. 쿼리 중에 semantic 쿼리는 기존의 3D instance segmentation 방법론들과 동일하게 랜덤하게 초기화되며 instance 쿼리는 쿼리 선택 알고리즘을 통홰 초기화 됩니다.

기존에 2D instance segmentation에서는 쿼리 선택이라는 알고리즘을 통해 쿼리를 초기화하였는데, 쿼리 선택이란 입력 쿼리를 학습 중에 objectness loss로 계산되는 objectness score을 기반으로 샘플링되어 트랜스포머 인코더 feature를 초기화합니다. 반면 3D segmentation에는 이와 유사한 방식이 적용된 적이 없기 때문에 2D와의 갭을 줄이기 위해서 3차원 데이터(포인트 클라우드)와 트랜스포머 구조가 아닌 인코더에 적합한 쿼리 선택 알고리즘을 새롭게 구성합니다. 어떠한 objectness score가 기준이 아니라 flexible pooling 후 다운샘플링 된 superpoint feature로 쿼리를 초기화하게 됩니다.

2.3. Training

cost function

트랜스포머 기반의 end-to-end 학습을 위해서 쿼리와 GT 물체 사이의 cost function을 정의해야 하는데, 이러한 cost function을 최소화할 수 있는 매칭 기법과 매칭된 쿼리-GT 페어에 적용되는 loss에 대해서 알아보려고 합니다.

• p_{i, c_k} : i번째 proposal이 c_k semantic 카테고리에 속할 확률

C_{ik}는 i번째 proposal과 k번째 GT 사이의 유사도를 측정하기 위한 매칭 cost로, 식(1)에서 정의한 것처럼 classification 확률과 superpoint 마스크 매칭 cost C_{ik}^{mask}로 이루어져 있습니다.

• m_i, m^{gt}_k : superpoint의 예측된 마스크, GT 마스크

여기서 C_{ik}^{mask}는 BCE와 Dice loss로 이루어져 있는데, 저도 Dice loss는 처음 보는데 segmentation task에서 널리 사용되는 loss라고 합니다.

Disentangled matching

기존의 2D, 3D에서 트랜스포머 기반 방법들은 대다수가 헝가리안 알고리즘에 기반한 bipartite 매칭 기법을 사용하였습니다. 그러나 이러한 방식은 큰 단점이 존재하는데, 바로 proposal과 GT instance 사이에 의미있는 매칭 수가 많아지면 학습 과정이 길어지며 불안정하는 것 입니다. 그래서 본 논문에서는 리소스가 많이 드는 헝가리안 매칭을 사용하지 않기 위해 간단한 트릭을 씁니다. instance 쿼리는 논문에서 제안한 쿼리 선택 알고리즘을 통해 superpoint의 feature로 초기화하기 때문에 이 instance 쿼리는 superpoint와 정확하게 일치합니다. superpoint가 하나의 instance에만 속한다고 가정하면, superpoint와 GT 물체 사이의 대응 관계를 성립할 수 있습니다. 이러한 가정으로 확장시켜 생각해보면 GT 물체와 superpoint, instance 쿼리, 그리고 instance 쿼리로부터 얻는 instance proposal 간의 대응 관계를 정의할 수 있는 것이죠. 이러한 대응 관계는 결국 instance proposal과 GT instance 사이의 매칭을 가능하게 하기 때문에 이를 논문에서는 새롭게 disentangles 매칭이라고 정의합니다. 그러나 disentangles 매칭은 proposal의 수가 GT instance의 수보다 많기 때문에 bipartite 매칭이 가능하려면 GT 물체와 일치하지 않는 proposal을 필터링해야 합니다. 이를 위해 cost 행렬에서 대부분의 가중치를 무한대로 설정하여 cost function 함수의 최적화를 단순화하고 있습니다.

이게 무슨 말이냐하면, 일반적으로 cost 행렬에서 대부분의 값은 유한한 값을 가지고 있습니다. 그래서 계산 복잡도가 O(K^3_{ins})인 헝가리안 매칭을 통해 최적의 솔루션을 얻을 수 있는데, disentangled 매칭은 O(K_{ins})의 복잡도를 가지고 있습니다. GT instance k에 대해서 C^{ik}가 가장 작은 proposal i를 선택하기만 하면 되는 것이죠. proposal당 무한대가 아닌 값이 하나만 존재하기 때문에 간단하게 수행할 수 잇습니다.

Loss

매칭이 끝나고 나면 instance loss가 계산되는데, 각 proposal과 GT instance 매칭 쌍에 대해서 BCE L_{bce}와 Dice loss L_{dice}의 합으로 계산합니다. semantic 쿼리 K_{sem}은 고정된 순서로 주어지는 GT 마스크와 대응되기에 특정한 매칭 알고리즘을 필요로 하진 않습니다. L_{sem}이 semantic loss로 BCE를 사용합니다.

3. Experiments

3.1. Experimental Settings

우선 실험은 ScanNet, ScanNet200, 그리고 S3DIS 이렇게 3개의 indoor 데이터셋에서 진행하였습니다. segmentation task의 실험을 처음 보기 때문에 Metric도 살펴보면 .. semantic segmentation의 성능을 측정하기 위해서 mIOU를 사용하고 instance segmentation에서는 mAP를 기본적으로 사용합니다. 추가적으로 S3DIS에서는 해당 벤치마크에서 정의한 평가 프로토콜에 따라 mean precision (mPrec), mean recall (mRec)을 계산합니다. panoptic segmentation에서는 정확도를 평가하기 위해 PQ score을 사용하는데, thing과 stuff 카테고리에 대해 각각 PQ_{th}, PQ_{st}를 계산하여 리포팅하고 있습니다.

3.2. Comparison to Prior Work

우선 제가 가장 집중해서 본 실험은 역시 3D detection method와의 비교 실험입니다. 예측된 3D instance를 축이 고정된 3D 바운딩 박스가 둘러싸는 방식으로 3D detection에 적용하였다고 합니다. OneFormer3D는 기존의 포인트 클라우드를 입력으로 하는 detection 방법론들 사이에서 ScanNet 데이터셋 기준 SOTA를 달성하였습니다. 이러한 실험 결과를 통해 기존에 indoor detection에서 포인트들끼리의 샘플링과 클러스터링을 통해 일정 영역을 물체가 존재한다고 판단하는 방식보다 포인트 각각의 instance level에 대한 정보를 제공하는 것이 detection에도 효과적임을 확인할 수 있었습니다.

이제부터는 segmentation 방법론들과의 비교로 먼저 Table2는 ScanNet 데이터셋에서의 실험 결과 입니다. ScanNet에서는 3가지 segmentation task에서 모두 SOTA를 달성하였으며 특히 SPFormer와 가장 최근의 Mask3D와 비교하였을 때 큰 폭의 성능 향상을 이루었습니다. instance와 semantic segmentation 대비 panoptic segmentation은 타 방법론에서 많이 리포팅되지 않았지만 리포팅된 방법론들과 비교하였을 때 최소 +15 PQ 이상의 차이를 보이고 있습니다.

다음은 S3DIS 데이터셋에서 역시 모든 task에서 SOTA를 달성하였는데 저자는 mPrec나 mRec 평가 메트릭이 representation이 부족하지만 이전 방법론들과의 공정한 비교하고, S3DIS에서 정의한 평가 프로토콜의 일관성을 유지하기 위해 리포팅하고 있다고 이야기 합니다.

마지막으로 ScanNet200에서의 실험 결과로, 마찬가지로 모두 SOTA를 달성하였네요. ScanNet200에서는 이전까지 panoptic segmentation에 대한 결과를 리포팅한 방법론이 존재하지 않는데 처음으로 OneFormer3D에서 리포팅하면서 해당 분야의 future 연구를 위한 기반으로 다질 수 있었다고 합니다.

3.3. Ablation Studies

Query selection & disentangled matching

ablation study로는 ScanNet 데이터셋에서 새로운 쿼리 선택 알고리즘과 매칭 기법이 성능에 미치는 영향에 대해 실험을 진행하였습니다. instance segmentation 메트릭 mAP로 비교를 하였는데 그 이유는 쿼리 선택과 매칭에 있어서 semantic segmentation의 영향이 없기 때문이라고 합니다. SPFormer는 랜덤한 쿼리 초기화와 헝가리안 매칭 기법을 사용하고 있습니다. 논문에서 제안하는 쿼리 선택 방식은 기존의 bipartite 매칭 기법과 함께 사용하면 성능 향상이 거의 발생하지 않지만 Discentangled 매칭과 함께 사용할 때 최소 1.3%의 개선이 되는 결과를 통해 제안하는 두 방식의 시너지 효과를 통해서 SOTA를 달성할 수 있었음을 확인할 수 있습니다.