[5주차 / 박민규 / 논문리뷰] ViTPose : Simple Vision Transformer Baselines for Human Pose Estimation

ViTPose

지금까지 Pose Estimation 작업에 Vision Transformer를 적용하는 방법은 다음과 같았다.

1. CNN을 BackBone으로 사용하여 추출된 특정을 Transformer에 사용하는 방법 (TokenPose, TransPose)
2. Keypoint들의 위치를 잘 찾아내기 위해서 Encoder와 Decoder를 통합하는 방법 (PRTR)
3. Transformer로 직접적으로 특징을 추출하도록 하고, 복잡한 구조를 통해 성능을 높히는 방법 (HRFormer)

기존의 방법들은 pose estimation에서 좋은 성능을 내는 것은 맞다.

하지만 feature extraction을 위해 추가적인 CNN 네트워크가 필요하거나, task에 맞는 transformer structure를 조심스럽게 디자인해야한다는 특징이 있다.

ViTPose의 장점은 다음과 같다.

1. 모델 구조가 간단하다. simplicity
2. ViTPose의 구조는 아주 단순하고 특정 domain에 dependent하지 않은 plain한 encoder로 구성되어 있으며, decoder 또한 단순한 구조를 가지고 있다.
3. 모델 크기를 쉽게 조절하능하다. scalability
4. Transformer의 layer를 쌓으면서 model size를 더 키울 수 있고 feature dimension 또한 쉽게 늘리거나 줄일 수 있다. → ViT-B, ViT-L, ViT-H
5. 학습에서의 유연성을 가진다. flexibilitySingle pose dataset으로 training을 하더라도, 추가적인 decoder들을 추가함으로써 multiple pose dataset에 잘 적응할 수 있도록 변경할 수 있다.
6. 사전학습을 smaller unlabeled dataset으로 했거나, frozen된 attention module를 사용하여 fine-tuning을 하더라도 여전히 SOTA 성능을 가진다.
7. input resolution과 feature resolution이 달라도 model이 잘 적응하며, pose estimation task에서 변함없는 정확도를 보인다.
8. 학습 내용에 대해 전이가능성이 있다. transferability
9. small ViTPose 모델로 학습되더라도 large ViTPose model의 지식을 사용하여 transferring한다면, 쉽게 성능을 향상시킬 수 있다.

Simplicity

ViTPose의 모델 구조는 기존 ViT의 CLS 토큰과 Decoder를 제외하면 동일하다고 할 수 있다.

1. Encoder

ViTPose는 다른 복잡한 구조를 사용하지 않고, ViT와 간단한 Decoder를 활용하여 좋은 성능을 냈다. 구조적인 간단함을 가지기 위해 Skip connection과 cross attention 등의 방법을 사용하지 않았다.

MHSA는 Multi Head Self Attention을 뜻하며, FFN은 Feed Forward Network를 의미한다.

i는 i번째 transformer layer를 말하며, $F_0$는 patch embedding layer를 말하며, d는 patch의 크기를 의미한다.

입력이미지 크기는 $X\in\R^{H\times W\times 3}$이고, patch embedding layer의 output은 $F_0\in\R^{\frac{H}{d}\times\frac{H}{d}\times C}$ 이다.

Encoder의 최종 output은 $F_0\in\R^{\frac{H}{d}\times\frac{H}{d}\times C}$ 이다.

2. Decoder

Decoder는 각 관절에 대한 Heatmap을 표현하기 위해 2가지 방법을 제안한다.

 

1. Deconvolution layer를 사용 (Classic Decoder)

두 번의 Deconvolution를 거친 후에 BatchNorm과 ReLU를 차례대로 적용한다.

2. Simple Decoder
Bilinear interpolation을 4번하여 feature map을 직접적으로 upsamplng 한 후, ReLU와 3x3 convolution을 차례대로 적용한다.

 

이때 $K\in\R^{\frac{H}{4}\times\frac{H}{4}\times N_k}$ 이며, $N_k$는 관절의 개수이다. 이렇게 간단한 구조의 decoder로도 좋은 성능을 낼 수 있었다.

