[ICML 2024] VISION TRANSFORMERS NEED REGISTERS

이번 논문은 Vision Transformers 기반 large model의 특징 표현력을 향상시키기 위해 원인을 찾아 분석하고 이에 대한 해결책을 제시한 논문입니다. 해당 기법에 주목하게 된 계기는 Vision Transformers 기반 large model의 일반화 능력 향상이 열화상에서의 성능 향상으로 이끌 것으로 판단되어 들고 오게 되었습니다. 분석 결과에 따라서도 어느정도 특징 표현력이 향상된 것으로 보아 기대해도 좋을 것 같습니다.

+ 결론적으로 해당 논문에서는 시사하는 바는 ‘Large ViT에 추가 토큰을 추가하는 것만으로도 표현이 강해진다’임다.

Intro

컴퓨터 비전 분야에서 다양한 목적으로 사용 가능한 일반화된 영상 특징을 추출하는 것은 오랜 숙원으로 치부되었습니다. 이러한 시도들은 handcrafted feature (e.g SIFT)에서 시작되어, 근래에는 지도 학습 기반의 딥러닝을 활용한 연구들이 진행되어 놀라운 결과물들을 보여줬습니다. 허나, 지도 학습 기반의 딥러닝 기법들은 가치 있는 데이터로부터 주석 처리된 정제된 데이터에 의존하기 때문에 특정 분야에서는 활용이 불가능하다는 한계를 가지고 있습니다. 최근, 라벨이 없는 대용량 데이터 셋을 자기 지도 학습을 통해 일반화를 보인 DINO와 비전과 언어 간의 일치 데이터 셋을 통해 의미론적인 표현 능력을 학습한 CLIP이 등장하면서 다양한 다운스트림 태스크에서 일반화된 성능을 보여주고 있습니다. 특히 DINO의 특징은 영상 내 의미론적인 영역을 명시적으로 표현하는 결과를 보여줍니다. 실제로 DINO를 활용한 LOST라는 알고리즘은 지도 학습 없이 물체를 찾는 object discovery 태스크에서 놀라운 성능을 보여주며, 컴퓨터 비전 분야에 새로운 지평을 열어주었습니다.

DINO의 후속 연구인 DINOv2는 해당 모델의 가중치를 얼려두고, head만 추가 학습하여 dense prediction tasks (semantic segmentation, depth estimation…)에서 SOTA에 달성하는 놀라운 성과를 보여주었습니다. 이러한 강력한 결과에도 불구하고, DINOv2를 LOST에 결합해 object discovery 태스크에 적용했을 때, 상대적으로 낮은 결과를 보여줍니다. 이러한 결과는 DINO와 DINOv2가 다름을 보여줍니다. 저자는 실험을 통해 DINO와는 다르게 DINOv2의 특징 맵에 인공물이 존재한다는 것을 확인했습니다. 놀랍게도 fig 2에서 보이는 바와 같이 지도 학습을 수행한 vision transformer (DeiT, OpenCLIP)에서도 인공물을 관찰 할 수 있었습니다.

저자는 이러한 현상을 이해하고 인공물을 검출하는 방법에 대해서 연구를 수행하였고, 이러한 현상들은 일반적인 토큰 대비 10배 높은 norm을 가지는 것을 관찰 할 수 있었습니다. 또한 이러한 토큰들은 vision transformer의 중간 레이어에서부터 등장 하고, 충분한 학습 시간과 충분히 큰 트랜스포머에서 등장한다는 것을 확인 하였습니다. 특히 이러한 outlier 토큰은 이웃 토큰과 유사한 패치, 즉 추가 정보를 거의 전달하지 않는 패치에서 나타나는 것으로 나타났습니다.

