[논문리뷰] Vision Transformers Need Registers (ICLR 2024) - ViT attention map artifact 원인과 register token 정리

ViT(Vision Transformer)의 attention map을 뽑아보면 물체와 상관없는 배경 위치에 이상하게 밝게 튀는 점(artifact)이 생기는 경우가 있음. ICLR 2024에 발표된 Vision Transformers Need Registers는 이 artifact가 왜 생기는지 분석하고, register 토큰이라는 아주 단순한 방법으로 해결한 논문임. 논문에서 말한 내용만 정리함.

논문 정보

• 제목: Vision Transformers Need Registers
• 저자: Timothée Darcet, Maxime Oquab, Julien Mairal, Piotr Bojanowski (Meta AI Research, Inria)
• 학회: ICLR 2024 (International Conference on Learning Representations, oral 발표)
• 링크: https://arxiv.org/abs/2309.16588

문제의식: attention map에 왜 이상한 점이 생기는가

• DINO(2021)는 별도 supervision 없이도 CLS 토큰의 attention map이 물체 윤곽을 깔끔하게 보여주는 것으로 유명했음.
• 그런데 후속 모델인 DINOv2의 attention map에는 배경 영역에 튀는 점 형태의 artifact가 나타남. 이 현상은 DINOv2뿐 아니라 supervised로 학습한 DeiT-III, 텍스트-이미지로 학습한 OpenCLIP에서도 관찰됨 → 특정 학습 방식의 문제가 아니라 ViT 전반의 현상임.
• 논문은 이 artifact 토큰들의 특징을 다음과 같이 정리함.

artifact 토큰의 특징

• high-norm 토큰: 해당 위치의 패치 토큰 출력 벡터 norm이 다른 토큰보다 훨씬 큼(약 10배 수준). 전체 토큰 중 약 2% 정도가 여기에 해당함.
• 발생 위치: 이웃 패치와 거의 똑같은, 정보량이 적은 배경 패치에서 나타남. 물체 위에는 잘 생기지 않음.
• 발생 조건: 큰 모델(ViT-L, ViT-H, ViT-g)에서 나타나고 작은 모델(ViT-S, ViT-B)에서는 거의 없음. 학습이 충분히 진행된 뒤(전체 학습의 약 1/3 지점 이후)에 나타나고, 네트워크의 중간~후반 레이어에서 발생함.

이 토큰들은 무엇을 담고 있는가

• local 정보는 적음: 토큰 하나만 보고 (1) 원래 이미지의 어느 위치 패치였는지 맞히기, (2) 원래 픽셀값 복원하기 실험을 하면, high-norm 토큰은 일반 토큰보다 성능이 크게 떨어짐 → 자기 위치의 정보를 버렸음.
• global 정보는 많음: 토큰 하나만으로 ImageNet 이미지 분류를 linear probing 하면, high-norm 토큰이 일반 토큰보다 정확도가 훨씬 높음(DINOv2 ViT-g 기준 일반 토큰 약 65.8% vs high-norm 토큰 약 78.5%).
• → 논문의 해석: 모델이 정보가 없는 배경 패치 토큰을 "재활용"해서 전역(global) 정보를 계산·저장하는 임시 공간으로 쓰고 있음. 문제는 그 토큰이 원래 담아야 할 local 정보를 잃어버리고, attention map까지 지저분해진다는 것.

방법: Register 토큰 추가

• 모델이 임시 저장 공간이 필요해서 패치 토큰을 희생시키는 거라면, 전용 저장 공간을 따로 주면 됨 → 이게 register임.
• 구현: 패치 임베딩 뒤에 학습 가능한 추가 토큰(register) 몇 개를 CLS 토큰과 마찬가지로 시퀀스에 붙여서 transformer에 통과시킴. 출력 단계에서는 register 토큰을 버리고 CLS와 패치 토큰만 사용함.
• 학습은 처음부터(from scratch) 다시 진행. 기존 모델에 사후로 붙이는 방식이 아님.
• 개념적으로는 NLP의 Memory Transformer(Burtsev et al.)와 유사하게, 시퀀스에 쓰기/읽기가 가능한 여분의 슬롯을 두는 것임.

