[NeurIPS 2023] Active Learning for Semantic Segmentation with Multi-class Label Query

안녕하세요. 허재연입니다. 오늘은 Segmentation분야의 Active Learning 논문을 가지고 았습니다. 포항공대에서 작성한 논문으로, 이번 KCCV2024 2일차에 포스터로 발표될 예정인 논문입니다. 지금까지 제가 읽은 AL분야 논문들에 모두 classification task를 segmentation에는 Active Learining이 어떻게 적용되는지 궁금해 읽어보게 되었습니다. 리뷰 시작하겠습니다.

Introduction

본 논문에서는 segmentation task에 대해 학습 알고리즘이 학습 데이터 중 높은 정보량을 가지고 있는 고가치 데이터를 선별하여 oracle에게 라벨링을 요청하는 AL for semantic segmentation을 다룹니다. 근 10년간 semantic segmentation분야는 지도학습을 이용한 딥러닝의 활용으로 크게 발전했지만, 픽셀 단위 라벨링의 비용이 높기 때문에 동시에 이를 실제 사례에 적용시키는데는 데는 제한이 있었습니다. 이 한계를 극복하기 위해 weakly supervised learning, semi-supervised learning, self-supervised learning, active learning 등 라벨링의 효율을 높일 수 있는 방법이 많이 연구되었습니다.

그 중 Active Learning에서는 주어진 예산 안에서 oracle이 라벨링한 데이터로 수행한 supervision의 양과 질을 높이기 위해 annotation query(쿼리 샘플 및 어노테이션 형식)의 설계가 중요했습니다. 기존에는 전체 이미지 하나를 샘플로 취급하고 이 이미지의 픽셀 단위 클래스 라벨링을 요청하는 방법과, 각각 픽셀을 선별하고 이를 오라클에게 클래스 라벨링을 요청하는 방법이 연구되었지만 전체 이미지를 샘플로 간주하면 샘플의 다양성이 부족해지고 픽셀 단위로 쿼리를 하면 예산 효율성이 너무 떨어지는 문제가 있었다고 합니다.

이 두 방법을 절충하기 위해 최근에는 서로 겹치지 않는 이미지의 부분 영역을 개개의 샘플로 취급하는 방법들이 연구되었는데, 이러한 영역 기반(region-based) 방법들은 수많은 이미지에서 다양한 맥락의 local region이 선택되어 샘플의 다양성을 보장할 수 있으며 동시에 영역 단위 segmentation label을 얻을 수 있도록 설계되어 효율도 좋다고 합니다. 좀 더 구체적으로 예시를 들면 oracle에게 이미지 패치 내에서 polygon 형태의 segmentation mask를 그리도록 요청하거나 superpixel의 dominant class label을 추정하여 해당 superpixel 내부 전체에 라벨을 할당하도록 합니다. 이러한 영역 기반 방법론들이 좋은 결과를 보여주긴 하였지만, 저자들은 이러한 쿼리 디자인에 여전히 개선의 여지가 많이 남아 있다고 합니다. Polygon labeling방법들은 여전히 많은 수의 쿼리 당 클릭수를 요구하고, dominant class labeling 방법들은 multi-class 영역 부분에 잘못된 supervision을 줄 수 있다고 합니다(객체의 경계를 위반하는 경우가 많고 다양한 클래스를 포함하는 경우가 많지만 러프하게 해당 영역에 대한 dominant class로 라벨링해버리니까요).

이러한 맥락에서, 저자들은 semantic segmentation에 있어 영역 기반 Active Learning을 위해 새로운 query design을 소개합니다. 이 방법은 oracle에게 주어진 영역에 있는 모든 class에 대한 multi-hot vector를 라벨링하도록 요청하는데, 이는 부분 영역의 annotation error를 막고 dominant class labeling의 이점을 유지할 수 있다고 합니다. 이 뿐만 아니라 multi-class labeling은 더 적은 클릭 당 어노테이션 시간(annotation time per click)을 필요로 하며 Figure1의 dominant class labeling모다 더 정확한 라벨링이 가능하다고 합니다.

