AI도 연습장이 필요해!!
VISION TRANSFORMERS NEED REGISTERS
AI(ViT) 에게 '연습장(Register)'이 필요해!!
1. 수학 시험지와 빽빽한 낙서의 추억
고등학교 때 치르던 수능 수학 시간을 떠올려 볼까요?
복잡한 미적분 문제를 풀 때, 연습장이 없어서 시험지 여백 귀퉁이에 깨알같이 계산식을 적었던 기억, 다들 있으시죠? 공간은 좁은데 적어야 할 계산은 많고... 결국 시험지 구석구석, 심지어 지문 위까지 숫자를 끄적이다 보면 시험지가 시커멓게 더러워지곤 했습니다. 나중에는 내가 쓴 글씨를 내가 못 알아봐서 계산 실수를 하기도 하고, 정작 중요한 문제(정답)가 낙서에 가려 잘 안 보이기도 했죠.
그런데 만약, 감독관이 "자, 여기 깨끗한 '연습장' 4장을 줄 테니 계산은 여기에 하세요"라고 했다면 어땠을까요? 시험지는 깨끗해지고, 계산은 연습장에 체계적으로 정리되어 정답률도 오르고 검토하기도 훨씬 편했을 겁니다.
오늘 소개할 논문이 바로 "AI에게 연습장(Register)을 쥐여주자!"라는 내용을 담고 있습니다.
2. AI의 시험지가 더러워지고 있다? (The Problem)
요즘 가장 핫한 이미지 인식 모델인 Vision Transformer(ViT), 그중에서도 DINOv2 같은 모델들을 뜯어보니 아주 이상한 현상이 발견되었습니다.
AI가 이미지를 볼 때 어디를 주목하는지 보여주는 '어텐션 맵(Attention Map)'을 그려봤더니, 이미지의 배경(하늘, 잔디 등)에 뜬금없이 강한 점(Artifacts)들이 찍혀 있는 겁니다.
위 그림을 보세요.
Input: 곤충, 철탑, 새 사진이 있습니다.
Without registers (가운데): 모델이 물체뿐만 아니라, 아무 의미 없는 배경(허공)에 밝은 점들을 찍고 있습니다. 마치 시험지 구석에 낙서를 해놓은 것처럼 지저분하죠.
With registers (오른쪽): 연습장(Register)을 사용한 결과로, input의 물체 윤곽을 잘 잡아내고 있습니다.
도대체 AI는 왜 아무 정보도 없는 맨땅에 점을 찍었을까요? 연구진이 분석해 보니, 이 점들은 AI를 위한 연습장이었습니다.. AI 모델은 이미지를 처리하면서 전체적인 정보(Global Information)를 어딘가에 저장하고 계산해야 하는데, 정보를 저장할 공간이 부족했던 겁니다. 그래서 AI는 이렇게 생각한 거죠.
"아, 이 복잡한 계산 값을 어디에 적어놓지? 에라 모르겠다. 저기 배경(하늘)은 별로 안 중요하니까 저기다가 적어놓자!"
즉, 시험지 여백(배경)을 연습장 대신 활용하여 계산식을 마구 적어놓은 것과 같습니다. 이것을 논문에서는 "재활용(Repurposed)"이라고 표현합니다.
3. 해결책: 야, 여기 '연습장' 써! (The Solution)
연구진이 내놓은 해결책은 너무나 심플하고 천재적입니다. 모델이 억지로 배경에 정보를 쑤셔 넣지 않도록, 연습장역할을 할 수 있는, 아무런 의미가 없는 빈 토큰(Blank Tokens) 몇 개를 슬쩍 넣어주는 것입니다.
논문에서는 이것을 "레지스터(Registers)"라고 부릅니다.
기존: 이미지 조각(Patches) + 정답 토큰([CLS])만 있음.
제안: 이미지 조각 + 정답 토큰 + [REG] (레지스터 토큰) 추가.
마치 시험 치는 학생에게 "계산은 시험지에 하지 말고, 이 노란색 연습장([REG])에 해!"라고 말하는 것과 같습니다. 이 레지스터 토큰들은 나중에 정답을 낼 때는 버려지지만, 중간 계산 과정에서는 아주 유용하게 쓰입니다.
4. 결과: 깨끗해진 시험지, 올라간 점수 (The Results)
AI에게 연습장(레지스터)을 쥐여주자 놀라운 일이 벌어졌습니다.
1) 시험지가 깨끗해졌습니다.
위 차트들을 보면, Register를 사용한 오른쪽(With registers)에서, 지저분함이 감소됨을 확인할 수 있습니다. 이를 통해서 배경의 불필요하던 점들이 싹 사라지고, AI가 정확하게 물체를 바라볼 수 있게 되었습니다. 계산은 '레지스터'에 따로 하고, 이미지 토큰에는 순수하게 이미지 정보만 남기게 된 것입니다.
2) 성능도 좋아졌습니다. 지저분한 노이즈가 사라지니, 물체를 찾거나(Object Discovery) 화면을 픽셀 단위로 분류하는 작업(Dense Prediction)에서 성능이 향상되었습니다. 심지어 연습장(레지스터)을 1개만 줘도 효과가 있었고, 4개 정도 주었을 때 가장 효율이 좋았다고 합니다.
5. 결론!: 단순함의 미학
이 논문은 복잡한 수식을 뜯어고친 게 아닙니다. 그저 "정보를 처리할 여유 공간(연습장)을 따로 만들어주자"는 아주 직관적인 아이디어 하나로 AI의 고질적인 문제를 해결했습니다. 수능 시험장에 연습장을 챙겨가는 마음으로, AI에게도 Register를 챙겨주는 센스! 이것이 바로 2024년 인기 있던 논문 Darcet, Timothée, et al. "Vision transformers need registers." arXiv preprint arXiv:2309.16588 (2023). 였습니다!!