논문 리뷰) Knowledge Distillation for Efficient Instance Semantic Segmentation with Transformers [CVPR 2024 Workshop]

Knowledge Distillation for Efficient Instance Semantic Segmentation with Transformers [CVPR 2024 Workshop]

https://openaccess.thecvf.com/content/CVPR2024W/Vision4Ag/html/Li_Knowledge_Distillation_for_Efficient_Instance_Semantic_Segmentation_with_Transformers_CVPRW_2024_paper.html

CVPR 2024 Open Access Repository

 

0. Abstract

• Instance Semantic Segmentation

각 객체(Instance)를 구분하여 의미 정보를 픽셀 단위로 제공. 정밀한 인식이 중요한 자율주행, 로봇 분야에서 필수적임.

Semantic vs Instance Segmentation

 그러나 Mask2Former와 같은 SOTA 모델들은 계산 비용이 매우 크며, 높은 정확도를 유지하면서도 계산 부담을 줄이는 것이 해당 Task의 도전과제로 남아있다.

 

 

• Knowledge distillation

 Knowledge Distillation는 큰 모델((teacher))의 지식을 작은 모델((student))에 전달하여 경량화하는 방법인데, 쿼리 기반 구조의 Mask2Former에서 그 방법은 많이 연구되지 않았다. 해당 연구는 Transformer 기반 학생 모델과 교사 모델의 출력 쿼리를 매칭하여 쿼리 기반 지식 증류 방식이 가능하도록 한다.

 

 각 모델을 독립적으로 GT값과 매칭하고, 이를 통해 Teacher-to-student 관계를 정의하게 되는데 이러한 방법론을 통해 학생 모델이 더 적은 수의 쿼리를 가지거나, backbone을 Transformer에서 CNN으로 변경하더라도 지식 증류가 가능함을 보인다.

 

 

• 핵심 기여

• Transformer 기반 모델을 사용하는 Instance semantic segmentation에서 지식 증류 기법을 제안한 최초의 연구.
• Teacher 모델 없이 학습했을 때보다, 지식 증류를 통해 student 모델이 평균적으로 성능 향상을 달성

 

---

1. Introduction

• Instance segmentation # 효용 및 한계

 Instance segmentation는 픽셀 단위 분류 및 객체의 위치 정보와 같은 풍부한 정보를 제공하며, 이는 농업 효율성 향상에 기여하고 인력 중심의 농업을 자동화로 전환할 수 있도록 한다.

 

 하지만 아직까지 조명 조건 그리고 시각적 혼잡 등의 복잡하고 예측 불가능한 요소들이나, 로봇의 제한된 연산 및 전력 자원은 비전 알고리즘을 적용하는 데 제약을 준다.

 

따라서 효율적이고 정확한 Instance segmentation 모델은 실제 현장 적용을 위해 필수적이다.

 

 

• CV에서의 Transformer 기반 모델 # 정확성 챙기기..

 NLP에서 Transformer가 주로 쓰였지만, 이제는 컴퓨터 비전에도 확장되고 있다. DETR는 Transformer 기반의 최초 객체 탐지 모델 중 하나였으며, 하나의 Instance에 해당하는 정보만 담은 object query 개념을 도입하였다.

 

 최근에는 Mask2Former라는 구조가 제안되었으며, 예측된 마스크 영역 내에서 masked attention을 추출함으로써, SOTA 성능의 semantic, instance, panoptic segmentation을 모두 수행할 수 있다.

 

 Transformer 기반 모델은 정확도를 향상시키긴 했지만, 복잡한 구조로 인해 연산 복잡도가 높아지고, 실시간 환경에서의 적용에 제약이 생긴다. → 정확도는 챙겼지만, 효율성은 챙기지 못했다.

 

 

• Knowledge Distillation # 효율성 같이 챙기기..

딥러닝 모델의 효율성을 향상시키는 가장 효과적인 방법 중 하나는 knowledge distillation이며, 고성능 모델의 특성을 경량 모델에 전달하여 효율성을 높이는 방식이다. 최종 또는 중간 레이어의 특징(feature)을 이용하는 방법, 특징들 간의 상호 관계를 활용하는 방법 등 다양한 방식이 존재한다.

하지만 지식 증류 방식 대부분은 CNN 기반 구조에 맞춰 설계되어 있어, 이러한 Transformer 기반 모델에 적용하기 어렵다. 그렇기에 해당 논문은 Transformer 기반 구조로부터 Instance-level의 지식을 증류하는 방법을 제시한다.

 

 

