해당 논문은 NeurIPS 2015에서 소개된 논문이다.
당시 CNN은 fast, scalable, end-to-end learning이 가능하다 라는 점에서 다양한 분야를 위한 네트워크로 각광받고 있었다. 이러한 배경에서 논문의 저자들은 CNN의 한계점에 주목하며 문제를 해결하고자 한다.
문제 정의
- CNN은 spatially invariance의 특성을 가지지 못한다는 한계점을 가진다. 즉 원본 이미지에 대해서 약간의 변형에 가해졌을 때 같은 결과를 출력하는 일관성이 떨어진다는 뜻이다.
- 이를 해결하기 위한 방법이 max pooling이었다. Max pooling의 사용이 주는 이점이 spatially invariance라는 점과 결과에서의 향상을 살펴보면 분명 의미가 있는 기법이다. 하지만 max pooling 또한 한계점이 존재한다. 먼저 이것을 위해 deep hierachy max pooling을 통해서만 invariance를 얻을 수 있다는 점이다. 다음으로는 max pooling이 넓은 영역에 대해 적용되면 information loss가 발생한다는 점에서 불가피하게 작은 영역만을 반영한다. 따라서 large transformations에 대해서는 invariance를 유지할 수 없다는 한계점을 가진다.
해결 방법 및 원리
- Spatial Transformer Networks(STN)
이것을 해결하기 위해 저자가 고안한 방법이 Spatial Transformer Networks이다. 논문의 저자가 생각한 방법은 조금씩 변형되어 알아보기 어려워진 이미지에서 인식해야하는 대상이 있는 영역을 뽑아내는 localisation network를 추가하여 여기서 학습한 파라미터를 기반으로, 원래 이미지를 알아보기 쉽게 바꾼 새로운 이미지로 학습하는 방법이다. 생각하기에는 쉬운 방법이지만 이것을 직접 구현하는건 어렵게 느껴진다. 먼저 아래는 고안한 네트워크의 모습이다.

U는 input feature map, V는 최종 output map을 의미한다. 네트워크의 단계를 살펴보면, input U가 들어가면 이 이미지의 feature map을 localisation network에 투입한다. 해당 네트워크는 파라미터 세타를 학습한다. 이 세타는 smpling grid를 만들게 되는데, 이를 통해 어떤 지점을 잡아야 원본 이미지를 알아보기 쉬운 상태로 만들 수 있는지 학습한다. 최종적으로 input image feature map과 sampling grid가 sampler의 input으로 들어가고 최종 변형된 output map이 나오게 된다.
특징을 살펴보자면,
- 해당 모듈은 differentiable 하다는 점
- CNN 사이에 삽입하는 것이 가능하다는 점
- 각 이미지마다 feature를 기반으로 localisation network parameter를 각각 학습한다는 점
- 단순화를 위해 single transform으로 single output을 만든다는 점이 있다.
각 이미지에 대해서 fixed된 local 방식으로 수행하는 max pooling과는 다르게 모든 이미지 각각에 대해서 개별적으로 학습이 가능해진다. 이것이 STN이 전반적인 이미지의 transformation을 학습할 수 있게 만들어주는 중요한 특징이다.
이제 각 단계를 조금 더 자세히 살펴보자.
1. Localisation Network

먼저 위와 같은 크기의 input feature map을 받게 된다. 여기서 localisation network를 단순히 아래와 같은 함수로 표현할 수 있고, 해당 network를 통과하여 얻어지는 파라미터를 세타라고 정의한다.

처음에는 이 세타가 의미하는게 무엇인지 잘 몰랐는데, 다음 단계를 보니 이해가 쉬웠다. 이 파라미터 세타는 이후 transformation에 적용하는 파라미터이다. 해당 논문에서는 affine transformation을 언급한다. 이 변환은 translation, scaling, rotation, shear, reflection 등을 사용하여 변환을 수행하는 방법이다. 총 6개의 세타 파라미터를 필요로한다고 한다. 이 변환을 통해 output map과 input map의 sampling point를 적절히 match시키는 것이다.
Localisation network는 이런 세타를 구하는 목적이기 때문에 최종 layer에는 regression layer를 포함해야한다.
2. Parameterised Sampling Grid
이제는 알아보기 쉬운 단계로 input image를 warping 시키기 위해서 적절한 sampling point를 찾아야한다. 따라서 두 번째 단계에서는 먼저 output map에서 픽셀의 위치를 의미하는 regular grid G={Gi}를 먼저 정의한다. 여기서 Gi는 아래와 같이 표현이 가능하다.

즉 이제 transform (여기서는 affine transformation)을 위한 파라미터 세타와 output map에서 regular grid를 알고 있다. 그렇다면 아래와 같은 식으로 input map에서 sampling grid를 구한다.

