비전 트랜스포머(ViT)에 레지스터 토큰이 필요한 이유

 (openreview.net)
1P by GN⁺ | ★ favorite | 댓글 1개
  • ViT 특징 맵의 고노름 토큰은 정보량이 낮은 배경 영역에 주로 나타나며, 패치 토큰이 내부 계산용 저장 공간처럼 재활용되는 현상으로 해석됨
  • 해결책은 이미지 패치와 별도의 레지스터 토큰을 입력 시퀀스에 추가해, 내부 계산 역할을 전용 토큰으로 분리하는 방식임
  • 지도학습·자기지도학습 ViT 모두에서 아티팩트가 줄고 특징 맵과 attention map이 매끄러워져 downstream 시각 작업에 유리함
  • object discovery와 dense prediction 개선이 보고됐지만, OpenCLIP에서는 object localization 성능이 소폭 하락해 “모든 모델에서 개선”된다는 표현은 수정됨
  • 4~16개 레지스터 토큰을 넣으면 2~6% 추가 계산량이 생기며, 토큰 수를 늘린다고 dense task 성능이 항상 오르지는 않음

ViT 특징 맵에서 드러난 아티팩트

  • Vision Transformer는 시각 표현 학습에서 강력한 도구로 쓰이지만, 지도학습 및 자기지도학습 ViT의 특징 맵에서 아티팩트가 확인됨
  • 아티팩트는 추론 중 나타나는 고노름 토큰으로 관찰됨
    • 주로 이미지의 정보량이 낮은 배경 영역에서 발생함
    • 네트워크가 해당 토큰을 내부 계산에 재활용하는 형태로 보임
  • 이 토큰들은 주변 패치와 비슷한 배경 패치에서 나오며, 지역 패치 정보는 적고 전역 이미지 정보는 더 많이 담는 것으로 평가됨
  • DINOv2와 다른 ViT에서는 attention map의 공간 해상도와 정확도 저하가 downstream object discovery와 localization 성능에 영향을 줄 수 있음

레지스터 토큰 방식

  • 제안된 해법은 Vision Transformer 입력 시퀀스에 이미지 패치와 별도의 레지스터 토큰을 추가하는 것임
  • 레지스터는 CLS token처럼 이미지 패치 토큰과 분리되지만, loss 계산에는 직접 쓰이지 않는 토큰으로 정리됨
  • 목적은 기존 패치 토큰이 내부 계산용 “저장 공간”처럼 쓰이지 않도록, 그 역할을 전용 토큰이 맡게 하는 데 있음
  • 이 방식은 기존 Memory Transformer와 관련이 있지만, 특징 아티팩트를 분석하고 이를 고치기 위해 새 토큰을 넣는 접근은 이 논문의 분석에 특화됨

특징 맵과 downstream 작업의 변화

  • 레지스터를 추가하면 지도학습 및 자기지도학습 모델 모두에서 해당 아티팩트가 제거되는 것으로 보고됨
  • 특징 맵과 attention map이 더 매끄러워져 downstream visual processing에 유리해짐
  • 자기지도 시각 모델의 dense visual prediction task에서 새로운 state of the art를 달성한 것으로 정리됨
  • 더 큰 모델에서도 object discovery 방법을 사용할 수 있게 됨
  • DINOv2와 DeiT-III에서는 object discovery 성능이 높아지고, 이미지 수준 task 성능은 대체로 유지됨

모델별 예외와 추가 분석

  • OpenCLIP에서는 레지스터를 추가해도 object localization 결과가 개선되지 않았고, 표의 수치상 소폭 성능 하락이 있음
    • 저자 답변은 “모든 모델에서 object discovery 성능이 개선된다”는 표현이 수치로 뒷받침되지 않는다고 인정함
    • 해당 내용은 Section 3.3에서 수정됨
  • OpenCLIP은 Figure 7에서 norm 감소가 보이지만, Table 3의 object localization 개선으로 이어지지 않는 불일치가 논의됨
    • LOST 평가에서는 모델별로 keys, queries, values 중 최고 성능 embedding을 선택함
    • CLIP에서는 values가 최고였고, artifacts는 keys나 queries에서는 보이지만 values에서는 보이지 않아 Table 3 및 정성 분석과 일관됨
  • MAE에서는 같은 “outlier patches”가 나타나지 않는 것으로 보임
    • MAE가 local loss만으로 학습돼 전역 정보를 모을 필요가 없거나 적기 때문일 수 있음
    • MAE-ViT-Large의 linear probing ImageNet 분류 정확도는 75%로 다른 SSL 방법보다 낮음

성능, 비용, 한계

  • 레지스터는 단순한 방법이지만 추가 토큰 때문에 계산량이 늘어남
    • 4~16개 레지스터 토큰을 추가할 때 2~6% 계산량 증가가 보고됨
  • 레지스터 개수 증가가 항상 dense task 개선으로 이어지지는 않음
    • 0개에서 1개 레지스터로 늘릴 때 아티팩트가 제거되고 segmentation 및 depth prediction이 개선됨
    • 1개 이후 더 많은 레지스터를 추가해도 feature map은 이미 깨끗해져 segmentation과 depth prediction 성능은 크게 오르지 않음
    • classification 성능은 레지스터가 더 많을수록 추가 개선되는 경향이 있었지만, 그 이유에 대한 명확한 직관은 아직 없음
  • NYU depth dataset에서는 8개에서 16개 레지스터로 늘릴 때 RMSE가 소폭 증가함
    • 시퀀스 길이 증가로 최적 hyperparameter가 달라졌을 수 있음
    • 8개에서 16개로 갈 때의 0.03 RMSE 차이는 0개에서 1개로 갈 때의 0.1 RMSE 차이에 비해 작아 단순 노이즈일 수도 있음