하지만 이런 multi-class labels는 선택된 영역 의 단일 픽셀에 여러 클래스 후보인 partial labels가 할당되기에 필연적으로 class ambiguity issue를 발생시키는데, 이러한 클래스 모호성 문제를 해결하기 위해 partial labels를 이용한 방법을 제안합니다. partial label learning은 각 학습 데이터에 후보 클래스들을 할당하는 weakly supervised learning 방법 중 하나로, ambiguous supervision을 유발한다고 합니다. 저자들은 2단계 학습 알고리즘을 도입해 수도라벨링을 사용하여 이 클래스 모호성 문제를 완화하였다고 합니다.

제안하는 방법론의 전반적인 파이프라인은 Figure2에 나와 있습니다. 해당 방법론은 2단계로 구성되어 주어진 부분 영역들과 이들의 멀티클래스 라벨들에 대해 segmentation network를 학습합니다.

첫 번째 단계에서는 직접적으로 영역 단위 멀티클래스 라벨로 네트워크를 훈련시키며, 이를 위해 partial-label learning과 multiple instance learning을 기반으로 label disambiguation을 없애기 위해 새로운 손실함수 2개를 도입합니다. 두 번째 단계에서는 지도학습 방법으로 segmentation network를 훈련하는 데 사용되는 pseudo segmentation label을 통해 partial labels를 disambiguate합니다. 좀 더 구체적으로 설명해면 첫 번째 단계의 모델을 사용하여 각 로컬 영역에서 class prototype feature들을 찾고 prototypes를 region-adaptive classifier를 사용하여 region의 픽셀 단위 수도라벨을 예측하도록 합니다. 또한 여기에 추가적으로 쿼리 당 supervision의 양을 늘리기 위해 수도라벨을 이웃 로컬 영역에 전파하도록 하여 멀티클래스의 수도라벨을 확장하도록 하였습니다.

마지막으로 영역 기반 AL에서 multi-class labels의 이점을 크게 할 수 있도록 습득 함수를 디자인 했다고 합니다. 해당 습득함수는 모델이 어렵고 underrepresented된 클래스를 포함하는 local region을 더 자주 선택할 수 있도록 uncertainty 및 선택된 영역의 class balance를 모두 고려하도록 하였습니다. 이렇게 고안된 프레임워크는 Cityscapes 및 PASCAL COV2012에서 각각 SOTA를 달성하였고 특히 전체 라벨링 비용의 4%만을 사용하여 cityscape에서 완전지도학습 학습 성능의 95%를 달성했다는 점을 달성했음을 강조합니다.

저자들은 본 논문의 contribution을 다음과 같이 요약했습니다 :

1. 우리는 semantic segmentation의 영역 기반 Active Learning에서 oracle에게 특정 영역에 포함된 모든 class의 multi-hot vector를 라벨링 요청하는 방법의 새로운 쿼리 방법을 설계하였다.
2. 우리는 멀티클래스 라벨의 supervision과 멀티클래스 라벨로부터 pseudo segmentation label을 생성하는 방법을 효과적으로 이용하는 새로운 손실함수 2개를 포함하는 새로운 AL 프레임워크를 제안했다.
3. 멀티클래스 라벨의 이점을 최대화하기 위해 로컬 영역의 multiple classes를 고려한 습득 함수를 설계했다.
4. 추가적인 실험 및 실세계 어노테이션 시나리오의 user study를 통해 멀티클래스 라벨링의 효과를 입증하였다.
5. 제안된 프레임워크는 Cityspaces 및 PASCAL VOC 2012에서 annotation cost를 상당히 낮추며 SOTA를 달성하였다.

Method