• 제시하는 방법론 개요

 해당 논문 목표는 경량화이지만 성능은 그대로 유지하는 것이다. 즉 복잡한 Transformer 기반 모델의 정확도를 유지하면서도 더 효율적인 네트워크로 압축한다.

1) Hungarian algorithm
→ Teacher와 student의 쿼리 간 최적 이분 매칭-bipartite matching 방법을 생성한다.

2) 추가 증류
→ 단순히 출력 결과뿐만 아니라 마스크 정보-mask map까지 함께 학습하여 성능을 높인다.

 

 

• 결과 요약

1) 실험 결과 속도는 2배 이상 빨라졌지만 성능 손실은 매우 작음, 오히려 특정 조건에서는 성능이 더 좋아지기도 한다.

 2) 백본 구조와 쿼리 수가 다른 Transformer 기반 모델 간에 쿼리를 쌍으로 매칭. 
->  Teacher로부터 student에게 쿼리 기반 지식을 증류하는 핵심 기여를 보인다.

 

 

---

2. Related Work

• Instance Segmentation

Instance segmentation는 각 픽셀의 Class와 Instance ID를 부여하는 과정이다. 다음과 같이 크게 두 가지 방법으로 나뉜다.

1) Two-stage # 현재 주로 사용됨 # 한 번 더 Top-down, Bottom-up 두 개로 나뉨
Top-down: BBOX를 예측한 후, 그 안에서 Instance mask를 생성 #
-> 디택션 성능에 의존하게 된다.

Bottom-up: 픽셀을 분류한 후, 후처리를 통해 인스턴스를 구성함. 
-> 후처리 성능에 의존하게 된다.

2) Single-stage # 탐지와 분할 '동시'에 수행 # 한 번 더 Anchor-based, Anchor-free 두 개로 나뉨
Anchor-based: 탐지기가 사전에 정의된 anchor에 의존. 앵커는 데이터셋에 따라 크기가 달라짐.
-> 겹치는 박스를 제거하는 NMS는 연산이 많고 최적화가 어려움.

Anchor-free: 예측된 위치와 형태를 기반으로 인스턴스 구별.
-> 각 그리드가 하나의 위치와 마스크만 예측할 수 있기에 겹쳐 있을 때 성능이 저하됨.

 

 

• Transformer-Based Dense Prediction

 Transformer 기반 구조는 최근 다양한 비전 태스크, 특히 인스턴스 세그멘테이션에도 적용되고 있음.

Swin-Transformer는 대표적인 멀티스케일 피처를 사용하는 backbone 구조, self-attention을 활용해 장거리 관계를 모델링.

Sparse R-CNN 기반의 쿼리 중심 방식, ISTR, 그리고 Dong et al. 의 unified query 방식 등 여러 접근 방식이 제안됨.

이러한 트랜스포머 기반 모델들은 정확도는 높지만, 복잡성과 연산량이 많아 로봇 플랫폼에서 실시간 처리에는 한계가 있음.

 

 

• Knowledge Distillation for CNN

   Knowledge Distillation은 성능을 유지하면서 모델을 압축하는데 효과적인 방법으로 입증되어 있다. 경량화된 Student model은 복잡한 Teacher model로부터 Soft label를 학습하여 성능을 향상시킬 수 있다. 

soft label은 확률 분포로 hard label인 GT 값 대비 정보가 더 풍부해 학습에 도움됨.

  이러한 지식 증류는 CNN 기반 모델에는 전통적으로 response-based knowledge, feature-based knowledge, and relation-based knowledge으로 나뉜다.

1) Response-based knowledge
-> Student model이 Teacher model의 마지막 레이어 아웃풋을 흉내내게 한다.

2) Feature-based knowledge
-> 중간 계층의 Feature를 사용하며, shape을 맞추기 위해 adaptation layer가 필요할 수 있다.

3) Relation-based knowledge
-> 특징 간 또는 샘플 간의 관계에 초점을 맞춘다. ex) 유사도, 거리 정보 전달

 

 이러한 방법들은 CNN에는 효과적이지만, 구조가 본질적으로 다른 트랜스포머 기반의 모델들에는 적용이 어렵다. 이러한 문제 때문에 self-attention 기반의 지식 증류 방법이 필요하다.

 

 

• Self-attention based knowledge distillation ***

1) Knowledge Distillation via the Target-aware Transformer [CVPR 2022]

https://arxiv.org/abs/2205.10793