저자는 간단한 linear model로 outlier 토큰을 평가하여 토큰에 포함된 정보를 파악합니다. outlier 토큰은 non-outlier 토큰에 비해 이미지의 원래 위치나 패치의 원래 픽셀에 대한 정보를 덜 보유하고 있는 것으로 분석되었습니다. 해당 분석으로 통해 추론하는 동안 패치에 포함된 로컬 정보를 버린다고 볼 수 있습니다. 반면에, 영상 분류를 수행했을 때는 outlier 토큰이 non-outlier 토큰 보다 좋은 결과를 보여줍니다. 이를 통해 영상에 전역 정보를 가지고 있는 것을 확인 할 수 있습니다.

이러한 분석에 따라 저자는 다음과 같은 해석을 제시합니다. 모델은 유용한 정보가 거의 없는 패치를 인식하고 해당 토큰을 재활용하여 글로벌 이미지 정보를 집계하는 동시에 공간 정보를 버리는 방법을 학습을 한다!

이러한 해석은 제한된 토큰 수 내에서 결과물을 뱉는 트랜스포머의 내부 메커니즘과 맥락이 일치합니다. 저자는 가설을 증명하기 위해서 입력 영상에 독립적인 추가 토큰인 registers token을 입력 토큰에 추가합니다. registers token 적용하거나 적용하지 않은 여러 모델을 훈련하고 정성적인 결과로부터 outlier 토큰이 시퀀스에서 완전히 사라지는 결과를 제시합니다. 해당 결과물은 dense prediction task에서 모델의 성능이 향상되고 특징 맵의 표현이 훨씬 더 부드러워진 결과물을 보입니다.

+ 해당 논문에서 관심을 가지는 특징: outlier = 인공물 = 아티팩트 = high-norm

PROBLEM FORMULATION

ARTIFACTS IN THE LOCAL FEATURES OF DINOV2

해당 섹션에서는 인공물이 나타나는 이유와 시기에 대한 연구 내용을 제시합니다. 모든 비전 트랜스포머의 아티팩트를 완화하는 데 초점을 맞추지만, 여기서는 DINOv2에 대한 분석에 초점을 맞춥니다.

Artifacts are high-norm outlier tokens. 우선 트랜스포머 내의 local features로부터 나타나는 인공물의 정략적인 특성을 찾는 법을 보입니다. 저자는 ‘아티팩트’ 패치와 이와 다른 패키 간의 가장 큰 차이점은 임베딩 토큰의 norm에 있다는 것을 발견합니다. Fig 3의 왼쪽에서 DINO와 DINOv2의 local feature의 norm을 확인 할 수 있으며 명확하게 DINOv2에서 더 높은 norm을 가진 것을 볼 수 있습니다. 또한, Fig 3의 오른쪽에서 작은 데이터 셋 내에서의 local feature의 norm에 대한 분포를 볼 수 있습니다. 여기서 DINOv2에서 이분화된 분포를 통해 인공물을 볼 수 있습니다. 저자는 이러한 분석을 토대로 norm이 150 이상인 토큰들을 high-norm=outlier 토큰이라고 명명합니다.
+ 추가로 저자는 hand-picked cutoff value는 DINOv2에 해당하며 다른 모델들은 최적화된 값들이 있을 것이라고 합니다.

Outliers appear during the training of large models. 저자는 DINOv2를 학습 과정 중에 outlier token이 어떤 조건에서 나타나는지를 분석합니다. 분석 결과는 Fig 4에서 확인 할 수 있습니다. (a) high-norm patch는 40개의 레이어를 가진 ViT에서 15번째 레이어 주변에서부터 등장, (b) 훈련 과정의 1/3에서부터 분포가 관찰되며, (c) 다양한 ViT 중 Large 모델부터 분포 특징이 나타나는 것을 보여줍니다. 이러한 결과는 outlier token은 중간 레이어(15 layers)에서 등장하며, 어느정도 학습(1/3)이 진행되었으며, 충분한 크기(Large)의 ViT에서 등장한다는 것을 보여줍니다.

