[논문리뷰] BootsTAP: Bootstrapped Training for Tracking-Any-Point

⚠️ 온전한 이해를 위해서 TAPIR 논문 리뷰 선행 권장

---

Paper link : https://arxiv.org/abs/2402.00847

BootsTAP: Bootstrapped Training for Tracking-Any-Point

---

➡️ Abstract

• Tracking-Any-Point(이하 TAP) : 비디오 데이터에서 조밀한 시공간 내의 특정 지점을 추적하는 알고리즘
• 해당 논문에서는 실제 large-scale, unlabeled, uncurated 데이터를 사용한 self-supervised learning을 통해서, 약간의 모델 구조 변화로도 TAP 모델의 성능 향상이 가능하다는 점을 시사함
• TAP-Vid benchmark의 기존 SOTA 모델 성능을 크게 뛰어넘었음
  • TAP-Vid-DAVIS : 61.3% → 66.4%
  • TAP-Vid-Kinetics : 57.2% → 61.5%

---

➡️ Introduction

• 최근 vision 모델의 발전이 지속되고 있으나, 여전히 물리 공간적인 추론 능력 부족이 SOTA 모델들의 약점
• 복잡한 움직임과 물리적인 상호작용을 이해해야 하는 robotics, video generation/editing, 3D asset creation와 같은 작업에서 TAP은 정확한 움직임을 표현하기 위한 유망한 알고리즘
  • 고체 표면에 대해서 일치하는 점을 추적하는 방식을 통해, 물체의 deformation과 motion에 대한 가치 있는 정보를 얻을 수 있음
• TAP은 비디오 데이터와 움직임 추적을 위한 query point 집합을 입력으로 받고, 비디오의 매 frame 마다 해당 query point들의 추적된 위치를 출력해야 함
  • 만약 중간에 query point가 사라지는 경우, 해당 frame에서 point가 투명해진(occluded) 정도도 함께 모델링함
• TAP 모델을 만들기 어려운 이유는 현실적인 training data의 부족
  • 이유 1: 부정확하고 균일하지 못한 수동 labeling 방식
  • 이유 2: 특수한 상황에서만 사용할 수 있는 3D sensing 방식
  • 따라서, 현재 SOTA 모델은 임의 생성(synthetic) 데이터에 의존
• 해당 논문은 unlabeled 데이터를 사용해서 TAP 성능을 향상하는 방식을 제안
  • 만약 주어진 video에 대한 추적이 잘 이루어졌다면 아래 3가지 항목들이 지켜져야 함 (지켜지지 않아서 생기는 error를 학습을 위해 사용)
    1. video와 trajectory에 대한 spatial transformation이 동일
    2. 동일한 trajectory 위에 존재하는 query point들은 동일한 추적값 산출
    3. 이미지 압축과 같이 비공간적인 데이터 augmentation을 한 경우, 동일한 결과 산출

 

---

➡️ Method

self-training TAP을 개발함에 있어서 저자는 optical flow와 비슷하게 TAP도 query point들에 대한 정확한 정답이 존재한다는 점에 주목했다. 이는 의미적으로 비슷한 이미지가 비슷한 표현을 갖는 한, 다른 self-supervised 방법론과는 구분되는 특징이다. 

 

제시한 방법론은 단단하고 불투명한 표면 위의 point 추적에 대해서 아래 두 가지 조건에 기반한다.

1. affine tranformation과 같은 spatial transformation를 video에 적용한 결과와 point tracks(trajectory)에 적용한 결과는 동일
2. track 위의 어떤 point를 query로 사용하더라도 산출되는 결과는 동일

저자는 위의 두 조건을 만족시키기 위해 Siamess 구조를 사용할 수 도 있지만, 이러한 방법은 예측 결과 성능이 좋지 않았다고 한다. 대신 Student-Teacher 구조를 사용했는데, student는 augmentation을 통해 더 어려운 데이터를 입력으로 받고 teacher은 gradient를 전달받지 않는 것이 특징이다.

 