Knowledge Distillation via the Target-aware Transformer

receptive field가 다름을 보이는 피규어/ 어텐션 기반 유사도 계산

 해당 방식은 티쳐 모델과 학생 모델의 수용 영역 receptive field이 다르기 때문에, 같은 위치에서도 의미 정보가 다를 수 있다고 가정한다. Teacher 모델의 각 픽셀과 Student 모델의 모든 공간 위치 간의 유사도 Similarity를 계산하는 과정을 추가하였으며, 이 방식은 실험적으로 우수한 성능을 입증하였다.

 

2) Training data-efficient image transformers & distillation through attention - Facebook AI

https://arxiv.org/abs/2012.12877

Training data-efficient image transformers & distillation through attention

 해당 논문의 저자들은 지식 증류의 가장 큰 영향을 주는 요소는 Teacher model의 성능이 아니라, Inductive bias라고 밝혔다. 티쳐 모델의 Inductive bias를 함께 전달하면 학생 모델은 더 다양한 지식을 학습할 수 있게 된다.

 

3) Crossarchitecture knowledge distillation [CVPR 2022]

https://arxiv.org/abs/2207.05273

Cross-Architecture Knowledge Distillation

 해당 논문은 아키텍쳐 간의 간극을 극복하기 위해,  학생의 feature를 교사의 attention 공간에 정렬한다.

 

4) TinyViT: Fast Pretraining Distillation for Small Vision Transformers [ECCV 2022]

https://arxiv.org/abs/2207.10666

TinyViT: Fast Pretraining Distillation for Small Vision Transformers

 교사 모델이 복잡할 경우 많은 리소스를 소모하므로, 교사의 출력을 미리 저장해두는 프레임워크를 제안한다. 

 

 하지만 위에서 언급하는 지금까지의 지식 증류 방법론들은 이미지 수준 또는 픽셀 수준 분류에 초점을 맞추며, Instance segmentation이나 dense prediction에 대한 증류 방식은 아직 충분하게 연구되지 않았다.

 

 해당 논문은 bipartite matching 방식으로 교사에서 학생으로 정보를 증류하여, dense prediction 문제를 다룬다.

cf. dense prediction: 모든 픽셀의 레이블을 예측하는 Task 

 

 

---

3. Proposed Approach

 이 논문에서는 Transformer 기반 모델을 활용한 Instance segmentation을 위해 query 중심의 knowledge distillation 방식을 제안한다. Teacher와 Student의 query 간 정렬된 매칭을 구성하고, 이를 기반으로 query-level 지식을 증류하게 된다.

• Teacher and Student Network Architectures ((3.1))

Mask2Former overview.

 SOTA 모델인 Mask2Former를 프레임워크로 사용하며, 이는 self-attention layer를 포함한 encoder-decoder 구조이다.

Mask2Former 구조 정리
1) Backbone: 이미지로부터 정보를 추출하고, 다양한 해상도의 feature를 출력한다. ResNet, Swin-Transformer를 사용함.

2) Pixel Decoder: 백본에서 나온 feature를 샘플링하여, 고해상도의 픽셀 단위 임베딩을 생성한다.

3) Transformer Decoder: self-attention을 활용하여 이미지 정보의 장거리 관계를 처리하고, object query를 생성한다.

-> 해당 구조로 최종 N개의 쿼리를 생성하여 클래스 확률 및 픽셀 단위 마스크를 예측한다.

 

 해당 연구에서 성능이 가장 좋은 Swin-L을 teacher의 백본으로 설정하고, 계산 효율성이 높은 ResNet-50 또는 Swin-T를 student의 백본으로 사용한다.

 

 Teacher 백본은 장면 내 더 깊은 관계를 학습할 수 있도록 선택되었으며, 200개의 query를 사용한다. 이와 달리 Student 백본이 더 단순하므로, query 수를 100개로 줄였다.

 

 

• Student-Teacher Query Matching *** ((3.2))

 지식 증류를 위해서는 teacher와 student 모델 사이에서 매칭되는 query들을 찾아야한다.  이를 위한 과정을 알아보자.

 

-> teacher-student 사이의 매칭 관계 ϕˆ 도출

 먼저 Teacher- GT, Student - GT query 연관성을 각각 계산한다. 이를 통해 간접적으로 teacher ↔ student 대응을 얻는다.