AL 과정은 총 R 라운드로 진행되며, 각 라운드마다 습득함수를 하용하여 배치 단위의 local region이 선택되고 oracle이 여기에 multi-hot vector(multi-class label)을 할당하게 됩니다. 라벨이 지정된 영역이 주어지면 학습 알고리즘은 Figure2와 같이 2단계로 동작합니다. 첫 번째 단계에서는 라벨의 모호성(ambiguity)를 처리하는데 특화된 두 손실 함수를 사용해 영역별 다중 클래스 레이블로 segmentation model을 직접 학습하고 두 번째 단계에서는 픽셀 단위의 수도라벨을 생성해 모호성을 완화하고 이들을 모델 학습에 추가적으로 사용하게 됩니다.

Acquisition of region-wise multi-class labels

우선, unlabeled image set에서 각 이미지를 서로 겹치지 않도록 분할합니다. 이 때 겹치지 않는 각 영역은 기존에 제안되었던 superpixel algorithm을 따른다고 합니다. superpixel이 무엇인지 와닿지 않으신다면, 아래 그림에서 좌하단 부분을 참고하시면 좋을 것 같습니다. 다음과 같이 영역을 분할한다고 생각하시면 됩니다.

수식 1에서 f_{t}(x)는 feature extractor이고, [w_{t,1}, w_{t,2}, … ]는 classifier의 weight matrix이며 τ는 temperature parameter입니다. 이 predictive probabilty를 이용하여 uncertainty를 측정하는데, x에 대해 가장 큰 predictive probability c_{b}와 두번째로 큰 predictive probability c_{sb}를 사용해 다음과 같이 uncertainty를 측정합니다.

이 때 샘플링 클래스의 밸런스를 맞추기 위해 라벨 분포 P_{θt}(y)를 다음과 같이 추정합니다:

여기서 X는 이미지 영역에 있는 픽셀들의 집합입니다. 습득함수는 희소한 클래스의 불확실한 영역을 얻기 위해 다음과 같이 정의됩니다 :

v는 클래스 밸런싱 효과에 규제를 주기 위한 하이퍼파라미터입니다. dominant class만을 고려했던 이전 습득함수와는 달리, 이 함수는 영역 안에 있는 모든 픽셀의 클래스들을 고려하여 multi-class labeling과 잘 align될 수 있다고 합니다.

Figure 2를 참고하며 아래 방법론을 읽으면 이해가 쉬울 것입니다

Stage 1: Learning with region-wise multi-class labels

전통적인 supervised learning에서는 segmentation model을 학습할 때 하나의 클래스로 지정된 label을 가지고 픽셀 단위 CE loss를 사용했었습니다. 로컬 영역 의 단일 클래스 라벨 집합은 다음과 같이 정의됐었고

픽셀 단위 CE loss는 다음과 같이 주어졌었습니다.

반면 Dm := D – Ds로 표현된 멀티클래스로 라벨링된 영역은 멀티클래스 라벨이 해당 픽셀에 대한 정확한 클래스를 가지고 있지 않기 때문에 픽셀 단위 CE loss 학습에 사용하는데에 제약이 있었습니다(weak label). Dm을 효과적으로 활용하기 위해 저자들은 2가지 손실함수를 소개합니다

Merged positive loss.

영역의 각 픽셀은 partial labels(후보 클래스 집합)가 할당되어 있으며 영역 당 각 픽셀 예측은 이들 후보 클래스 중 하나가 되어야 합니다. 이 개념으로 고안된 것이 merged positive loss로, 다음과 같이 정의됩니다 :

이 손실함수는 모든 후보 클래스의 predictive probability가 positive로 간주되기 때문에 후보 집합의 모든 클래스를 예측하도록 합니다. 멀티클래스 라벨의 예측 분포를 합산해서 CE를 진행했다고 생각하시면 됩니다.

Prototypical pixel loss

multi-class labels 학습에 있어서 각 영역 안의 픽셀 중 적어도 하나는 각 후보 클래스에 대해 positive여야 합니다. 저자들은 이러한 픽셀을 prototypical pixel이라 하며 해당 영역의 각 후보 클래스에 대해 가장 confidence한 예측을 하는 해당 픽셀을 prototypicla pixel로 설정합니다.