모델은 Kubric 데이터 셋으로 pre-trained된 TAPIR network로 시작된다. 모델의 예측값을 \(\hat{y}=\{\hat{p}[t], \hat{o}[t] , \hat{u}[t]\}\)라 하자. 이때 \(\hat{p}\in\mathbb{R}^{T\times2}\)는 위치값(position), \(\hat{o}\in\mathbb{R}^{T}\)는 투명도(occlusion logit), \(\hat{u}\in\mathbb{R}^{T}\)는 불확실도(uncentainty logit)이고 \({T}\)는 frame수 이다. \({p[t]}\)와 \({o[t]}\)는 frame \({t}\)에서의 ground truth이고, 하나의 trajectory에 대한 TAPIR loss는 아래와 같이 정의된다.

 

$ \begin{array}{rlr}{{\mathcal{L}_{tapir}}(\hat{p}[t],\hat{o}[t],\hat{u}[t])} & {=\operatorname{Huber}({\hat{p}}[t],p[t])(1-\hat{o}[t])} & \text{Position loss}\\ {} & {\quad\operatorname{+BCE}({\hat{o}[t], o[t]})} & \text{Occlusion loss}\\ {} & {{\quad\operatorname{+BCE}({\hat{u}}[t],u[t])(1-{\hat{o}}[t])}} & \text{Uncertainty loss}\end{array}  \tag{1}$

 

불확실도(uncerntainty logit)는 \({u[t]}= \unicode{x1D7D9} (d(p[t],\hat{p}[t] > \delta))\)으로, \({d}\)는 \({L}_2\) 거리이며 \({\delta}\)는 6 pixel로 설정된 거리의 threshold 값이고 \({ \unicode{x1D7D9} }\)는 특성 함수(indicator function)이다. 즉, 불확실도는 모델의 예측값이 ground truth의 근처(threshold 만큼)에 위치하게끔 학습하게 하는 효과를 준다. 

 

pre-training 이후에는 추가 convolutional residual layer를 추가해서 추가 학습 데이터를 입력받을 수 있게 한다. \(\hat{y}_\mathcal{S}=\{\hat{p}_ \mathcal{S} , \hat{o}_ \mathcal{S} , \hat{u}_ \mathcal{S} \}\)를 student의 예측값이라고 하고, teacher의 예측 \(\hat{y}_\mathcal{T}=\{\hat{p}_ \mathcal{T} , \hat{o}_ \mathcal{T} , \hat{u}_ \mathcal{T} \}\)으로 부터 pseudo-labels \({y}_\mathcal{T}=\{{p}_ \mathcal{T} , {o}_ \mathcal{T} , {u}_ \mathcal{T} \}\)를 뽑는다. ( \({t}\)는 시간 index )

 

$ p_\mathcal{T}[t]=\hat{p}_\mathcal{T}[t]\quad;\quad o_\mathcal{T}[t]=\unicode{x1D7D9}(\hat{o}_\mathcal{T}[t]>0);\quad u_\mathcal{T}[t]= \unicode{x1D7D9}(d(p_\mathcal{T}[t],\hat{p}_\mathcal{S}[t])>\delta) \tag{2} $

 

video frame \({t}\)에서의 loss값 \({\ell_{s s l}(\hat{p}_\mathcal{S}[t], \hat{o}_\mathcal{S}[t], \hat{u}_\mathcal{S}[t])}\) 는 TAPIR loss로부터 구할 수 있고, 이는 pseudo-label로써 ground truth로 사용하며 다음과 같이 정의된다.

 

$  \begin{array}{rl}{{\ell_{tapir}}(\hat{p}_\mathcal{S}[t],\hat{o}_\mathcal{S}[t],\hat{u}_\mathcal{S}[t])} & {=\operatorname{Huber}({\hat{p}}_\mathcal{S}[t],p_\mathcal{T}[t])(1-{o}_\mathcal{T}[t])} \\ {} & {\quad\operatorname{+BCE}({\hat{o}_\mathcal{S}[t], o_\mathcal{T}[t]})} \\ {} & {{\quad\operatorname{+BCE}({\hat{u}}_\mathcal{S}[t],u_\mathcal{T}[t])(1-{{o}}_\mathcal{T}[t])}} \end{array}  \tag{3} $

 

