[논문 리뷰] K-Net: Towards Unified Image Segmentation

Abstract

 

본 논문은 Semantic, instance, panoptic segmentations의 유사한 task들에 대해 효과적인 framework을 제시한다. 이 Framework의 이름은 K-Net으로 instance와 semantic 카테고리에 대해 segments를 수행한다. 이는 learnable kernels 그룹을 통해 이루어지며 각 kernel은 잠재적 instance와 stuff class에 대한 mask를 유동적으로 생성해낸다. 다양한 instances를 구분하는 데 어려움을 해결하기 위해 우리는 input image에 있는 의미있는 group에 대해 각 kernal이 dynamic하고 conditional하게 update할 수 있는 전략을 제시한다. K-net은 bipartite matching과 함께 end-to-end로 학습될 수 있고, training과 inference는 모두 NMS-free, box-free이다. K-net은 이전의 모든 SOTA single-model의 성능을 제쳤고, inference speed 또한 더 빠르다.

 

Introduction

 

Image Segmentation은 coherent pixels를 찾는 것을 목표로 한다. 크게 semantic categories 또는 instance categories를 찾는 task로 나누어져있고, 이 두가지를 포괄하는 시도(panoptic) 또한 이루어지고 있다. Pixels를 semantic categories로 grouping하는 것은 dense classification problem으로 정의될 수 있다.

Fig-1에서 볼 수있듯, 최근의 모델에서는 convolutional kernels(본 논문에서는 semantic kernels라고 명명)의 집합을 통해 미리 정의된 pixels이나 regions의 class를 학습한다. 그러나, 이러한 개념을 instance로 확장하면 이미지 내 instance수가 다양하기 때문에 쉽지 않다. 결과적으로 instance segmentation은 object detection 또는 embedding generations등과 같은 추가적인 단계를 거쳐야하는 한계가 있다. 이러한 방법은 NMS나 pixel grouping과 같은 복잡한 post-processing이나 extra componets에 의존할 수 밖에 없다. 최근의 접근에서는 dense feature grids에서 kernels를 생성하고 segmenation을 위한 kernels를 선택하는 방법으로 frameworks를 단순화하고 있다. 그럼에도 불구하고, kernels들을 열거하고, 선택하기 위해 밀집된 그리드를 기반으로 구축되기에, 여전히 중복된 마스크나 kernel을 제거하기 위해 post-processing이 의존한다.

 

본 논문에서는 kernels라는 개념을 통해  semantic, instance, panoptic tasks를 모두 합쳐서 수행할 수 있는 효과적인 framework을 만들었다고 말한다. 이는 K(ernels)-Net 이라 명명한다. K-Net은 랜덤하게 초기화된 convolutional kernels로 시작하여 semantic categories를 위해서는 semantic kernels를 통해 학습하고, instance identities를 위해서는 instance kernels를 통해 학습한다. 이는 자연스럽게 panoptic task까지 가능하게 한다.

K-Net의 다양성과 단순성은 두가지 디자인을 통해 구현된다. 먼저 K-Net의 kernels들을 동적으로 업데이트하여 image의 activations에 대해 조건을 지정할 수 있도록 한다. 이러한 content-aware 매커니즘은 각 kernel, 특히 instance kernel이 다양한 객체에 대해 정확하게 인식하는지 확인하는데 중요하다. K-Net은 이러한 적응형 kernel update 전략을 반복적으로 적용함으로써 kernel의 식별 능력을 향상시킨다. 이 전략은 모든 segmentation tasks에 universally하게 적용된다는 점이 주목할만한 성과이다.

두 번째로는, 각 kernel이 학습할 targets을 bipartite matching을 통해 지정해주는 것이다. 이는 kernels와 image 내 instance를 1 대 1로 매핑해준다는 점에서 기존의 training 전략들에도 유리하다. 따라서 이미지의 다양한 instance 수를 처리하는 문제를 해결할 수 있다. 또한 box를 포함하지 않고, mask-driven 구동 방식으로 이루어지기 때문에 NMS-free와 box-free를 달성할 수 있어, real-time applications에 효과적이다.

K-Net은 SOTA single-model의 결과를 모두 앞지르는 결과를 보이며, expensive 한 모델과도 비슷한 성능을 보인다.

 

Related Work

 