여기서 c는 multi-class 집합 중 하나입니다. segmentation model은 각 prototypical pixel에 할당된 class가 참이라는 가정을 가지고 CE loss를 적용하여 학습됩니다. loss 구성은 다음과 같습니다 :

프로트타입 펙셀이 항상 gt값과 일치하는 것은 아니지만 다양한 지역의 많은 프로토타입 픽셀을 사용한 훈련을 통해 모델이 각 클래스에 대한 기본적인 이해를 갖출 수 있을 것으로 기대합니다. 또한 이 손실함수는 모든 후보 클래스가 학습에 동일하게 기여하기에 class imbalance를 완화할 수 있다고 합니다.

요약하면, 첫번째 단계에서 학습 loss는 다음과 같이 사용됩니다. λ_{CE}와 λ_{MP}는 가중치 조절을 위한 하이퍼파라미터입니다.

stage 2: Learning with pixel-wise pseudo labels

2번째 단계에서는 픽셀 단위의 원핫 라벨을 생성하고 활용하여 partial label의 모호성을 완화합니다. 수도라벨 생성 과성은 각 팍셀에 수도 클래스 라벨을 할당하는 intra-region label localization과 수도라벨을 인접한 unlabeled region으로 확장하는 label expansion의 2step으로 나뉩니다.

Intra-region label localization

레이블이 지정된 각 영역에 대해 주석이 달린 각 클래스의 프로토타입을 클래스의 프로토타입 픽셀에 위치한 특징 벡터로 정의합니다(stage 1의 수식 8에 의해 추정). 이러한 프로토타입 집합은 지역 내 픽셀 수준 분류에 사용되는 region-daptive classifier로 사용됩니다. 각 픽셀에 그것의 feature space에서 최근접 프로토타입 클래스를 할당하는 것입니다. 픽셀 x에 할당된 수도라벨은 다음과 같이 추정됩니다.

x_{s,c}^{*}는 클래스 c의 prototypical pixel이며 cos(f,f’)는 두 feature vector f와 f’ 간 코사인 유사도입니다.

Label expansion

label localization과 유사하게 Label expansion 단계는 레이블이 지정되지 않은 영역 집합 N(s)의 클래스 레이블을 할당하는 것을 목표로 합니다. 하지만 L(s)에 관련 없는 클래스가 포함될 수 있어 label을 N(s)의 모든 픽셀로 라벨을 전파하면 잘못된 수도라벨이 발생하여 모델 성능이 저하될 수 있습니다. 따라서 수도라벨은 feature space에서 적어도 하나의 prototype에 충분히 가까운 관련 픽셀에만 제안됩니다. 좀 더 구체적으로는 region-adaptive하고 class-adaptive한 방법으로 연관성을 계산하기 위해 prototype-adaptive threshold를 사용할것을 제안합니다. class c ⊂ Y에 대한 prototype-adaptive threshold는 다음과 같이 정의됩니다.

med(·)는 집합의 중앙값을 반영하고, x_{s,c}^{*}는 클래스 c의 prototypical pixel이고, y\hat(x)는 x에 대한 수도라벨입니다. labeled region의 수도라벨을 인접 영역의 픽셀에 다음과 같이 전파합니다.

여기서 Y\hat(x) :={c : cos(f_{θ}(x), f_{θ}(x_{x,c}^{*}))}이고 threshold 이상입니다. prototype-adaptive thresdhold를 이용한 필터링으로 각 영역의 라벨 확장의 양을 손으로 직접 튜닝할 필요 없이 조절할 수 있다고 합니다.

그런 다음 label localization 및 expansion 단계 모두에서 생성된 pseudo segmentation label을 사용하여 픽셀 단위 CE loss를 사용하여 segmentation 모델을 추가로 학습하게 됩니다.

Experiment

