스윈 트랜스포머 구조 해설: 이동 윈도 셀프어텐션과 계층적 텐서로 CNN 백본 대체하기

핵심 메시지

• 스윈 트랜스포머는 이미지 하나에서 여러 해상도의 텐서를 만들어 내는 계층적 비전 트랜스포머로, 객체 탐지와 분할처럼 물체 크기가 제각각인 과제에서 합성곱 신경망 백본을 대체하려고 설계됐다.
• 셀프어텐션을 이미지 전체가 아니라 M×M 윈도 안에서만 계산해, 연산량이 픽셀 수의 제곱이 아니라 픽셀 수에 비례해 늘도록 바꾼 것이 핵심이다.
• 윈도를 고정하면 경계에 걸친 이웃 픽셀끼리 서로를 볼 수 없기 때문에, 다음 블록에서는 이미지를 왼쪽 위로 밀어 윈도를 다시 잡는 이동 윈도 어텐션을 이어서 수행한다.
• 단계 사이마다 2×2 이웃 패치를 합쳐 해상도를 절반으로 줄이고 채널 깊이를 늘리면서, 크기가 다른 물체를 잡아낼 수 있는 텐서 계층이 만들어진다.
• 발표에서 비교한 객체 탐지 실험에서는 ResNet 백본을 스윈 트랜스포머로 바꿨을 때 성능이 거의 전 항목에서 앞섰지만, 추론 속도는 다소 느렸다.

쉽게 이해하기

영상은 스윈 트랜스포머를 "합성곱 신경망 백본을 대체하려고 만든 계층적 비전 트랜스포머"라는 한 문장으로 정의한 뒤, 그 문장에 들어 있는 용어를 하나씩 풀어 간다. 계층적이라는 말은 이미지 한 장에서 해상도가 서로 다른 여러 개의 텐서를 뽑아낸다는 뜻이고, 백본이란 객체 탐지 같은 파이프라인의 맨 앞에서 이미지를 텐서로 바꿔 주는 부분을 가리킨다. 왜 이런 모델이 필요했는지는 역사적 맥락으로 설명된다. 2015년 Faster R-CNN이 객체 탐지의 최신 기법으로 자리 잡았을 때 그 앞단의 백본은 합성곱 신경망이었고, 2017년 번역 문제를 풀기 위해 등장한 트랜스포머가 2020년 무렵 BERT와 GPT를 거치며 자연어 처리를 장악한 뒤에도 상황은 크게 달라지지 않았다. 트랜스포머를 객체 탐지에 쓴 DETR조차 백본은 여전히 ResNet 계열 합성곱 신경망이었다.

트랜스포머를 백본 자리에 그대로 넣기 어려웠던 이유는 두 가지다. 첫째, 셀프어텐션은 질의 행렬과 키 행렬을 곱하는 구조라 입력 개수의 제곱에 비례해 연산이 늘어나는데, 이미지에서는 입력 단위가 픽셀이라 고해상도 이미지에서는 계산이 사실상 불가능해진다. 둘째, 트랜스포머가 텐서를 하나만 내놓으면 큰 물체와 작은 물체가 섞여 있는 이미지에서 탐지 성능이 떨어진다.

스윈 트랜스포머는 두 문제를 각각 겨냥한다. 이미지를 4×4 픽셀 묶음으로 잘라(픽셀당 값 3개이므로 깊이 48) 임베딩한 뒤, 텐서를 M×M 윈도로 나누고 각 윈도 안에서만 어텐션을 계산한다. 이렇게 하면 윈도 하나에서 생기는 어텐션 행렬 크기가 고정되고, 전체 연산은 이미지의 가로×세로에 선형으로 비례하게 된다. 어텐션 헤드는 8개를 쓰고, 소프트맥스 직전에 상대 위치 인코딩을 더한다.

다만 윈도를 고정해 두면 바로 옆에 붙어 있어도 다른 윈도에 속한 픽셀끼리는 서로를 참조하지 못한다. 그래서 두 번째 블록에서는 이미지를 M의 절반만큼 왼쪽과 위로 순환 이동시킨 뒤 윈도를 다시 나눈다. 이때 화면 밖으로 밀려 반대편으로 감겨 들어온 픽셀들은 원래 이웃이 아니므로, 텐서를 아홉 개 구역으로 나눈 마스크를 소프트맥스 앞에서 더해 그런 조합의 어텐션을 막는다. 계산이 끝나면 이동과 분할을 정확히 역순으로 되돌린다.

이 블록들은 단계별로 2개, 2개, 6개, 2개씩 쌓이고, 단계 사이에서는 2×2 이웃 패치를 하나로 합쳐 가로세로를 절반으로 줄이면서 채널 깊이를 키운다. 그 결과 해상도와 깊이가 다른 텐서들의 계층이 만들어지고, 여기에 특징 피라미드 네트워크를 붙여 채널 깊이를 맞추면 크기가 제각각인 물체를 탐지하는 백본으로 그대로 쓸 수 있다. 백본만 ResNet에서 스윈 트랜스포머로 바꾼 비교에서는 성능이 대체로 올라갔고, 대신 추론은 조금 더 느려졌다.

주요 인사이트

• "어텐션을 이미지 전체에 걸지 않는다"는 선택 하나가 제곱 복잡도를 선형으로 바꿨다. 성능을 위해 구조를 키우는 대신, 계산 범위를 국소적으로 제한하고 그 손실을 다른 장치로 메우는 설계다.
• 이동 윈도는 그 손실을 메우는 장치다. 윈도를 절반만큼 밀어 다시 나누는 것만으로 이전 블록에서 끊겼던 이웃 관계가 다음 블록에서 연결되므로, 전역 어텐션 없이도 넓은 문맥이 층을 거치며 전달된다.
• 순환 이동은 공짜가 아니다. 반대편에서 감겨 들어온 픽셀이 가짜 이웃이 되는 부작용이 생기고, 이를 막기 위한 아홉 구역 마스크가 반드시 필요하다. 영상의 퀴즈에서도 이 마스크의 필요성이 정답 항목으로 제시된다.
• 패치 병합으로 만들어지는 텐서 계층은 합성곱 신경망이 원래 자연스럽게 갖고 있던 다중 스케일 특성을 트랜스포머 쪽에서 다시 확보한 것이다. 특징 피라미드 네트워크와 결합하기 쉬운 것도 그래서다.
• 정확도가 올라가도 추론이 느려질 수 있다는 점은 실제 도입 시 그대로 트레이드오프가 된다. 백본 교체는 성능 수치만이 아니라 지연 시간 예산까지 함께 봐야 하는 결정이다.

자주 묻는 질문

관련 AI 소식