2019년 ICCV에 발표된 논문이다. Segmentation task를 목표로 하고, 기존 2018년 CVPR에 발표된 "Non-local neural networks" (Wang, Xiaolong, et al.)를 후속 연구한 논문이라고 볼 수 있다. 간단하게 발전 과정을 적을까 한다.
CNN(Convolution Neural Network)
우선 "Non-local neural networks"를 보기 전에 CNN의 단점부터 보고 넘어가자. MLP(Multi-Layer Perceptron)에서 CNN으로 발전되면서 이미지의 spatial 정보를 효과적으로 전달할 수 있게 되었다. 모든 task에서 엄청난 성능 증가를 보여주었고, 여전히 많이 사용되고 있는 방식이다. 하지만 충분한 receptive field를 담기에는 한계가 존재했고, 다양한 방식으로 해결하려는 시도가 있었다. Classification task의 경우에는 feature map에 global average pooling을 적용해 해결을 하는 방식을 취하였지만, Segmentation task에는 한계가 있었다.
Convolution Filter
이러한 문제는 전통적인 영상처리 방법인 convolution filter에도 존재했다. noise를 제거하기 위해 Mean filter, Gaussian filter 등 다양한 convolution filter를 적용하였지만, 주변 정보(local한 정보)만 참고한다는 한계가 있었다.
Non-local Mean Filter (NLM Filter)
그래서 탄생한 filter가 non-local mean filter이다. 같은 영상내에서 해당 픽셀 주변 영역과 비슷한 패턴을 갖는 영역들을 찾아내 그 영역들의 정보도 참고하는 방식이다. 그림에서 q1과 q2 부분이 denoising하기 원하는 부분과 비슷한 패턴을 갖고 있으므로 참고를하고, q3는 참고를 하지 않는 것이다. (실제로는 weight를 주어 참고하는 정도를 다르게 한다고 생각하면 편하다.) 이러한 아이디어를 통해 global한 정보를 참고할 수 있게 된 것이다.
Non-local neural Networks
해당 논문에서는 비디오 데이터를 다뤄 time 축 포함 3D 데이터였지만, 우리가 볼 CCNet은 2D 이미지를 다루므로 2D라고 생각하고 보자. THW -> HW
non-local mean filter에 영감을 받아 만들어진 것이 non-local neural networks이다. 내부 구조를 보면 위와 같은데, transformer의 attention을 생각하면 쉽다. θ가 query, Φ가 key, g가 value이다. CNN을 타고 온 Featuer map X는 각 point의 spatial 정보를 담고 있다고 볼 수 있는데, X의 한 point마다(θ:query) 모든 point(Φ:key)들 간의 similarity를 계산하여 다른 point로 부터 얼마만큼의 정보를 얻어 올 것인지 결정하고(similarity가 높을 수록 많은 정보를 얻어온다.), g:value로 부터 정보를 비율에 맞게 취합(aggregation)한다.
위 그림의 빨간 박스 부분이 similarity를 구하는 부분이다. 각 point 간의 similarity를 구할때 feature map X를 그대로 사용할 수도 있지만( similarity를 구한다는 의미상으로는 이 방법이 더 맞지만), 1x1 convolution으로 projection하여 query map θ와 key map Φ를 생성해 주었다. 이렇게함으로써 학습가능한 weight를 추가해준 것이다.
두 벡터의 similarity를 구하는 방법은 다양한데, 계산 편의상 dot product(내적)를 이용하여 구하였다고 한다. 모든 point에 대해 similarity를 구하게되면 HW개의 point 마다 HW개의 similarity map이 생기게 되고 이 map의 총합이 1이되게 해주기 위해서 softmax를 거치게된다. 그후 value map g에서 point 별로 유사한 정도만큼의 정보만 가지고와(similarity map * value map), HW*512의 global 정보를 담고있는 새로운 feature map을 생성한다. 이를 1x1 convolution을 거쳐 원래 dimension으로 돌려주고 local 정보만 가지고있는 feature map X와 aggregation하게 된다.
SENet을 channel attention으로 볼 수 있듯이 non local 정보를 aggregate 하는 과정을 pointwise Attention으로 볼 수 있겠다. 단점은 similarity를 구하는 cost가 크다는 점이다. 대략적으로 계산해봐도, feature map에서 구하기 때문에 HW가 100*100이면 similarity를 구하기 위해 1억번의 dot product를 진행하게된다.
CCNet
드디어 본론으로 돌아오게 되었다. CCNet은 non-local neural network의 pointwise attention 부분의 cost를 줄였고, 정확도도 높였다고 주장한다. 어떻게 가능했는지 살펴보자.
전체적인 module의 구조는 중간중간 디테일한 부분을 제외하고는 non-local neural network의 구조와 동일하다. 이름에서 볼 수 있듯이 십자가 모양을 사용하는데, pointwise attention map을 feature map 전체에서 계산하지 않고, 십자기 형태에서만 계산한다는 것이 이 논문의 핵심이다. 그러면 기존의 HWxHW 번의 계산에서 HWx(H+W-1) 번의 계산으로 계산량이 확 줄게된다.(100x100 기준으로 1/50 수준)
한 point(i) 마다 십자가(j)와만 similarity를 계산하는 것이다. 여기서 자연스럽게 의문이 들게 되는데 기존의 방식은 전체 point에대한 정보를 다 수합할 수 있었는데, 이 경우에는 십자가를 제외한 부분의 정보는 반영되지 못하게된다. 본 논문에서는 이를 간단하게 모듈을 2번 이상 연속하게 쌓는 것으로 해결했다. 이 연속한 모듈들은 weight를 공유한다.
이와 같이 2번 연속으로 쌓은 Recurrent Criss Cross Attention(RCCA)를 구성하게 되면 십자가에 위치하지 않은 point들에 대한 정보도 반영할 수 있게된다. 그 이유는 위의 그림을 보면 알 수 있다. 첫번째 module에서 파랑색 point의 정보가 연두색 points로 흘러들어가게 되고, 두번째 module에서 연두색 points의 정보가 녹색 point에 흘러들어가기 때문에 결과적으로 파랑색 point의 정보가 녹색 point에 존재하게된다. 마찬가지로 임의의 다른 point들의 정보들도 다 반영이 되기때문에 global한 정보를 반영한 feature map이되게 된다.
Category Consistent Loss - CCL
이 파트는 해당 논문이 2020년 TPAMI에 기존 논문을 revision하여 발표할 때 추가된 부분이다.
RCCA의 경우 동일한 weight를 recurrent하게 사용하기 때문에 over-smoothing 문제가 발생할 수 있다고 주장한다. (하지만 주장일 뿐 실제로 실험 결과는 없었다는 점은 짚고 넘어가자.) RCCA를 제대로 학습 시키기 위해 직접적인 loss를 추가해 주었다. RCCA를 거쳐 나온 feature map에서 branch를 따와 Category Consistent Loss를 추가해 주었다.
아이디어는 같은 label(class)을 가지는 feature들은 비슷할 것이고(거리가 가까울 것이고), 다른 label을 가지는 feature들은 거리가 멀 것이라는 것이다. 아래와 같이 같은 class들 끼리는 모아주고, 다른 class들끼리는 떨어뜨려주는 loss를 구상하였다.
class가 같은 feature들에는 현재 학습 중인 mini-batch의 속하는 같은 class의 feature들의 mean값과의 거리에 따라 좌측 loss를 주게되고, class가 다른 feature들은 각 class의 feature들의 mean 값과의 거리에 따라 우측 loss를 주게된다. 그래서 최종 loss는 segmentation loss와 합쳐서 아래와 같게 구성했다.
Experiments
Sematic segmentation 뿐만 아니라 Instance segmentation도 진행하였다.
attention을 visualize한 것을 보면 R=1 일때는 십자가 부분만 attention 된것을 볼 수 있고, R=2일때는 전체적으로 attention 된것을 볼 수 있는데, query point가 차인 경우에는 멀리 떨어진 다른 차에 attention 되었고, 나무인 경우에 나무에 된 것을 확인 할 수 있다.
COCO dataset에서 instance segmentation을 보면 NL(Non-local neural network) 보다 성능이 좋은 것을 확인 할 수 있다.
Ablation Study
RCCA
CCNet vs NL
Category Consistent Loss
결론
criss-cross 형태를 사용하여 cost를 줄인 아이디어는 괜찮았고, 심지어 accuracy도 높아졌다는 게 놀라웠다. 하지만 여전히 cost efficiency하지 않다는 점은 생각해야한다. ablation study에서도 그러한 단점 때문에 일부러 baseline을 0으로 잡고 table을 작성한 것 같다.
category consistent loss는 accuracy가 높아진 사실에는 놀라웠지만, 본인들이 세운 가정( over-smoothing )에 대한 experiment나 feature들의 visualization을 해줬으면 좋지 않았을까 싶다. revision하여 논문을 재제출하다 보니 힘이 많이 빠진게 느껴졌다.
이미지 출처:
- Wang, Xiaolong, et al. "Non-local neural networks." Proceedings of the IEEE conference on computer vision and pattern recognition. 2018.
- Huang, Zilong, et al. "Ccnet: Criss-cross attention for semantic segmentation." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2019.
- Van der Maaten, Laurens, and Geoffrey Hinton. "Visualizing data using t-SNE." Journal of machine learning research 9.11 (2008).