TAPIR의 loss는 여러 번 refinement iteration을 거치게 되고, teacher의 마지막 단계의 refine 된 예측값을 ground  truth로 사용하기 때문에 refine 되지 않은 데이터가 더 빠르게 수렴할 수 있게 돕는 역할을 할 수 있다.

 

만약 위의 식이 잘 정의되었다면 student와 teacher가 동일한 video query point를 입력받았을 때, loss 값은 아주 작아야 한다. 이러한 특성을 학습하게 하기 위해서 input video query에 대해서 affine tranformation과 같은 \({\Phi}\)를 pixel에 적용시켜서 낮은 resolution을 갖도록 사이즈를 줄이고, 예측을 더욱 어렵게 하기 위해 JPEG degradation을 적용한 이미지를 student에 전달하여 track을 연산하게 한다. 출력 값에 \({\Phi}^{-1}\)를 적용시켜 입력 좌표공간으로 되돌린 다음 teacher의 track과의 비교를 통해 loss를 계산한다.

 

 

또한 저자는 동일한 같은 track 위의 query points 짝에 대해서 teacher의 두 예측값이 같은 값을 출력하도록 하도록 frame별 loss 함수를 수정하는 두 가지의 계수를 추가했다. 첫 번째로, 먼저 \(Q_1 = (q_1, t_1)\)을 뽑았는데 \(q_{1}\)은 \({(x, y)}\)형식의 좌표이고 \({t_1}\)은 frame index이며 두 값은 랜덤하게 sampling한다. 그 후에 student의 query는 teacher의 trajectory 위에서 랜덤하게 뽑는다. (\({i.e.}\quad Q_2=(q_2,t_2)\in\{(p_\mathcal{T}[t],t);t \quad s.t. \quad o_\mathcal{T}[t]=0\}\))

 

하지만 teacher가 query point를 잘 추적(track) 하지 못하는 경우가 있을 수 있는데, 그럴 경우 student는 더더욱 추적을 잘하지 못할 것이다. 이러한 경우를 막기 위해서 저자는 cycle-consistency를 도입했다. valid 한 trajectory를 구분하는 mask를 계산해서 valid 하지 않은 경우는 loss 계산을 하지 않고 무시하는 방법인데, 이때 사용되는 mask는 다음과 같이 정의된다. ( \(\delta{cycle}\)은 거리 threshold이며, 저자는 4 pixel을 사용했다, 위 Fig. 2의 주황 원형 영역 참고 )

 

$ m_{c y c l e}= \unicode{x1D7D9} \left(d({\hat{p}}s[t_{1}],q_{1}\right)<\delta_{c y c l e})\quad*\quad \unicode{x1D7D9} \left(\delta s[t_{1}]\leq0\right) \tag{4} $

 

두 번째로, frame index \({t}\)의 위치에 따른 정확도 차이를 고려하는 proximity를 도입했다. 저자는 student의 query frame에 가까운 포인트들에 대해서는 teacher의 예측값보다 student의 예측값이 더 정확할 것이라고 생각했다. 이를 반영해서 정의한 mask는 다음과 같다. ( Fig. 2의 회색 사각 영역 참고 ) 

 

$ m^t_{proximity}= \unicode{x1D7D9}(|t-t_1| \leq |t-t_2|) \tag{5} $

 

위 \({(4)}\)와 \({(5)}\)를 통해 계산한 최종 loss는 아래와 같다.

 

$ \mathcal{L}_{S S L} = \sum\limits_{t}{m}_{\scriptscriptstyle{c y c l e}}^{t} \ast m_{p r o x i m i t y}^{t}\ast\ell_{\scriptscriptstyle{S S l}}^{t} \tag{6} $