Masked-attention Mask Transformer for Universal Image Segmentation

이번에 소개 드릴 연구는 Facebook Research에서 진행한 MaskFormer의 후속 연구인 Mask2Former로 논문의 공식 명은 Masked-attention Mask Transformer for Universal Image Segmentation 으로 CVPR 2022에 소개 되었습니다. 기존 Image segmentation의 테스크인 Semantic/Instance/panoptic Segmentation을 하나의 아키텍쳐에서 학습 할 수 있도록 Universal Architecture를 제안합니다.

TL;DR

• 기존 segmentation 연구들은 각 테스크(Semantic/Instance/Panoptic)를 독립적으로 연구해왔음
• 하지만 각 테스크들은 Semantics만 다를 뿐 나머지는 동일함
• 본 연구는 모든 테스크를 하나의 architecture에서 train/infer 할 수 있는 Architecture와 방법론을 제안함
• 결론적으로 다양한 Benchmark Dataset에서 State-of-the-art를 달성함

Image Segmentation은 대표적으로 Semantic / Instance / Panoptic Segmentation Task로 이루어져 있습니다. 본 연구에서는 3가지 Task 들 모두 Semantics만 다르지 그 이외에는 모두 동일하다고 이야기합니다. 그럼에도 불구하고 현재까지의 Fragmentation 연구들은 각 테스크에 맞는 specialized architecture로 연구되고 개발되어왔고, 이로 인하여 서로 다른 Task에 대한 유연성이 결여되어왔다고 주장합니다. Mask2Former는 이러한 문제를 극복하고 모든 Segmentation task에 대한 capacity를 갖는 universal architecture를 제안합니다.

기존의 MaskFormer 역시 universal architecture로 동작할 수 있습니다. 하지만 위 그림과 같이 specialized architecture와 비교했을 때 현저하게 낮은 성능을 보입니다.(특히 Instance Seg.) Mask2Former는 MaskFormer의 연구를 개선하여 Mask2Former을 제안합니다. MaskFormer와 Mask2Former의 차이는 다음과 같습니다.

• Masked-attention: foreground object feature에 집중하기 위하여 기존 cross-attention을 자체적인 masked-attention으로 교체함.
• Multi-scale high-resolution features: small object에 강인할 수 있도록 multi-scale feature를 사용함. 모두 high-resolution을 사용하는 것이 아닌 low/high-resolution 을 적절히 사용하여 메모리 효율성을 높임.
• Optimization improvements: (1) self-attention과 cross-attention(masked-attention) 의 자리를 바꿈 (2) query를 학습가능하도록 변경함. (3) sampling된 point만 선별적으로 loss를 계산하여 메모리 효율성을 높임

Masked-attention Mask Transformer

Mask classification preliminaries

Mask2Former는 Universal Architecture인 MaskFormer를 베이스라인으로 활용한 연구입니다. MaskFormer는 다음과 같은 구조로 이루어져 있습니다.

• Backbone:  이미지로부터 low-resolution features를 추출함
• Pixel Decoder: low-resolution features로 부터 점진적으로 high-resolution features로 upsample하여 high-resolution per-pixel embedding을 생성함
• Transformer Decoder: image feature를 처리하기 위한 객체 쿼리를 사용하는 Transformer Decoder. 이를 통해 object queries와 mask를 생성함

Transformer decoder with masked attention

Mask2Former와 MaskFormer의 주요 차이점은 Transformer Decoder이며, 차이점은 다음과 같습니다:

1. Masked-attention: Mask2Former에서는 모든 이미지 특성에 대한 Cross-Attention 대신, 예측된 foreground mask region에서 Cross-Attention을 수행하기 위해 Masked-Attention을 사용합니다. 이로써 전체 이미지에 대한 Cross-Attention을 수행하는 것이 아니라 Foreground 예측에 더 집중할 수 있게 됩니다.
2. Small object handling(Multi-scale high-resolution features): 작은 객체를 다루기 위해 Mask2Former는 Pixel Decoder의 Pyramid에서 나오는 고해상도 특성을 연속적으로 feed하여 처리합니다.