실험은 대표적인 semantic segmentation dataset인 Cityscape와 PASCAL VOC 2012에서 수행되었습니다. 모델로는 이미지넷에서 사전학습한 ResNet/DeepLabv3+를 사용했으며 수식 10의 λCE, λCE는 각각 16과 8로 설정해주었다고 합니다. temperature parameter는 0.1로 고정되었습니다.

기존의 AL for segmentation과 같이 클릭 수를 labeling cost로 가정하지만, 이게 완벽히 통제되어있는 세팅은 아닙니다(라벨링 방법에 따라 클릭 당 요소 시간이 차이나기도 한다고 하네요). 여기서는 cityscape와 pascal에 대해 각각 round당 100k 및 10k클릭의 budget을 가정하고 5라운드의 연속 데이터 샘플링 및 모델 업데이트를 수행하였습니다. initial round에서 region은 무작위로 선택되고 모델은 각 라운드당 이미지넷 사전학습 가중치로 재초기화됩니다. 모든 실험은 세번씩 수행된 후 평균 성능으로 보고되었습니다.

저자들의 multi-class labeling(Mul)은 다양한 선택 전략으로 dominant class labeling(Dom)과 비교되었습니다. 슈퍼픽셀을 무작위로 선정하는 Random과 uncertainty 기반 BsSB라는 방법, 그리고 BvSB에 추가적인 클래스 밸런싱을 고려한 방법이라고 하는 ClassBal이라는 방법도 있다고 합니다. 저자들의 샘플링 전략은 PixBal이라고 합니다. segmentation쪽 AL은 처음 읽어봐서 전부 처음 보는 방법들이네요.

Figure 4의 결과를 참고하면 multi-class labeling이 대부분 상황에서 dominant class labeling을 능가하는 것을 확인할 수 있습니다. 또한 Pixbal방법의 효율에 대해서도 확인할 수 있는데, cityscapes에서는 PixBal이 일관적으로 budget size에 상관없이 모든 다른 샘플링 전략들을 능가하는것을 확인할 수 있습니다.

table 1에서는 multi-class labeling과 다른 쿼리 디자인에 대해서 비교를 수행했습니다. Patch+polygon은 이미지 패치에 폴리곤 마스크를 그린 것이고, Spx+Dominant는 superpixel에 dominant class법을 조합한 것입니다. 이 실험에서는 사용된 클릭의 비율을 나타낸 것으로, fully-supervised model의 mIoU의 95% 성능을 달성하는데 필요한 각 방법의 클릭 비율을 나타냈습니다. segmentation분야 AL에서는 어노테이션 효율을 이렇게 측정하는 것으로 보이는데, 저자들이 제안한 방법은 기존 방법론들보다 현저히 적은 클릭 수를 필요로 했습니다.

Table2는 저자들이 제안한 프레임워크의 각 요소들에 대한 ablation study입니다. 2개의 loss functionrhk, localization 및 expansion을 모두 사용하는게 가장 좋은 성능을 달성할 수 있음을 확인할 수 있습니다.

Figure6은 수도라벨링에 대한 정성적 평가입니다. 저자가 제안하는 방법의 pixel-wise pseudo labeling과 기존의 region-wise dominant class labeling을 비교했는데, 저자가 제안하는 방법이 주어진 영역에서 비교적 넓은 영역의 다양한 클래스를 정확하게 다룰 수 있음을 볼 수 있습니다.

Conclusion

저자들은 segmentation의 Active Learning을 위한 새로운 Multi-class label query 전략을 제안했으며, 새로운 2단계 학습 전략 및 습득함수는 semantic segmentation의 두 벤치마크에서 좋은 성능을 보이는데 성공했습니다.

AL 분야의 논문을 읽는게 오랜만이고, 특히 segmentation 분야는 처음 읽어봤는데 굉장히 어렵네요. 읽는데 상당히 힘들었습니다. 특히 segmentation은 저에게 낯선 task라 중간중간 나오는 프로토콜들이 낯설어서 그냥 그렇구나..하고 넘어간것들도 많았습니다. 앞으로 부지런히 followup해야할 것 같습니다.

감사합니다.