Semantic segmentation task에서 incremental learning은 dense annotation에 기반해서 새로운 카테고리를 학습한다. 하지만 pixel-by-pixel 라벨링은 비용이 많이들고 time-consuming하다.
Weakly incremental learning for semantic segmentation(WILSS)는 image-level annotation만 활용하기 때문에 비용이 저렴하고 넓게 적용할 수 있다. 하지만 이런 image-level label은 segment 각각을 locate 하는 디테일한 정보를 주지 못해서 성능이 제한적이다.
따라서 어떻게 image-level labels만 주어진 new classes를 활용해서 forgetting을 피하고 효과적으로 성능을 높일 수 있을지를 고민했다.
본 논문에서는 teacher-student 구조를 활용해서 noisy pseudo masks를 최적화해간다.
그리고 memory-based copy-paste augmentation을 활용해서 catastrophic forgetting 문제를 해결한다.
Introduction
Sematic segmentation의 성능이 매우 좋음에도 불구하고 하나의 데이터셋으로 pre-train 된 하나의 모델을 가지고 새로운 클래스를 가지는 다른 데이터셋에 retrain 시키면 학습한 지식을 까먹는 문제가 있다.
이러한 문제를 catastrophic forgetting이라고 하며, 과거 샘플에 접근하지 않고 새로운 카테고리로 모델링할 때 모델의 파라미터가 너무 크게 바뀌는 문제로 인해 발생한다.
이를 해결하기 위한 접근법은 incremental learning이다. Incremental Learning Semantic segmentation(ILSS)에서는 현재 스텝에서 학습 샘플의 new class pixel annotation으로 예측하고 이전 스텝에서의 old class는 background로 예측되는 식이다.
현대 ILSS에서는 크게 두 가지 접근 방식으로 나눌 수 있는데, 1) Regularization-based, 2) Replay-based이다.
1) Regularization-based 방법은 knowledge distillation 방식이다.
2) Replay-based 방법은 memory를 따로 두고 저보를 저장해뒀다가 새로운 학습 스탭에서 활용하는 방식이다.
그러나 위 두 방식의 한계는 new class에 대한 pixel level annotation이 필요하다는 점이다. 그래서 WILSON은 최초로 weakly incremental learning for semantic segmentation task를 제안했지만 이미지 레벨의 라벨이 유의미한 성능 향상에 도움이 되지 못했다.
본 연구에서는 old 데이터에 대한 지식을 보존하면서 이미지 레벨의 label이 주어진 new class에 대해 효과적으로 supervision을 활용할 수 있는 방법을 연구했다.
일단, fig1처럼 최초로 pre-trained foundation model을 활용해서 training-free 방식으로 주어진 image-level label supervision을 개선했다.
Pre-training based co-segmentation을 활용해서 vision-language pretraining model(CLIP)의 knowledge를 distill한다.
먼저 clip으로 category-aware dense mask generation을 통해 initial mask를 뽑아낸다. 그럼 각 new class의 initial mask를 가지고 self-supervsed model를 활용한 seeds guidance로 compact category-agnostic attention map을 뽑아내서 mask fusion을 통해 pseudo mask를 refine한다. 그래도 여전히 noisy한 pseudo mask를 최적화하기 위한 teacher-student구조를 제안하는데, plug-in teacher는 dense contrastive loss로 최적화된다. 이러한 방법으로 pseudo dense supervision을 더 효과적으로 활용할 수 있다.
Related Work
3. Method
일단, incremental learning에는 multiple step(t)가 존재한다. 하지만 original incremental setting 과는 조금 다르게 novel weakly incremental learing for semantic segmentation (WILSS) 셋팅에서는 t=0 즉, 초기 스텝에서만 pixel annotations가 주어진다. 이후 스텝에서는 new class에 대해 오직 image-level annotations로만 접근할 수 있고 이전 스텝에서의 training sample에는 더이상 접근할 수 있다. 목표는 old class에 대한 forgetting 없이 모델을 학습시키고 업데이트하는 것이다.
Pre-training Based Co-segmentation
오직 image-level label만 사용해서 dense prediction tasks를 supervise하기에는 여전히 한계가 있다.
Q. 그렇다면 어떻게 image-level abel로 새로운 클래스에 대한 supervision을 개선할 수 있을까?
본 논문에서는 pre-training based co-segmentation method를 활용해서 두 종류의 foundation model(vision language pretrained model로는 CLIP, self-supervised pretrained model로는 iBOT, DINO)의 knowledge를 distill한다.
전체적인 과정은 위와 같다.
Initial Mask
Pre-trained vision-language model(CLIP)에 태워서 주어진 image-level labels로 category-aware pixel annotations를 뽑는다. Fig 2처럼 구체적으로 보면 t시점에 이미지가 image-level label인 C^T와 함께 주어지면 처음에 CLIP encoder(f_I)로부터 image feature(F)를 뽑고, final linear projection layer(f_proj)에 통과시킨다.
여기서 L2 normalization은 channel 차원으로 진행되고 d는 feature dimension을 의미한다.
그리고 나서, t step에서의 target class C^T를 input prompt로 해서 text embedding을 계산한다.
위 식에서 f_T는 CLIP text encoder를 말하고, prompt는 CLIP에서 multiple prompt templates를 ensemble하는 prompt engineering을 의미한다.
그러고 나서 language-compatible image feature embedding, text embedding으로 pixel-text score map를 계산한다. (dot product 연산)
M_init은 initial segmentation results가 될 것이다.
Refine Mask via Seeds Guidance
앞에 CLIP을 통해 만든 pseudo mask M_init은 rich-category-aware pixel annotations 일 수 있지만, image-text pairs로 학습된 CLIP은 segmentation task보다 instance-level classification에서 유리하기 때문에 mask가 noisy하다. 이 마스크 퀄리티를 높이기 위해, local feature modeling에서 좋은 성능을 낸 다른 종류의 foundation model(self-supervised pre-traininig models)에 knowledge를 distill하는 것을 제안한다. 이 모델들은 compact category-agnostic attention map을 얻을 수 있다.
Q. 하지만, 이미지가 여러 개의 object를 포함할 수 있는 경우에는 어떻게 해야 할까?
-> 이 문제를 다루기 위해 category-specific seeds guidance로 initial mask를 refine할 것을 제안한다. 우선 각 target class마다 initial mask에서 랜덤 시드 포인트를 고르고 pre-trained self-supervised model로부터 해당하는 category-aware attention map을 뽑는다. 이 때의 attention map은 last self-attention block에서 뽑아온다.
여기서 n은 attention head의 수를 의미한다. N은 seed point 수를 의미하며 각각의 새 클래스, training step마다 binarized M_init의 foreground로부터 랜덤으로 N개만큼 샘플링한다. ([]는 binarization operation을 의미하며, 특정 threshold보다 크면 1, 아니면 0으로 설정한다. 이 때의 threshold는 averaged attention map이 top K(k=70%) locations를 유지하며 dynamically updated된다.)
FIg 5는 2개의 class(horse, dog)에 대한 attention map을 시각화한 것인데 시드 수(N) = 9 일 때 좋은 clustering performance를 보였다.
마지막으로, 간단한 mask fusion으로 union연산 거쳐서 refine 된 마스크를 얻는다.
3.3 Pseudo Label Optimization
여전히 noisy한 pseudo mask를 더욱 최적화하기 위해 teacher-student구조를 사용할 것을 제안한다.
구체적으로, segmentation model은 student model로 함으로써, 학습 중에 더 나은 pseudo mask를 dynamically learn하기 위해 plug-in-teacher module(ASPP)를 도입한다. teacher module를 학습하기 위해 먼저 pixel-wise binary cross-entropy loss를 사용할 것을 제안한다.
BCE는 t 단계에서 new clas prediction을 supervise하고, 이전 클래스 최적화는 다음 섹션에 남긴다.
여기서 P^c,i는 pixel i의 new class c에 대한 predicted probability를 의미한다. 그러나 pixel-wise BCE loss는 주로 현재 이미지에 대한 foreground class c의 최적화에 초점을 맞추고 다른 클래스들은 모두 배경으로 처리하기 때문에 픽셀 간의 상관관계를 무시한다. 따라서 unsupervised representation learning에서 쓰이는 InfoNCE loss에서 영감을 받아서 전체 데이터셋 간에 multi-class predictions와 해당하는 pixel-wise pseudo labels을 더 잘 활용하기 위해 dense contrastive learning을 수행할 것을 제안한다.
fig3과 같이 새로운 클래스 이미지의 픽셀 i와 해당 pseudo annotation에 대해 i와 동일한 class를 가지는 모든 픽셀을 수집하여 positive sample Pi를 구성하고, 다른 클래스 point들은 negative sample Ni로 구성했다. 이 dense contrastive loss는 아래 식 7과 같이 계산된다.
타우는 temperature term으로 0.1을 디폴트로 한다. 학습 효율성을 위해 현재 미니배치에 포함된 class들에 대해 ten point만 랜덤으로 샘플링해서 계산한다.
pixel-wise BCE loss(식 6)과 dense contrastive loss(식 7)은 서로 보완적이다. teacher module이 discriminative pixel features를 잘 학습할 뿐만 아니라 intra-class와 inter-class pixel feuatre modeling을 통해 pixel feature regularize 또한 할 수 있다.
3.4 Memory-based Copy-Paste Augmentation
추가로 new class에 대한 supervision을 개선하고 효과적으로 활용하는 것 외에도, 학습을 안정화하고 segmentation model 성능을 더욱 향상시킬 수 있는 memory-based copy-paste augmentation strategy를 제안한다.
Fig4처럼, 각각의 old class마다 memory bank를 두고 각 class archive는 base model training 중에 B개의 foreground instaces와 segmentation labels를 저장한다. 그 다음 t 단계에서 랜덤으로 선택된 이전 class archive에서 foreground image와 label 쌍을 랜덤으로 선택해서 new class image에 랜덤으로 paste한다. 이제 training sample에는 t 단계에서 new class image 뿐만 아니라 t-1 단계에서 old class image와 pixel label이 포함되어 있다. 따라서 teacher module의 old class lerning을 다음과 같이 최적화한다.
g는 logistic function이다.
3.5. Overall Optimization
이전 단계(t-1)에서 학습된 모델과 t시점에서 동적으로 업데이트 된 teacher module에 대한 knonwledge를 distill해서 segmentation model을 최적화한다. 이 teacher module은 주로 BCE loss로 최적화되기 때문에 우리는 BCE loss를 사용해서 teacher module의 prediction을 segmetation model로 distill한다. 학습된 pseudo mask가 완벽하지 않다는 점을 고려해서 soft pixel label을 new class에 대한 최종 supervision으로 사용하고, 이전 모델과 teacher module 출력의 가중 평균 값을 old class에 대한 supervision으로 사용한다.
여기서 hard는 one-hot operation을 의미해서 각 픽셀 별로 maximum score를 설정하고 나머지는 0으로 한다. 알파와 베타는 trade-off 파라미터이고 알피는 0.5, 베타는 0.9를 디폴트로 설정했다.
segmentation model의 BCE loss는 아래와 같다.
y^t는 all seen classes를 의미하고 p는 t시점에서의 segmentation model의 output을 의미한다.
따라서 최종 objective는 다음과 같다.
(loss 총정리)
Experiments
4.1. Datasets and Protocols
공정한 비교를 위해 SOTA 실험 셋팅과 데이터셋, 프로토콜을 따랐다. new class에 대해 pixel-wise annotation을 사용했던 기존 연구들과 달리 우리는 novel class에 대해 image-level labels만 사용했다.
Datasets : 2개의 standard evaluation benchmarks PascalVOC2012, COCO를 사용했다.
Protocols: 이전 연구들은 disjoint , overlap 이 2개의 incremental learning protocols를 소개한다. disjoint는 각 training step에서는 이전에 본 클래스랑 현재의 클래스만 포함한다. overlap인 경우에는 각 training step이 모든 이미지를 포함하고 픽셀은 어느 class에도 속할 수 있다. 따라서 overlap protocol이 disjoint에 비해 더 현실적이고 challenge한 task이다.
15-5 VOC: 15 base class는 처음 training step에서 학습되고, 5개의 new classes는 second step에서 연속적으로 학습된다.
10-10 VOC: 10 base class는 처음 스텝에서 학습되고 나머지 10개의 new class는 두번째 step에서 추가된다.
COCO-to-VOC: 첫 step에서는 60개의 COCO class로 학습하고(VOC랑 겹치는 20개 버림). 그리고 연속적으로 VOC의 20개의 class를 학습시킨다.
평가지표로는 mIoU를 썼다.
4.2. Implementation Details
VOC에 대해서는 image pretrained Deeplab V3 with a Resnet-101 backbone을 썼고 COCO에 대해서는 image pretrained Wide-ResNet-38 dataset을 썼다. 모든 model은 40에폭에 배치사이즈 24로 학습했고 optimizer로는 SGD, lr는 1e^-3, momentum 0.9, weight decay 1e^-4로 학습했다. segmentation model을 학습하기 전에 warm up 과정을 거쳐서 teacher module을 먼저 5에폭 학습시켰다. 람다는 0.1, 알파는 0.5, 베타는 0.9로 학습시켰다.
이전 연구들과 공정한 비교를 위해 똑같이 5개의 class(plant, sheep, sofa, train, tv monitor)를 incremetal step으로 두었다. table 1을 보면 disjoint, overlap 셋팅 모두에서 SOTA를 달성했다.
10-10 VOC나 COCO-to-VOC 에서도 마찬가지로 SOTA를 달성했다.
4.5. Ablation Studies
M_seeds를 계산할 때 number of seeds를 바꿔가면서 실험한 결과이다. 더 많은 random seeds를 guidance가 더 compact한 attention maps를 만들어주고 9개정도가 좋은 학습 효율성과 함께 좋은 클러스터링 결과를 보여주었다. 그리고 category-agnostic M_seeds가 M_init을 보완해준다는 걸 red box로 표시하고 있다.
mask fusion의 두 가지 operation을 비교하고 있고 fusion 없는 경우도 비교하고 있다. Union fusion 방식이 눈에띄는 성능 향상을 보였다. Table 5에서는 람다를 0.1로 셋팅한 이유에 대한 ablation이다. 그리고 0인 것과도 비교한 결과, 우리가 제안하는 Dense contrastive learning 방식이 더 효과적으로 촤적화할 수 있고 dense supervision을 더 잘 할 수 있다는 것을 보여준다.
Abstract
Introduction
Sematic segmentation의 성능이 매우 좋음에도 불구하고 하나의 데이터셋으로 pre-train 된 하나의 모델을 가지고 새로운 클래스를 가지는 다른 데이터셋에 retrain 시키면 학습한 지식을 까먹는 문제가 있다.
이러한 문제를 catastrophic forgetting이라고 하며, 과거 샘플에 접근하지 않고 새로운 카테고리로 모델링할 때 모델의 파라미터가 너무 크게 바뀌는 문제로 인해 발생한다.
현대 ILSS에서는 크게 두 가지 접근 방식으로 나눌 수 있는데, 1) Regularization-based, 2) Replay-based이다. 1) Regularization-based 방법은 knowledge distillation 방식이다. 2) Replay-based 방법은 memory를 따로 두고 저보를 저장해뒀다가 새로운 학습 스탭에서 활용하는 방식이다. 그러나 위 두 방식의 한계는 new class에 대한 pixel level annotation이 필요하다는 점이다. 그래서 WILSON은 최초로 weakly incremental learning for semantic segmentation task를 제안했지만 이미지 레벨의 라벨이 유의미한 성능 향상에 도움이 되지 못했다.
본 연구에서는 old 데이터에 대한 지식을 보존하면서 이미지 레벨의 label이 주어진 new class에 대해 효과적으로 supervision을 활용할 수 있는 방법을 연구했다.
일단, fig1처럼 최초로 pre-trained foundation model을 활용해서 training-free 방식으로 주어진 image-level label supervision을 개선했다.
먼저 clip으로 category-aware dense mask generation을 통해 initial mask를 뽑아낸다. 그럼 각 new class의 initial mask를 가지고 self-supervsed model를 활용한 seeds guidance로 compact category-agnostic attention map을 뽑아내서 mask fusion을 통해 pseudo mask를 refine한다. 그래도 여전히 noisy한 pseudo mask를 최적화하기 위한 teacher-student구조를 제안하는데, plug-in teacher는 dense contrastive loss로 최적화된다. 이러한 방법으로 pseudo dense supervision을 더 효과적으로 활용할 수 있다.
Related Work
3. Method
일단, incremental learning에는 multiple step(t)가 존재한다. 하지만 original incremental setting 과는 조금 다르게 novel weakly incremental learing for semantic segmentation (WILSS) 셋팅에서는 t=0 즉, 초기 스텝에서만 pixel annotations가 주어진다. 이후 스텝에서는 new class에 대해 오직 image-level annotations로만 접근할 수 있고 이전 스텝에서의 training sample에는 더이상 접근할 수 있다. 목표는 old class에 대한 forgetting 없이 모델을 학습시키고 업데이트하는 것이다.
Pre-training Based Co-segmentation
본 논문에서는 pre-training based co-segmentation method를 활용해서 두 종류의 foundation model(vision language pretrained model로는 CLIP, self-supervised pretrained model로는 iBOT, DINO)의 knowledge를 distill한다.
Initial Mask
그리고 나서, t step에서의 target class C^T를 input prompt로 해서 text embedding을 계산한다.
위 식에서 f_T는 CLIP text encoder를 말하고, prompt는 CLIP에서 multiple prompt templates를 ensemble하는 prompt engineering을 의미한다.
그러고 나서 language-compatible image feature embedding, text embedding으로 pixel-text score map를 계산한다. (dot product 연산)
M_init은 initial segmentation results가 될 것이다.
Refine Mask via Seeds Guidance
앞에 CLIP을 통해 만든 pseudo mask M_init은 rich-category-aware pixel annotations 일 수 있지만, image-text pairs로 학습된 CLIP은 segmentation task보다 instance-level classification에서 유리하기 때문에 mask가 noisy하다. 이 마스크 퀄리티를 높이기 위해, local feature modeling에서 좋은 성능을 낸 다른 종류의 foundation model(self-supervised pre-traininig models)에 knowledge를 distill하는 것을 제안한다. 이 모델들은 compact category-agnostic attention map을 얻을 수 있다.
Q. 하지만, 이미지가 여러 개의 object를 포함할 수 있는 경우에는 어떻게 해야 할까?
여기서 n은 attention head의 수를 의미한다. N은 seed point 수를 의미하며 각각의 새 클래스, training step마다 binarized M_init의 foreground로부터 랜덤으로 N개만큼 샘플링한다. ([]는 binarization operation을 의미하며, 특정 threshold보다 크면 1, 아니면 0으로 설정한다. 이 때의 threshold는 averaged attention map이 top K(k=70%) locations를 유지하며 dynamically updated된다.)
FIg 5는 2개의 class(horse, dog)에 대한 attention map을 시각화한 것인데 시드 수(N) = 9 일 때 좋은 clustering performance를 보였다.
-> 이 문제를 다루기 위해 category-specific seeds guidance로 initial mask를 refine할 것을 제안한다. 우선 각 target class마다 initial mask에서 랜덤 시드 포인트를 고르고 pre-trained self-supervised model로부터 해당하는 category-aware attention map을 뽑는다. 이 때의 attention map은 last self-attention block에서 뽑아온다.
마지막으로, 간단한 mask fusion으로 union연산 거쳐서 refine 된 마스크를 얻는다.
3.3 Pseudo Label Optimization
여전히 noisy한 pseudo mask를 더욱 최적화하기 위해 teacher-student구조를 사용할 것을 제안한다.
여기서 P^c,i는 pixel i의 new class c에 대한 predicted probability를 의미한다. 그러나 pixel-wise BCE loss는 주로 현재 이미지에 대한 foreground class c의 최적화에 초점을 맞추고 다른 클래스들은 모두 배경으로 처리하기 때문에 픽셀 간의 상관관계를 무시한다. 따라서 unsupervised representation learning에서 쓰이는 InfoNCE loss에서 영감을 받아서 전체 데이터셋 간에 multi-class predictions와 해당하는 pixel-wise pseudo labels을 더 잘 활용하기 위해 dense contrastive learning을 수행할 것을 제안한다.
fig3과 같이 새로운 클래스 이미지의 픽셀 i와 해당 pseudo annotation에 대해 i와 동일한 class를 가지는 모든 픽셀을 수집하여 positive sample Pi를 구성하고, 다른 클래스 point들은 negative sample Ni로 구성했다. 이 dense contrastive loss는 아래 식 7과 같이 계산된다.
타우는 temperature term으로 0.1을 디폴트로 한다. 학습 효율성을 위해 현재 미니배치에 포함된 class들에 대해 ten point만 랜덤으로 샘플링해서 계산한다.
pixel-wise BCE loss(식 6)과 dense contrastive loss(식 7)은 서로 보완적이다. teacher module이 discriminative pixel features를 잘 학습할 뿐만 아니라 intra-class와 inter-class pixel feuatre modeling을 통해 pixel feature regularize 또한 할 수 있다.
구체적으로, segmentation model은 student model로 함으로써, 학습 중에 더 나은 pseudo mask를 dynamically learn하기 위해 plug-in-teacher module(ASPP)를 도입한다. teacher module를 학습하기 위해 먼저 pixel-wise binary cross-entropy loss를 사용할 것을 제안한다. BCE는 t 단계에서 new clas prediction을 supervise하고, 이전 클래스 최적화는 다음 섹션에 남긴다.
3.4 Memory-based Copy-Paste Augmentation
추가로 new class에 대한 supervision을 개선하고 효과적으로 활용하는 것 외에도, 학습을 안정화하고 segmentation model 성능을 더욱 향상시킬 수 있는 memory-based copy-paste augmentation strategy를 제안한다.
Fig4처럼, 각각의 old class마다 memory bank를 두고 각 class archive는 base model training 중에 B개의 foreground instaces와 segmentation labels를 저장한다. 그 다음 t 단계에서 랜덤으로 선택된 이전 class archive에서 foreground image와 label 쌍을 랜덤으로 선택해서 new class image에 랜덤으로 paste한다. 이제 training sample에는 t 단계에서 new class image 뿐만 아니라 t-1 단계에서 old class image와 pixel label이 포함되어 있다. 따라서 teacher module의 old class lerning을 다음과 같이 최적화한다.
g는 logistic function이다.
3.5. Overall Optimization
이전 단계(t-1)에서 학습된 모델과 t시점에서 동적으로 업데이트 된 teacher module에 대한 knonwledge를 distill해서 segmentation model을 최적화한다. 이 teacher module은 주로 BCE loss로 최적화되기 때문에 우리는 BCE loss를 사용해서 teacher module의 prediction을 segmetation model로 distill한다. 학습된 pseudo mask가 완벽하지 않다는 점을 고려해서 soft pixel label을 new class에 대한 최종 supervision으로 사용하고, 이전 모델과 teacher module 출력의 가중 평균 값을 old class에 대한 supervision으로 사용한다.
여기서 hard는 one-hot operation을 의미해서 각 픽셀 별로 maximum score를 설정하고 나머지는 0으로 한다. 알파와 베타는 trade-off 파라미터이고 알피는 0.5, 베타는 0.9를 디폴트로 설정했다.
segmentation model의 BCE loss는 아래와 같다.
y^t는 all seen classes를 의미하고 p는 t시점에서의 segmentation model의 output을 의미한다.
따라서 최종 objective는 다음과 같다.
(loss 총정리)
Experiments
4.1. Datasets and Protocols
공정한 비교를 위해 SOTA 실험 셋팅과 데이터셋, 프로토콜을 따랐다. new class에 대해 pixel-wise annotation을 사용했던 기존 연구들과 달리 우리는 novel class에 대해 image-level labels만 사용했다.
4.2. Implementation Details
VOC에 대해서는 image pretrained Deeplab V3 with a Resnet-101 backbone을 썼고 COCO에 대해서는 image pretrained Wide-ResNet-38 dataset을 썼다. 모든 model은 40에폭에 배치사이즈 24로 학습했고 optimizer로는 SGD, lr는 1e^-3, momentum 0.9, weight decay 1e^-4로 학습했다. segmentation model을 학습하기 전에 warm up 과정을 거쳐서 teacher module을 먼저 5에폭 학습시켰다. 람다는 0.1, 알파는 0.5, 베타는 0.9로 학습시켰다.
4.5. Ablation Studies
Conclusion