Masked attention

Mask2Former의 중요한 특징 중 하나는 cross-attention 대신 masked-attention을 사용하여 수렴 속도를 향상시키는 것입니다. 기존의 Transformer-based 모델은 전체 이미지에 대한 cross-attention을 수행하기 때문에 모든 위치의 정보를 고려하기에 수렴 속도가 느립니다.

Mask2Former는 이러한 한계를 극복하기 위해 가설을 제시합니다. 이 가설에 따르면, self-attention을 사용하여 로컬 특성만으로도 객체 쿼리(object query)를 업데이트하고 필요한 컨텍스트 정보를 수집하는 것이 가능하다는 것입니다. 이를 실현하기 위해 Mask2Former에서는 예측의 Foreground 영역만을 고려하는 masked-attention 메커니즘을 제안합니다. 이렇게 하면 전체 이미지에 대한 정보를 고려하지 않고도 필수적인 정보만을 처리할 수 있으며, 이로 인해 모델의 수렴 속도가 크게 향상됩니다.

기존의 cross-attention과 다른 부분은 cross-attention에서 softmax 함수 안에 𝓜_(l-1)을 추가하여 마스크가 있는 구역에서만 Cross-Attention을 수행하도록 강제하는 부분입니다. 여기서 𝓜_(l-1)은 l-1 번째 (down sampled) mask의 임계값(threshold)이 0.5 이상인 경우에는 0을, 그렇지 않은 경우에는 음의 무한대를 입력(softmax->0)으로 사용합니다.

이렇게 하면 Cross-Attention에서 이전 단계(l-1)에서 얻은 마스크 정보를 사용하여 현재 단계(l)에서 마스크가 있는 위치에 집중할 수 있습니다. 이로써 모델은 관련성이 높은 정보에 더 집중하고 계산 리소스를 절약할 수 있으며, 수렴 속도와 정확도를 향상시킬 수 있습니다.

High-resolution features

Segmentation task에서 high-resolution는 성능향상에 도움이 되지만, 이는 너무나 많은 computational resource를 사용합니다. Mask2Former는 computational resource를 제어하기 위해 모든 layer에 High-resolution를 사용하는 대신 pixel decoder의 피라미드 레이어에서 추출되는 lowest-resolution부터 highest-resolution을 Transfermer Decoder에 순차적(1/32, 1/16, 1/8)으로 feed하여 효율적인 연산을 구성하였습니다.

Optimization improvements

기존의 Transformer Decoder는 self-attention, cross-attention, 그리고 FFN(Feed-Forward Network)으로 구성되어 있었고, 이에 residual connection과 attention에 Dropout을 사용해왔습니다. 그러나 Mask2Former에서는 self-attention이 이미지와 무관하며 어떠한 이미지 신호도 사용하지 않기 때문에, self-attention과 cross-attention(mask-attention)을 switching했습니다. 또한, 본 연구에서는 Transformer Decoder 내에서 Dropout이 효과가 없을 뿐더러 오히려 성능 저하를 가져온다는 사실을 발견하여 Dropout을 제거하였습니다. 마지막으로, 초기에 0으로 채워진 X_0의 쿼리도 학습 가능하게(SIREN and learnable scale-level embedding) 하였으며, 이는region proposal과 유사하게 동작했다고 합니다.

Improving training efficiency

기존 MaskFormer의 주요 문제 중 하나는 고해상도 마스크 예측으로 인한 높은 메모리 사용량입니다. 이러한 고해상도 마스크 예측은 단일 이미지에서 32G의 메모리를 사용했다고 합니다. 그러나 Mask2Former는 PointRend와 Implicit PointRend에서 영감을 받아서 일부 포인트만을 손실(loss) 계산에 활용합니다. matching cost를 계산하는 부분에서는 일괄적으로 K개의 포인트를 선택하여 계산했고, 최종 손실(loss) 계산에는 PointRend의 Important Sampling를 사용합니다.

Experiments

Ablation studies