Semantic Segmentation - SS는 주로 FCN을 기반으로 하여 dense classification으로 다뤄진다. 이 framework을 기반으로 하여, dilated convolution이나, pyramid pooling, context representations와 attention mechanisms등을 활용해 feature representation을 높이기 위한 시도가 이루어지고 있다. 논문이 나온 시점에서는 SETR이 ViT를 활용해 sequence-to-sequence 형태로 재구성하여 접근했다. 비록 다양한 시도가 이루어졌으나, 이러한 접근법은 static semantic kernels(고정된 input 크기에 대해서만 동작하도록)을 활용하여 예측을 수행한다. 이에 반해, K-Net은 kernels를 dynamic하고 conditional하게 NN의 activations에 적용한다.

 

Instance Segmentation - IS는 top down, bottom up 접근 방식으로 나뉘는데, top down의 경우에는 먼저 bbox를 찾아낸 후, 각 bbox에 맞는 mask를 잘 생성하도록 하는 방법이다. Bottom-up 방식은 미리 SS를 진행한 후에 pixels를 각기 다른 instances에 grouping 하는 방법이다. 이에 반해, K-Net은 각 kernel이 하나의 객체에 대해 한번에 하나의 마스크를 예측하도록 제한한다. 이를 통해 bbox 또는 grouping process가 추가로 필요하지않다.
또한, 최근에는 task를 dense mask prediction으로 적용하여 instance segmentation을 one stage로 하는 방법도 존재한다. 하지만 이러한 methods들 또한 NMS나 중복된 instance masks를 제거해야해 end-to-end training이 힘들다.

 

Panoptic Segmentation - PS는 scene을 더 잘 이해하기 위해 SS와 IS를 combine 해서 사용하는 task이다. 주로 instance segmentation framework에 SS branch를 더하거나 또는 pixel grouping 전략을 다양하게 한다. 최근에는 DETR이 framework을 Transformer를 통해 좀 더 단순화하려는 시도가 있지만, box 예측, stuff classes들 예측 등을 추가로 수행해야해서 object detection과 embedding generations 등이 필요한 상황이다. 하지만 K-Net은 semantic kernels와 instance kernels를 활용해서 image를 semantic region로 나누고, object instance로 나누기 때문에 kernels의 관점에서 다시 볼 수 있다.

 

Dynamic Kernels - Convolution kernels은 주로 static하다. 따라서 제한된 representation 능력을 가진다. SS에서도 dynamic kernels를 도입하여 receptive fields를 넓히거나, multi-scale을 다루게 하여 model의 representation을 넓힌 사례가 있다. 이와 다르게 K-Net에서는 dynamic kernels를 활용해 kernels의 input feature보다 discriminative 기능을 더 향상시키는데 목적이 있다. IS나 PS에서도 이를 도입한 적있으나, NMS나 extra kernel fusion 등의 추가적인 과정이 필요하다.

 

Methodology

 

3.1 K-Net

 이 섹션에서는 제안된 K-Net이 각 pixel이 잠재적 instance나 semantic class에 assign 되도록하는 kernels 설명한다.

 모든 segmentation task는 기본적으로 각 pixels들을 미리 정의된 meaningful group에 할당한다. 이미지의 그룹 수를 N이라고 설정할 수 있다. 예를 들어, N개의 semantic classes가 있을 수 있고, 또는 N개의 Object가 image에 존재할 수 있다. PS에서는 stuff classes와 image 내 Object 개수의 합이 N이 될 수 있다. 따라서 N개의 커널을 사용해 이미지를 N개의 그룹으로 분할 할 수 있으며, 각 커널은 해당 그룹에 속하는 pixel을 찾는 역할을 담당한다. 구체적으로 B개의 images에 대해 (B, C, H, W)의 input feature map이 주어졌다고 하면, convolutions을 하기 위해 (N, C) 형태의 N개의 kernels만 필요하다. 예측 M을 얻기 위해서는 이 (N, C) 커널과, F, 즉 (B, C, H, W)가 필요하다. 이는 아래 수식으로 표현될 수 있다.

σ함수는 SS에서는 softmax 함수가 될 것이고, multiple masks를 사용하거나 binary masks가 필요하다면 sigmoid도 사용 가능하다. 이 수식은 semantic segmentation을 오랫동안 지배했다. SS에서는 각 kernel은 모든 pixels에 유사한 class를 매핑하는데 사용되어 왔다. 하지만 IS에서는 각 pixel 그룹이 객체에 해당한다. 이전 방법에서는 kernels 대신 다른 extra steps를 통해 instances를 분리했으나, 이 논문에서는 처음으로, 이러한 kernels의 개념을 확장하여 SS의 개념이 IS에도 적용 가능해 PS까지 달성할 수 있다고 본다. Kernels를 통해 instances를 구분하기 위해서 K-Net의 각 kernel들은 한 이미지의 하나의 object만 segment 해낸다. 이를 통해 K-Net은 extra steps 없이 instances를 분리해내고, segmentation 또한 동시에 해낸다. 단순히 하기 위해 논문에서는 semantic kernels와 instance kernel라고 부른다. 이 둘의 조합은 PS 또한 달성해낸다고 말한다.

 

