비전 트랜스포머(ViT) 작동 원리: 이미지 패치 분할, 위치 임베딩, 분류 토큰까지 정리

핵심 메시지

• 비전 트랜스포머는 2020년 구글이 발표한 모델로, 이미지를 작은 패치로 잘라 트랜스포머에 그대로 넣는다는 발상에서 출발한다.
• 각 패치는 선형 변환을 거쳐 정해진 크기의 벡터가 되고, 실제 구현에서는 커널 크기를 패치 크기와 같게 맞춘 합성곱 층 하나로 모든 패치를 한꺼번에 처리한다.
• 패치 벡터만으로는 순서 정보가 없기 때문에 자연어 처리와 똑같이 위치 임베딩을 더해 각 조각이 이미지의 어느 자리에 있었는지 알려 준다.
• 분류 작업에는 별도의 분류 토큰을 앞에 붙이고, 인코더를 통과한 뒤 그 토큰에 대응하는 벡터 하나만 예측 헤드에 넣어 클래스를 맞힌다.
• 인코더 블록 구조는 어텐션 논문의 것을 그대로 유지했고, 그 결과 오랫동안 컴퓨터 비전을 지배해 온 합성곱 신경망과 대등하거나 더 나은 성능을 냈다.

쉽게 이해하기

트랜스포머와 셀프 어텐션은 원래 자연어 처리 모델들의 약점을 보완하려고 설계된 구조였다. 그런데 구글 연구자들은 텍스트 대신 이미지를 입력으로 넣어도 같은 구조를 쓸 수 있다는 것을 확인했고, 그렇게 나온 것이 2020년 발표된 비전 트랜스포머다. 합성곱 신경망이 컴퓨터 비전을 오래 지배해 온 상황에서, 구조를 거의 손대지 않고도 대등하거나 더 나은 성능을 냈다는 점이 이 논문이 놀라웠던 이유다.

핵심 질문은 이미지를 어떻게 트랜스포머에 넣느냐다. 방법은 이미지를 작은 패치 여러 개로 쪼개고, 각 패치를 선형 변환으로 정해진 크기의 벡터로 바꾸는 것이다. 이렇게 만들어진 벡터들의 나열이 모델의 입력이 되며, 각 패치를 시각적인 토큰이라고 생각하면 이해가 쉽다. 어텐션 메커니즘은 이 토큰들 사이의 관계, 즉 이미지의 서로 다른 조각들이 어떻게 상호작용하는지를 학습한다.

패치를 하나씩 변환하는 것은 개념 설명에는 좋지만 실제로는 비효율적이다. 구현에서는 합성곱 층 하나로 모든 패치를 동시에 처리한다. 커널의 크기를 패치 크기와 같게 두면 커널이 이미지 위를 미끄러지며 이미지를 패치 단위로 나누는 효과가 나고, 커널 개수를 원하는 벡터 차원과 같게 두면 결과 벡터의 크기까지 통제할 수 있다. 합성곱 층은 결국 행렬들의 모음이라 병렬화도 쉽다.

이미지를 이 합성곱 층에 통과시키면 하나의 텐서가 나오는데, 텐서의 한 방향은 벡터 차원이고 나머지 방향은 이미지에서 잘려 나온 패치의 배치와 대응한다. 가로로 세 개, 세로로 두 개의 패치가 나오는 이미지라면 텐서 안에는 여섯 개의 벡터가 자리 잡고, 첫 번째 벡터는 첫 번째 패치에, 두 번째 벡터는 두 번째 패치에 차례로 대응한다. 합성곱 층이 하는 일은 결국 단순한 선형 변환이며, 이미지를 트랜스포머가 받을 수 있는 모양으로 정리해 주는 역할이다.

여기까지 만든 벡터들에는 순서 정보가 없다. 그래서 자연어 처리에서 하던 것처럼 위치 임베딩을 각 패치 벡터에 더해 준다. 첫 번째 패치에는 위치 임베딩의 첫 번째 벡터가, 두 번째 패치에는 두 번째 벡터가 대응하는 식이다. 이렇게 시각 토큰과 위치 정보를 합치면 모델이 조각들의 배치를 이해하고 원래 이미지 구조를 재구성할 수 있게 된다.

분류 작업을 위해서는 흔히 분류 토큰을 하나 더 붙인다. 이는 모델이 학습하는 벡터일 뿐이며, 자연어 처리에서 문장 앞에 특수 토큰을 붙여 그 벡터로 예측하는 방식과 같다. 인코더에 들어갈 때 이 토큰은 다른 벡터들과 아무 관계가 없지만, 여러 인코더 블록을 지나면서 이미지 전체의 정보를 담게 되고 최종적으로 예측에 쓰인다.

인코더 내부는 어텐션 논문의 구조를 그대로 따른다. 멀티 헤드 어텐션 층을 지난 결과에 잔차를 더하고 정규화해 기울기가 폭발하거나 사라지지 않게 한 뒤, 피드포워드 신경망을 통과시키고 다시 잔차를 더해 정규화한다. 이런 인코더 블록이 여러 개 쌓이고, 앞 블록의 출력이 다음 블록의 입력이 되면서 데이터 안의 관계를 학습한다.

마지막은 예측 헤드다. 보통 선형 층 하나 또는 여러 개로, 모델 내부 차원의 벡터를 예측하려는 클래스 개수만큼의 차원으로 옮긴다. 출력 벡터의 각 원소는 자동차, 사람, 트럭 같은 개별 대상에 대응하고 그 대상일 확률을 담는다. 분류에서는 인코더 출력 중 분류 토큰에 해당하는 첫 번째 벡터만 사용한다. 결국 텍스트 데이터였다면 했을 일과 똑같으며, 이미지가 아닌 다른 종류의 데이터를 넣어도 같은 방식이 성립한다는 점이 이 구조의 핵심이다.

주요 인사이트

• 비전 트랜스포머의 기여는 새로운 아키텍처를 발명한 것이 아니라 입력을 바꾼 것이다. 이미지를 패치의 나열로 재정의하는 순간 기존 트랜스포머가 그대로 컴퓨터 비전 모델이 된다.
• 패치를 벡터로 만드는 선형 변환을 합성곱 층으로 구현한 것은 성능 최적화이면서, 합성곱이 트랜스포머의 반대편이 아니라 그 안에 부품으로 들어갈 수 있음을 보여 준다.
• 커널 크기로 패치 크기를, 커널 개수로 벡터 차원을 정한다는 설계는 하이퍼파라미터 두 개로 입력 표현 전체를 통제할 수 있게 만든다.
• 분류 토큰은 이미지 전체를 요약하는 자리를 모델 안에 미리 만들어 두는 장치다. 인코더 출력의 다른 벡터들이 모두 이미지의 서로 다른 측면을 담기 때문에, 요약용 벡터를 따로 두는 편이 예측에 유리하다.
• 이 구조가 데이터 종류에 무관하다는 점은 이후 멀티모달 모델들이 텍스트와 이미지, 그 밖의 신호를 같은 인코더에 태우는 흐름의 출발점이 된다.

자주 묻는 질문

관련 AI 소식