ResNet을 Backbone으로 한 것과 달리, ViTPose는 Simple Decoder($AR_{50}$)와 Classic Decoder($AR$)의 성능 차이가 별로 없다.

ViTPose는 feature extract를 매우 잘하며, 이로인해 복잡한 형태의 decoder를 필요로 하지 않는다는 결론을 도출할 수 있다.

Scalability

ViTPose 모델이 간단하기 때문에 여러 모델 크기를 선택할 수 있다. 다음 그림과 같이 모델 구조의 큰 변경 없이 layer 개수 또는 channel 개수만을 변경하여 쉽게 다른 형태의 ViTPose 모델을 구현할 수 있다.

Flexibility

Pre-training data flexibility

ViT는 특성상 많은 양의 데이터를 학습해야 좋은 성능을 낼 수 있다. ⇒ inductive Bias가 부족하기 때문

이러한 특성으로 인해 Masked AutoEncoder 방식을 활용하여 pre-training을 진행한다.

MAE; Patch의 일부를 masking하면서 encoder의 끝 부분에서 masking 된 부분을 추측하도록 하는 방식

Resolution flexibility

ViTPose가 higher resolution image을 input하게끔 하기 위해선

• 단순히 각 패치의 size를 유지한채 patch embedding layer의 stride를 변경하여, input image size를 resize한다.
• 그 후 resize된 image를 패치로 나누어서 일반적인 방법처럼 모델을 훈련시킬 수 있다.

Attention type flexibility

이미지 전체에 대해 attention을 수행하는 경우(patch size=1) computational cost가 크게 발생한다.

position embedding을 추가한 attention(ViT)를 이용하면 computational cost가 출어들긴 하지만, global context 정보가 줄어들기 때문에 성능의 저하가 발생할 수 있다.

• Shift Window (Swin Transformer에서 제안된 기법)
• 고정된 window attention 대신 shift 연산을 통해 window 정보가 흐를 수 있도록해, edge부분의 성능을 높히는 방법이다.
• Pooling Window (ViTDet에서 제안된 기법)
• Shift Windwo 방법과는 다르게 Pooling 방법을 통해 정보가 흐를 수 있도록 하는 방법이다.

Fine-tuning flexibility

Transformer 기반 모델이 그렇듯 ViTPose는 pre-trained 후 fine-tuning을 하면 더 좋은 성능을 얻을 수 있다.

Task flexibility

Decoder의 구조가 간단하기 때문에 pose estimation을 위해 구조를 바꿀 필요없이, 동일한 encoder를 가지고 여러 자세 추정 데이터에 맞게 heatmap을 구성할 수 있다.

Transferability

transfer learning의 방법 중 하나는 잘 학습된 큰 모델을 학습 해야할 작은 모델에게 모방하도록 학습하는 것인데 이를 distillation 방법이라고 한다.

잘 학습된 큰 모델의 output을 $K_t$, 학습 해야할 작은 모델의 output을 $K_s$라고 한다면 transfer learning을 위한 loss function은 다음과 같다.

본 논문에서는 output result만 동일하도록 transfer하는 것이 아닌 token-based distillation 방식을 사용하였다.

1. 학습가능한 extra knowledge token을 randomly initialize하고 난 후, teacher model($K_s$)의 patch embedding layer 이후에 visual token에 append한다.
2. 잘학습된 teacher model의 parameter들을 freezed(고정)하고 knowledge token을 학습시킨다.

$K_{gt}$는 GT heatmap, X는 input image, $T(t;X)$는 teacher model의 predictions, $t^*$는 loss를 최소화시키는 optimal token을 나타낸다.

 

     3. 이후 t*를 frozon한 후 student model($K_s$)내의 visual token과 함께 concat하여 transfer learning을 진행한다.

 

최종적인 Loss function은 다음과 같다.

$L^{td}{t->s},L^{tod}{t->s}$ 는 각각 token distillation loss, (output distillation + token distillation loss)의 조합을 의미한다.

---

참고 문헌

ViTPose 논문 리뷰