이러한 transformation을 통하면 앞서 언급했듯이 crop, translation, rotation, scale, skew 등의 변환이 가능하다. 추가적으로 논문에서는 cropping이 가능한 이유를 언급한다. 예를 들어 affine transformation에서 왼쪽의 2x2 sub matrix가 unity matrix의 크기보다 작다면 기존 input feature map보다 작은 평행사변형이 sampling 되고 이것이 결국 cropping을 의미한다. 아래 사진에서 (b)와 같은 경우이다.
만약 파라미터 세타가 identity라면 (a)와 같이 input feature map이 그대로 output feature map으로 출력 될 것이다.

3. Differentiable Image Sampling

이제 output feature map을 구하기 위해서 필요한 값들을 모두 구했다. 이제는 input feature map에 sampling grid를 결합하여 최종 output feature map을 구하는 sampling 과정만 거치면 된다. 여기서 중요한 점은 여전히 differentiable한 방법으로 sampling을 진행해야, STN이 전체적으로 학습 가능한 네트워크가 된다는 점이다.
위의 식은 sampling 하는 식을 일반화하여 나타낸 것이다. Φx와 Φy는 generic sampling kernel k()의 parameter로 이해하면 된다. U^c_nm은 input feature map의 channel c에서 (n,m)위치에 존재하는 value를 의미한다. 구하는 값인 V^c_i는 output map의 channel c에서의 pixel i 값을 의미한다. Differentiable을 위해서 우리에게 중요한 것은 어떤 sampling kernel을 사용하는가 이다.
예를 들어integer sampling kernel을 사용한다고 가정하자.

즉 반올림을 통해 정수로 변형하는 sampling kernel을 통과시킨 것이다. 잘 생각해보면 이미 정수 값만 가지게 되었다는 점에서 비선형적인 처리가 된 것을 알 수 있고 미분은 불가능해지는 것을 생각할 수 있다. 비슷한 개념으로 아래와 같은 이미지 신호에서의 aliasing을 생각할 수 있을 것 같다. 선을 확대하면 비선형적인 연결처럼 계단식으로 표현되는데 비슷한 느낌이라고 생각한다.

따라서 여기서는 bilinear sampling을 적용한다.

따라서 미분이 가능해지며 역전파 시의 미분은 아래와 같이 구할 수 있다.

미분이 가능한 smapler까지 사용하였다면 이제 아까 보았던 network 구조의 그림을 생각해보자. Output feature map은 sampling grid와 input feature map과 연결되어 있었다. 따라서 양쪽 모두의 back propagation이 가능해진다. 이것은 localisation network의 학습을 가능하게 만든다.
- STN + CNN의 장점
1. Image classification: 이미지 분류에서 같은 class에 대해서 서로 다양하게 변형된 이미지 샘플을 학습해야하는 경우 CNN의 분류 성능을 높여줄 수 있다.
2. Co-localisation: 같은 class이지만 서로 다른 예시를 가지는 경우 해당 이미지들에 대해서 localisation이 가능해진다.
3. Spatial attention: Attention을 통해 저해상도에서의 transformation을 고해상도에 적용할 수 있어지면 계산 효율이 좋아진다는 장점이 있다. STN은 이러한 attention mechanism이 필요한 경우에 유연하게 사용 가능하다.
- STN의 특징
- 이 모듈은 computationally very fast하다, 따라서 training speed에 영향을 미치지 않는다
- CNN에 any point에 삽입되어도 무관하다
- CNN에 사용되면 STN은 training 과정에서 overall cost function을 minimise 하도록 feature map을 transform하는 것을 돕는다
- Downsampling, oversampling이 모두 가능하다 (output dimension이 input과 같지 않아도 되고, 사용자가 직접 정의할 수 있기 때문)
- 하지만 sampling kernel이 fix된다면, downsampling 과정에서 aliasing이 발생할 수 있다 (samll spatial support kernel의 경우)
- Multiple STN을 CNN에 적용하는 것이 가능해짐 (abstract representation의 증가를 유도)
- Multiple STN을 parallel하게 연결하는 것도 가능 (multiple object가 있는 경우 유리, but STN의 수가 network가 model하는 객체의 수를 제한함)
결과
1. Distorted MNIST