구분	기존 ViT	ViT + Registers
입력 시퀀스	[CLS] + 패치 토큰	[CLS] + [REG × n] + 패치 토큰
전역 정보 저장 위치	배경 패치 토큰을 임의로 재활용	register 토큰 (전용 슬롯)
출력에 사용하는 토큰	CLS, 패치 토큰	CLS, 패치 토큰 (register는 버림)
attention map	배경에 high-norm artifact 발생	artifact 사라지고 물체 중심으로 깔끔해짐
추가 비용	-	토큰 n개 추가 (n=4 기준 연산량 증가 미미)

register 개수

• register를 1개만 추가해도 artifact는 거의 사라짐.
• 개수를 늘리면 dense 태스크 성능이 조금씩 더 좋아지는 경향이 있어, 논문은 기본값으로 4개를 사용함.
• 학습 후 register들의 attention을 보면 각 register가 이미지의 서로 다른 영역을 보는 등 약간의 역할 분화가 관찰됨(논문은 이를 정성적으로만 보고함).

결과

• artifact 제거: DINOv2, DeiT-III, OpenCLIP 세 모델 모두 register를 붙이면 high-norm 토큰이 사라지고, 토큰 norm 분포가 정상 범위로 돌아옴. attention map도 물체 중심으로 깔끔해짐.
• 분류 성능: ImageNet linear probing은 기존과 거의 같거나 소폭 향상(DINOv2 ViT-g 기준 84.3% → 84.8%).
• dense prediction: ADE20k semantic segmentation(mIoU 46.6 → 47.9), NYUd depth estimation(RMSE 0.378 → 0.366)에서 향상. 패치 토큰이 자기 위치 정보를 온전히 갖게 되면서 픽셀 단위 태스크가 좋아진 것으로 설명함.
• object discovery(LOST): 패치 feature 유사도 기반으로 물체를 찾는 LOST 알고리즘에서 가장 큰 차이가 남. DINOv2는 원래 artifact 때문에 DINO보다 LOST 성능이 크게 떨어졌는데, register를 붙이면 VOC2007 corloc 기준 35.3 → 55.4로 회복됨.

태스크 (DINOv2 ViT-g)	register 없음	register 4개
ImageNet linear probing (Top-1)	84.3	84.8
ADE20k segmentation (mIoU)	46.6	47.9
NYUd depth (RMSE, 낮을수록 좋음)	0.378	0.366
LOST VOC2007 (corloc)	35.3	55.4

※ 수치는 논문 본문 표 기준이며, 모델 크기·설정에 따라 달라짐.

의료영상·segmentation 관점에서의 의미

• segmentation처럼 패치 토큰 하나하나의 local 정보를 그대로 쓰는 태스크에서는 artifact 토큰이 곧 예측 오류로 이어짐. CT/MRI처럼 균일한 배경(공기, 테이블 영역)이 넓은 영상은 논문이 말한 "정보량이 적은 패치"가 많아서, 같은 현상이 나타날 조건에 해당함.
• ViT 기반 foundation model을 backbone으로 쓸 때, register가 포함된 체크포인트(DINOv2 with registers 등)를 쓰는 것이 dense 태스크에 유리하다는 것이 논문 결과의 직접적인 함의임.
• 최근 국가 단위로 자체 foundation model을 만드는 흐름(소버린 AI 정리 글)에서도 이런 backbone 수준의 설계 디테일이 그대로 적용됨.

핵심 요약

• 큰 ViT(DINOv2, DeiT-III, OpenCLIP)는 학습 중반 이후 배경 패치 토큰을 전역 정보 저장용으로 재활용하면서 high-norm artifact 토큰이 생김.
• 이 토큰은 local 정보(위치·픽셀)는 잃고 global 정보(이미지 분류)는 많이 담고 있음.
• 해결책은 학습 가능한 register 토큰을 시퀀스에 추가하고 출력에서는 버리는 것. 1개만으로도 artifact 제거, 논문 기본값은 4개.
• 분류 성능은 유지하면서 segmentation·depth 등 dense prediction과 object discovery(LOST) 성능이 향상됨.
• ICLR 2024 oral 논문. 구현 비용이 거의 없어 이후 DINOv2 공개 체크포인트에 register 버전이 포함됨.

reference. https://arxiv.org/abs/2309.16588 (Darcet et al., Vision Transformers Need Registers, ICLR 2024)

함께 보면 좋은 글

• [IT] 소버린 AI(Sovereign AI)란 무엇인가 - 국가가 자체 AI 모델을 만드는 이유 정리
• [python(파이썬)] 파일명 숫자 순서대로 정렬하기 - sorted key, natsort로 1, 10, 2 순서 문제 해결