해당 매칭을 위해 Hungarian 알고리즘을 사용하며, query 간 유사도 기반 cost를 최소화하는 매칭을 찾는다.

 객체가 있는 경우에만, 정답 클래스((C_i))에 대한 예측 확률을 손실로 계산하며, Dice Loss와 Cross Entropy Loss를 함께 사용하여 마스크 예측 손실을 계산한다.

Teacher - Student Query 정렬 정리 ***

- 쿼리는 "찾고자 하는 객체" 이다
즉, 쿼리는 이미지 안에서 “내가 찾고자 하는 물체가 있을까?” 하고 묻는 질문이다.
ex) 쿼리 1 → “이 이미지에 사람 있니?”
ex) 쿼리 2 → “강아지는 있니?” 
즉, Mask2Former에서는 여러 개의 쿼리를 이미지에 던져서 물체를 찾아내는 방식

- 왜 Student와 Teacher의 쿼리 정렬이 중요한가?
쿼리는 객체를 찾기 위한 단위니까 → 서로 같은 객체에 대해 질문하는 쿼리끼리 잘 매칭시켜야 함

- 그래서 어떻게 정렬할 건데? # 간접적으로 정렬
GT를 기준점으로 삼아서, Teacher ↔ GT, Student ↔ GT 간 매칭을 먼저 하고,
→ 그 결과를 통해 Teacher ↔ Student 쿼리 정렬을 간접적으로 유도한다.

이를 계산하는게 Hungarian 알고리즘이며 Cost는 위에서 제시한 Loss 함수가 들어가게 된다.

예를 들어:

• Teacher 쿼리 5번 ↔ GT 객체 2번
• Student 쿼리 3번 ↔ GT 객체 2번
  → 이 경우 Teacher 쿼리 5번과 Student 쿼리 3번은 같은 객체를 가리키는 쿼리로 간주됨

 

 

• Knowledge Distillation Loss function ((3.3))

 Student와 Teacher model의 쿼리 정렬이 끝났으므로, query 중심의 knowledge distillation을 진행하는 단계이다. Student model은 증류 과정에서 teacher의 정보 뿐 아니라, GT를 모두 사용한다. 다음은 학습에 사용되는 전체 Loss이다.

α 는 데이터셋에 따라 달라지는 가중치

-> Student 모델이 GT값을 잘 예측할 수 있도록, Teacher model를 잘 모방할 수 있도록 하는 Loss라고 생각하면 된다.

 

 1) Groundtruth Loss: GT에 대하여 클래스를 잘 예측했는지, 마스크를 잘 예측했는지의 가중합이다. (Supervised)

 

2) Distillation Loss: 3.2 과정을 통해 매칭한 쿼리가 잘 정렬되었다고 생각하고, 이제 그에 대응하는 Teacher의 예측 결과를 Student가 모방하도록 학습하는 로스이다. (Distillation)

 

 

---

4. Experimental Setup and Results

• 다양한 Task에 대한 실험 결과

BUP20/ SB20/ Cityscape

st:  Swin-Tiny 기반/  r50:  ResNet50 기반

 세 가지 데이터셋에서 평균적으로 Sr50 모델은 AP 기준 1.8점, Sst 모델은 1.9점 향상되었으며, 추론 속도는 각각 2.3배, 2.0배 빠름.

-> 복잡한 Teacher 모델을 통해 기존 ResNet, Swin-T가 도달할 수 없었던 SOTA 수준의 성능을 distillation으로 학생 모델에서 구현할 수 있었고, 동시에 빠른 추론 속도도 확보할 수 있었다.  

 

---

5. Summary

• 본 논문에서 Mask2Former의 쿼리 기반 출력에 대해 bipartite matching을 적용한 knowledge distillation 기법을 제안함.
• Teacher와 Student의 쿼리를 Hungarian 알고리즘을 통해 정렬하고, 이에 따라 클래스 확률과 인스턴스 마스크를 효과적으로 증류.
• 학생 모델로는 ResNet-50과 Swin-Tiny를 사용하였고, 구조가 전혀 다른 백본에서도 적용 가능함을 입증함.
• 실험 결과, 모든 데이터셋에서 AP 기준 평균 1.8~1.9 포인트 향상되었고, 특히 AP75 향상이 두드러졌음.
• 따라서 Teacher의 성능을 유지하면서도 더 빠른 추론 속도를 가진 Student 모델을 만들 수 있음을 입증함.
• 추후 연구에서는 백본뿐만 아니라 Pixel Decoder 교체에 따른 효율-정확도 간 Trade-off를 추가로 탐색할 수 있음.