[Paper review] Vision Transfomers for Dense Prediction

Computer Vision 분야에 transfomer을 적용해 classification task에서 SOTA를 달성한, Vision Transfomer는 큰 파장을 일으켰다. 그리고 classification 외에 다른 task에도 이를 적용하려는 움직임이 나타났다. 오늘 리뷰하는 논문에서 소개하는 모델인 DPT(Dense Prediction Transfomer) 역시 이러한 움직임 속에서 나타난 모델로써, Dense Prediction하는 작업에서 Vision Transfomer를 어떻게 적용할 수 있을까라는 질문에서 시작했다. 

 

Dense Prediction에서 좋은 성능을 내는 모델들은 다음과 같은 공통된 특징을 가지고 있었다. 첫번째는 encoder-decoder 구조를 가지고 있다는 점이고, 두번째는 CNN backbone을 encoder에서 쓴다는 점이다. 여기서 주목해야 할 점은 두번째 특징이다. encoder에서 CNN backbone을 쓸 경우, downsampling 하는 과정을 필연적으로 거치게 된다. 왜냐하면 메모리상의 한계가 명백하게 있기 때문에, input image의 resolution을 그대로 가져가면서 층을 쌓기엔 현실적으로 불가능하기 때문이다. 하지만 이러한 과정을 거칠경우, Dense Prediction에서 중요한 resolution과 관련된 정보들을 잃게되는 문제가 발생한다. 이후 dilated convolution과 skip connection 등 다양한 방법들이 시도되었지만, 문제를 완전히 해결 할 수는 없었다. (성능이 향상시킬 수는 있었음)

 Vision Transfomer를 사용해 Down sampling을 해결

DPT는 이러한 문제를 근본적으로 해결하기 위해, encoder backbone으로써, 'Vision Transfomer'를 사용한다. Transfomer는 CNN과 달리 downsampling과정이 없기 때문에, 초기 input image의 resolution 정보를 계속해서 가지고 갈 수 있다. 추가로, Transfomer를 사용할 경우, MHSA operation을 통해 CNN과 달리 초기부터 global 한 receptive field를 계속 갖게된다. 이러한 특징은 곧 이미지에 대한 global하게 일관된 예측에도 도움이 된다고 논문에서는 설명한다. 

논문 참조

DPT의 구조를 좀 더 자세히 살펴보자. 구조자체는 크게 encode-decoder의 구조를 따른다. encoder는 앞서 설명했듯이 Vision Transfomer을 그대로 사용한다. 참고로 Vision Transfomer에서 token embedding 이후 하나의 special token을 더해준 과정까지 DPT에서도 그대로 따른다. (위 그림에서 빨간색 막대기) decoder는 기본적으로 Convolution decoder를 사용한다. (Fusion 과정 참조)

DETR에는 Reassemble이라는 과정이 있다. 위 그림을 통해 확인할 수 있듯이 Reassemble 과정을 통해 나온 output이 Fusion의 input으로 들어가 순차적으로 합쳐진다는 것을 알 수 있다. Reassemble 과정은 크게 Read -> Concatenate -> Resample으로 이루어져 있고, 이 과정을 통해 Transfomer의 특정 깊이의 layer에서 나온 토큰들을 조합하여 특정 해상도의 feature representation을 만들어 낸다고 생각하면 된다. 가령 위 그림에서 마지막 Transfomer layer을 거친 토큰들은 Reassemble32을 거치는 것을 알 수 있는데, 이는 곧, 토큰들을 input image의 1/32의 해상도를 가지는 feature representation으로 만들어준다는 뜻이다. 최종적으로 DETR는 prediction head를 거치면서, Dense를 예측하게 된다. 

DPT의 성능은 위 표를 보면 알 수 있듯이 다양한 Dense prediction task에서 SOTA를 달성한다는 것을 알 수 있다. 이외에도 작은 데이터셋에도 fine tuning했을때도, SOTA를 달성했고(논문참조) 이는 DPT의 강력한 representatinal power를 보여준다. 

추가로 논문에서는 Semantic Segmentation에도 DPT를 적용했다. (Computer Vision에서 Dense prediction에 Semantic Segmentation은 포함되는 작업이다.) 결과는 쉽게 예상할 수 있듯이, 두 가지 실험한 데이터셋 모두에서 SOTA를 달성했다. 이외에도 기존 CNN모델에 비해 파라미터 수가 3배이상 많은데 inference speed는 거의 차이가 없다는 결과 등 여러측면에서 기존 CNN backbone 기반 모델들을 상회하는 성능지표을 보여준다는 것을 확인할 수 있었다.

 

이번논문을 읽고서 Vision Transfomer을 활용한 모델이  Dense prediction에서 좋은 성능을 낸다는 사실을 알 수 있었다.  요즘 Computer Vision의 뜨거운 감자이기도 한 Vision Transfomer의 엄청난 잠재력을 다시 한번 느낄 수 있었던 논문인 것 같다. (왜 관련 최신논문들이 쏟아지는지 알겠다...)

 

리뷰한 논문의 주소는 다음과 같다.

https://arxiv.org/abs/2103.13413