3.2 Group-Aware Kernels

 Kernels의 분할 능력을 높이기 위해 pixel groups 내의 contents를 가지고 static kernels를 update하는 방법을 설명한다.

Instance를 kernels 별로 직접 분리하는 것은 쉽지 않다. Instance kernels은 이미지 내에서 다양한 scale과 appearance를 가진 objects를 구분해야하기 때문이다. Semantic categories와 같이 공통적이고 특이한 특징을 가지고 있지 않으면 instance kernels는 static kernels보다 더 강력한 discriminative 능력이 필요하다.

 이러한 문제를 극복하기 위해 해당 pixel 그룹에 대해 kernels를 조건부로 만드는 접근 방식을 제시한다. Kernel Update Head 인 fi에는 3가지 단계를 내포하고 있다. Group feature assembling, adaptive kernel update, 그리고 kernel interaction이다. 첫번째로, Mask 예측 Mi-1을 사용하여 각 pixel 그룹에 대한 group feature Fk를 조합한다. Fk는 각 개별 그룹이 다른 그룹과 구별되는 feature를 가지고 있기 때문에, kernel Ki-1이 update 되는데 사용할 수 있다. 그 후에 Kernel은 interaction하여 image context를 포괄적으로 모델링한다. 마지막으로, 획득된 group-aware kernels인 Ki는 feature map F와 convolution 연산을 수행해 좀 더 명확한 Mi 예측을 이루어낸다. 아래 그림을 보면 이러한 프로세스는 반복적으로 수행될 수 있다. 왜냐하면 더 미세한 partition은 group feature의 노이즈를 줄여 더 discriminative kernels를 만들어 낼 수 있기 때문이다.

특히, kernel update Head는 kernel의 특성에 의존하지 않는 다는 점에서 universal하다. 따라서 Instance kernels 뿐 아니라 semantic kernels 또한 향상시킬 수 있다.

 

Group Feature Assembling -  Kernel update Head는 먼저 각 group의 features를 조합하고 추후에 kernels가 group-aware가 되도록 돕니다. 각 Mi-1의 kernels에 있는 mask들은 본질적으로 pixel이 kernel에 속하는지 아닌지 정하는 값이므로, feature F와 Mi-1를 곱하여 Ki-1에 대한 feature Fk를 조합할 수 있다. 

Adaptive Feature Update - Kernel update Head는 Fk를 가지고 kernels의 representation 능력을 높이려고 한다. Mi-1이 부정확할 수 있기 때문에, 각 group의 feature에는 다른 group으로부터 온 노이즈가 포함되어 있을 수 있다. 이러한 부작용을 줄이기 위해 adaptive feature Update를 제안한다. 구체적으로는, 먼저 Fk와 Ki-1를 element-wise로 곱한다. φ1 는 선형 변환이다.

Kernel update head는 GF와 GK를 학습하고, 이는 Fk와 Ki-1에 GF와 GK를 곱하고 더해 Kernel K~를 만든다. K~는 이제 kernel interaction에 사용되게 된다. 아래 수식에서 ψn (n:1~4)는 LayerNorm(LN)과 sigmoid function이 포함된 fully connected layer이다. 

GF와 GK는 여기서 Transformer에서 나온 self-attention과 비슷한 역할을 한다. Transformer의 self-attention은 일반적으로 Query와 Key의 내적을 통해 계산된다. 마찬가지로 Adaptive Feature Update 또한, Kernel 특성인 Ki-1과 그룹 특성 FG의 가중치 합산을 수행한다. 가중치 GK와 GF는 element-wise 곱셈으로 계산된다.

Kernel Interaction - Kernel 간 interaction은 각 kernel이 다른 group의 context 정보를 알리는데 매우 중요하다. 이러한 정보는 kernel이 image의 그룹 간의 관계를 암시하고, 모델링해 활용할 수 있다. 이를 위해 update된 kernel K~가 주어지면 새로운 커널 Ki를 얻기 위해 kernel interaction을 추가했다. 여기서는 Multi-Head Attention과 Feed-Forward NN을 선택했다. 

3.3 Training Instance Kernels

 bipartite matching 전략을 통해서 end-to-end 로 만들어 instance kernels를 학습할 수 있도록 하였다.

 

3.4 Applications to Various Segmentation Tasks

 K-Net은 semantic, instance, 그리고 panoptic segmentation 모두에 적용될 수 있다.