High-norm tokens appear where patch information is redundant. 주변에 중복된 패치 정보들이 존재했을 때, 등장한다는 것을 증명하기 위해서 저자는 high-norm 토큰과 4개의 이웃 패치로부터 cosine similarity 측정합니다. 해당 결과는 fig 5-(a)에서 확인 가능합니다. 해당 결과는 high-norm 토큰이 중복 정보를 억제하고 모델의 영상 표현의 품질을 손상시키지 않고 정보를 버리도록 학습이 됨을 시사합니다. 이에 대한 정성적인 관찰은 fig 2에서 이들이 uniform, background 영역에서 나타나는 것으로부터 확인 가능합니다.

High-norm tokens hold little local information. 해당 토큰의 특성을 더 이해하기 위해서 저자는 두 가지 태스크 (position prediction and pixel reconstruction)를 적용합니다. Position prediction: 영상 내 각 패치의 position을 예측하는 간단한 linear model을 학습 시키고 정확도를 평가합니다. position은 첫 레이어의 position embedding을로부터 획득합니다. Fig 5-(b)로부터 normal 토큰들이 더 좋은 결과를 보여줍니다. 이를 통해 outlier 토큰들이 local information이 소실된 것을 시사합니다. pixel reconstruction: 패치로부터 pixel value를 예측하는 linear model을 학습하고 정확도를 측정합니다. 해당 결과에서도 outlier 토큰이 더 낮은 결과를 보여줍니다. 두 결과를 통해 outlier 토큰이 local information이 덜 존재함을 확인 할 수있습니다.

Artifacts hold global information. high-norm 토큰이 얼마나 global information을 가지고 있는지를 평가하기 위해서 저자는 standard image representation learning benchmark로부터 평가를 진행합니다. 해당 결과는 tab 1에서 확인 가능합니다. [CLS]를 기존 classfication에서 활용되는 토큰으로 upper로 보시면 됩니다. outliner와 normal 토큰들은 각 요소 평가에 따라 랜덤하게 선정되어 classfication에 활용됩니다. 실험 결과, normal token을 활용한 경우에는 저조한 성능을 보여주는 반면에 outliner 토큰을 활용한 경우에는 [CLS]에 준하거나 뛰어넘는 결과를 보여줍니다. 해당 분석을 통해 outliner 토큰에 global information이 존재함을 시사합니다.

HYPOTHESIS AND REMEDIATION

위 와 같은 분석을 통해 저자는 다음과 같은 가설을 세웁니다. 이는 large, sufficiently trained ViT 모델은 redundant 토큰을 인식하고 이를 global information를 store, process 및 retrieve하는 장소로 사용하는 방법을 학습한다는 것입니다. 이러한 현상 자체는 나쁜 현상으로만 보긴 어렵지만, 영상 특징 맵의 표현력을 헤친다는 점에서 바람직하지 못한다고 봅니다. 실제로 tab 5-(b)에서 보인 바와 같이 local information을 버리도록 유도하기 때문에 dense prediction task에서는 성능이 저하될 수 있습니다.

저자는 해당 문제를 해결하기 위해 아주 간단한 수정을 제안합니다. 토큰 뭉치에 새로운 토큰을 registers로써 모델이 학습하도록 추가합니다. 이는 학습 가능한 토큰으로 patch embedding layer 추가되며, [CLS]와 유사하게 구성됩니다. ViT의 최종 결과에서는 해당 토큰을 폐기합니다.

이러한 관찰을 통해 다음과 같은 가설을 세웠습니다. 이는 large, sufficiently trained 모델은 redundant토큰을 인식하고 이를 global information를 store, process 및 retrieve하는 장소로 사용하는 방법을 학습한다는 것입니다. 또한, 이러한 행동 자체는 나쁘지 않지만 패치 토큰 내부에서 이러한 행동이 발생한다는 사실은 바람직하지 않다고 가정합니다. 실제로 이 동작은 모델이 로컬 패치 정보를 폐기하도록 유도하여(표 5b) 밀도가 높은 예측 작업에서 성능이 저하될 수 있습니다.

EXPERIMENTS

해당 섹션에서는 제안한 솔루션, [reg] register tokens을 추가한 ViT을 정량/정성적 분석을 통해 검증하고자 합니다.