OpenReview 채택과 리뷰 쟁점

  • OpenReview 결정은 Accept (oral) 임
  • 메타 리뷰에서는 네 명의 리뷰어가 모두 8점을 줬고, 논문의 기술적 기여가 인정됨
  • 리뷰어들은 성능 불일치, dataset bias와 SSL 모델의 관계, 최적화 설정 선택, gradient clipping 영향 등을 우려함
  • 저자들은 반박 과정에서 추가 실험과 appendix 분석을 더함
  • 이후 논의에서는 레지스터가 outlier의 동작을 이어받는지, image token이 더 local한 성격으로 돌아가는지, CLS token norm이 어떻게 변하는지 같은 내부 메커니즘 분석이 보강됨

함께 보면 좋은 글 β

GN⁺   [-]
Hacker News 의견들

  • 논문에 따르면 register는 Vision Transformer 모델 학습 중 입력 시퀀스에 덧붙이는 추가 학습 가능 토큰임
    패치 임베딩 층 뒤에 [CLS] 토큰처럼 학습 가능한 값으로 추가되고, Vision Transformer 끝에서는 register 토큰을 버린 뒤 [CLS] 토큰과 패치 토큰만 이미지 표현으로 사용함
    이 토큰은 모델이 순전파 중 전역 정보를 저장·처리·조회할 공간을 제공해, 패치 토큰을 그 용도로 억지로 쓰지 않아도 되게 함
    그 결과 학습된 Vision Transformer의 특징 맵에 나타나는 아티팩트와 높은 노름의 이상치 토큰이 사라지고, 특징 맵이 더 매끄러워지며 dense prediction 성능과 비지도 객체 발견도 좋아짐
    추론 비용이 2%만 늘어나는데 ViT 성능을 꽤 개선하니 거의 공짜 점심에 가까운 깔끔한 결과로 보임

    • 토큰을 여기저기 더했더니 마법처럼 되는 이 흐름이 꽤 수상해 보이고, 확률론까지 엮이면 이산수학 쪽 사람들은 가볍게 넘기기 어려울 듯함
      확률적 정규식도 한 가지인데, 이건 수학적 방탕함이 완전히 다른 수준임
      그래도 이렇게 동작한다니 정말 놀라움
    • BERT에 [CLS] 토큰을 여러 개 추가하려던 시도도 있었지만 결과는 덜 인상적이었음: https://arxiv.org/pdf/2210.05043
    • 일반적인 텍스트 Transformer에도 여기서 배울 점이 있는지 궁금함
      attention sink나 pause token과 비슷하게 들림
    • 이런 토큰을 쓰려면 최종 층의 결과를 먼저 토큰으로 되돌리고, 다음 패스에서 다시 임베딩해야 함
      한 패스의 원시 출력을 다음 입력으로 그대로 넘기는 방식을 조사한 사람이 있는지 궁금함

  • 관련 있어 보임: "Let's Think Dot by Dot: Hidden Computation in Transformer Language Models" https://arxiv.org/abs/2404.15758
    이 논문은 언어 모델의 chain-of-thought 성능 향상이 사람 같은 작업 분해 때문인지, 아니면 추가 토큰이 주는 계산량 때문인지가 불명확하다고 봄
    Transformer가 의미 없는 채움 토큰, 예를 들어 ......만으로도 중간 토큰 없이 답할 때 못 풀던 어려운 알고리즘 과제를 풀 수 있음을 보였고, 다만 그런 채움 토큰 사용을 학습시키려면 구체적이고 촘촘한 감독이 필요하다고 함
    그래서 텍스트 LLM에도 특별한 계산/register 토큰을 넣으면 이득이 있을지 궁금함
    추가 논의: https://news.ycombinator.com/item?id=40182695, https://www.reddit.com/r/LocalLLaMA/comments/1cf2w5a/transfo...

  • 지난 몇 년간 비전과 멀티모달 인식 분야에서 논문·아이디어를 수십 개, 어쩌면 100개 넘게 시험했는데, 이건 드물게 모든 것이 잘 동작한 사례였음
    예를 들어 이 모델은 register 토큰 4개와 Matryoshka식 손실을 함께 써서 학습하고, CLIP 대안을 찾는 사람에게 유용할 수 있는 매우 작은 64차원 임베딩을 만들어냄: https://huggingface.co/unum-cloud/uform3-image-text-english-...

  • ICLR에서 봤는데 올해 최고 논문 중 하나였고, 포스터 세션에서도 그 점이 분명히 드러났음
    저자들에게 축하를 보냄

  • 몇 달째 DinoV2를 쓰고 있고, [CLS] + 패치 토큰과 함께 register 토큰 4개가 있는 모델도 시험해 봤음
    기존 register 없는 모델에서 나온 여러 임베딩 토큰이 이미 솔루션 일부라서, 0 register 모델과 4 register 모델 사이에 [CLS] 토큰이 정렬되지 않아 새 register 모델은 채택하지 않았음
    두 모델 사이에서 [CLS]와 패치 토큰이 어떻게든 정렬되면 좋겠음

  • 비전 모델용 attention sink처럼 보임

답변달기