기존의 MNIST dataset에 대해서 다양한 왜곡을 진행한 데이터셋에서 실험한 결과이다. Baseline으로는 FCN, CNN을 사용하였다고 한다. 모든 네트워크는 SGD를 사용하여 최적화를 하고, 비슷한 파라미터 수, scheduled learning rate decrease를 사용하였다고 한다. 또한 multinomial cross entropy loss를 손실함수로 적용하였다.
CNN에 대해서는 2 max pooling layer를 추가하였다. 그리고 STN를 적용한 CNN, FCN의 결과와 비교하였다. STN을 위한 transformation에서는 affine transformation(Aff), projective transformation(Proj), thin plate spline transformation(TPS)을 사용하였다.
결과는 위의 사진과 같다. 먼저 RTS 변형에서 CNN과 ST-FCN의 Aff 버전을 살펴보면 0.8%로 같은 error rate를 보여준다. 이것으로 STN이 max pooling과 같이 invariance를 부여한다는 것을 알 수 있다. ST-CNN과 ST-FCN 중에서는 ST-CNN이 우수한 것을 알 수 있다. 이것은 CNN에 추가적으로 사용된 max pooling layer의 영향과 CNN의 우수함이 복합적으로 나타난 결과라고 생각된다. Transformation 측면에서는 TPS가 미세하게 우수한 결과를 내는 것을 확인할 수 있다.
추가적으로 60x60의 noisy한 환경에서 실험한 결과, FCN(13.2%), CNN(3.5%), ST-FCN(2.0%), ST-CNN(1.7%)의 순서로 error율을 구했다고 한다.
2. Street View House Numbers
해당 데이터에는 1~5개의 집이 각 이미지에 존재한다고 한다. 따라서 CNN의 경우는 11개의 hidden layer와 5-class softmax classifier를 사용하였다고 한다. 여기서는 ST-CNN만 사용하였는데 1개만 사용한 single과 여러개를 사용한 multi를 모두 시도하였다. Mutiple STN은 아래 사진과 같이 각 conv 앞에 1개씩 추가하였다고 한다.
결과를 살펴보면 기존 CNN이나 DRAM보다 ST-CNN single의 성능이 더욱 좋은 것을 알 수 있다. 또한 multiple STN의 경우 더욱 좋은 성능을 보여주는 것을 확인하였다. 이것은 여러 STN을 적용함에 따라서 첫 STN은 input feature map을 사용하지만, 이후 STN은 conv 이후 얻어지는 각 feature map에 적용되기 때문에 feature가 풍부해지는 효과를 얻었기 때문으로 해석된다. 따라서 deeper spatial transformation이 가능해지는 것이다.
속도 면에서는 ST-CNN multi가 기존 CNN 방법보다 6% 정도만의 속도 저하를 보였다고 한다.

3. Fine-Grained Classification

여기서는 CUB-200-2011 dataset을 사용했다고 한다. 여기서 적용한 ST-CNN multi의 경우는 parallel 모델을 적용하였다고 한다. Baseline model로는 inception architecture와 batch normalisation을 적용한 모델을 사용했다고 한다. 이것은 ImageNet에서 pretrain 하였다.
최종적으로 4개의 parallel STN을 사용하여 baseline CNN 모델을 1.8% 능가하여 84.1%의 성능을 낸 것을 확인할 수 있다. 하지만 ImageNet과 CUB test set에 overlap이 존재한다는 점 때문에, 이 overlap을 제거하고 실험하니 84.0%로 1.7% 능가함을 확인하였다.
또한 input image로 224px 이외에도 더 큰 448px 이미지를 사용하더라도 결과에 영향이 별로 없다는 점이 특징이다. 왜냐하면 STN을 통해 downsample이 가능하기 때문이라고 설명한다. 조금 신기한 점은 오히려 448px이 input인 경우 더욱 높은 성능을 낸다는 점인데, STN을 downsample로 사용하는 것이 유리한 것인가 하는 생각이 들었다. 또한 STN으로는 upsampling도 가능한데 결과에 추가하지 않은 것을 보면 성능이 나빠진 것은 아닌지 궁금했다.
추가적으로 논문의 저자가 한 해석이 있다. 위의 사진에서 볼 수 있듯이 parallel STN을 적용하는 경우 각 STN이 뽑아내는 feature map에 특징이 생긴다는 것이다. 2개의 STN의 경우 빨간색은 새의 머리를, 초록색 부분은 새의 몸통을 뽑아내는 특징을 가지게 되었다. 즉 pose normalised representation을 나타내는 것이다. 이것은 추가적은 supervision을 적용하지 않았고 data에서 뽑아지는 방식으로 굉장히 신기했던 점이다.
결론
나는 기존에 데이터의 invariance를 위해 사용하는 pooling 혹은 사전에 데이터를 처리하는 data augmentation 등만 알고 있었다. 이것은 network에서 직접 각 data에 맞게 학습하여 transform하는 방법을 채택하였다는 점에 놀랐다. 이것이 추가적은 label을 필요로하지 않고 막대한 계산량을 사용하지도 않으면서, 마지막 새 사진에서와 같이 feature를 뽑아낼 수 있다는 것이 강력한 결과라고 생각한다. 이것이 object detection과 같은 dense level prediction에 응용될 수는 없을까 하는 궁금증이 생기기도 하는 논문이었다.