TRAINING ALGORITHMS AND DATA
DEIT-III: 지도 학습 기반의 ViT. (ViT-B setting, ImageNet-22k)
OpenCLIP: text-image aligned models (지도 학습, ViT-B/16)
DINOv2: 자기 지도 학습 기반 ViT (ViT-L, ImageNet-22k)

EVALUATION OF THE PROPOSED SOLUTION

Fig 1에서 보이는 바와 같이 register 토큰이 추가해 학습된 모델에서 outliner 토큰이 완화된 결과를 보여줍니다. 정량적인 측정을 수행하기 위해서 각각의 모델로부터의 norm of features를 측정합니다. 해당 결과는 fig 7에서 확인 가능하며, 세 가지 모델에서 모두 norm의 분포도가 줄어든 결과를 보여줍니다.

Performance regression. 해당 실험에서는 register token이 representation 품질을 저하시키지 않는다는 것을 보입니다. 여기서는 linear probing으로 ImageNet (classification), ADE20k (segmenation), NYUd (monocular depth estimation)을 학습하고 평가를 수행합니다. 해당 실험은 tab 2-(a)에 해당합니다. 실험 결과, 제안한 기법을 적용할 경우에 기존 성능에 준하거나 뛰어넘는 결과를 보여줍니다. 추가로 zero-shot 성능을 확인하기 위해 추가 실험 (fig-(b))를 진행했으며, 개선된 결과릅 보여줍니다.

Number of register tokens. 해당 실험에서는 register token의 수가 local feature와 dense downstream task에 영향을 주는 가를 진행합니다. 이를 위해 register token이 0, 1, 2, 4, 8, 16로 추가된 DINOv2 ViT-L/14을 개별적으로 학습하여 실험을 진행합니다. 정성적인 결과는 fig 8-top에서 확인 가능합니다. register를 한개라도 추가하는 경우에 인공물이 완화된 결과를 보여줍니다. 정량적인 결과는 fig 8-bottom에서 확인 가능하며, 각 태스크마다 최적의 토큰 수가 존재하는 것으로 보입니다. 하지만 토큰 한 개라도 추가하는 경우, 성능이 향상된 결과를 보여줍니다. 저자가 진행하는 모든 실험은 4 register tokens을 유지합니다.

OBJECT DISCOVERY

앞서 언급한 LOST를 적용했을 때, 개선된 결과를 보여줍니다. 이를 통해 특징 표현력이 개선됨을 보여줍니다. 추가로 저자가 밝히길, DINO를 이용한 LOST 보다는 성능이 낮다고 합니다. 이는 후속 연구로 남겨둔다고 합니다.
+ 또한 OpenCLIP에서는 오히려 성능이 하락되는 결과를 보여줍니다. 해당 이유도 후속 연구로 남겨둔다고 하네요.

QUALITATIVE EVALUATION OF REGISTERS. 해당 실험에서는 register tokens의 수에 따른 정성적인 분석을 수행합니다. Fig 9에서 확인 가능하며, 놀랍게도 [CLS]와 유사한 결과를 보여줍니다. 또한 의도하지 않았지만, register token에 따라 서로 다른 물체가 attention이 형성되는 결과를 보여줍니다. 해당 현상에 대해서도 후속 연구로 남겨 둡니다.

+ 마지막으로 다양한 영상 도메인에서의 정성적 결과를 보시라고 자료 남깁니다.

---

제가 해당 논문을 리뷰한 이유는 열화상을 이용한 6D Pose Estimation에 foundation model을 적용하되, 개선을 어떻게 시킬까 하다가 찾은 논문입니다. 해당 논문의 분석 중 Fig 19-21에서 real-world의 object가 아닌 결과들과 color가 아닌 gray-scale 영상에서도 강화된 표현 능력을 보고 개선된 결과를 주지 않을까란 기대에 리뷰를 작성했습니다. 이후 연구에서 해당 기법과 또 다른 기법을 활용해서 열화상에서의 foudation model을 구현해보고자 합니다.