An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale (VIT) (부록 추가 필요)

https://ar5iv.labs.arxiv.org/html/2010.11929

An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

 

초록

Transformer 아키텍처는 자연어 처리 작업에서 사실상의 표준이 되었지만, 컴퓨터 비전 분야에서의 응용은 여전히 제한적입니다. 비전에서는 주로 주의(attention)를 합성곱 신경망(CNN)과 결합하여 사용하거나, CNN의 일부 구성 요소를 대체하는 방식으로 사용하되 전체 구조는 유지하고 있습니다. 우리는 CNN에 대한 이러한 의존이 필수적이지 않으며, 이미지 패치의 시퀀스에 직접 적용되는 순수한 Transformer도 이미지 분류 작업에서 매우 우수한 성능을 낼 수 있음을 보여줍니다. 대량의 데이터로 사전 학습한 후, 중형 또는 소형 이미지 인식 벤치마크(예: ImageNet, CIFAR-100, VTAB 등)로 전이할 경우 Vision Transformer(ViT)는 최신의 CNN과 비교했을 때 훈련에 훨씬 적은 계산 자원을 요구하면서도 탁월한 결과를 얻습니다.

 

사전 학습된 모델 및 세부 튜닝 코드는 여기에서 확인할 수 있습니다.

 

1. 서론

자기 주의(self-attention) 기반 아키텍처, 특히 Transformer(Vaswani et al., 2017)는 자연어 처리(NLP)에서 선택되는 모델로 자리 잡았습니다. 주요 접근 방식은 대규모 텍스트 말뭉치에 대해 사전 학습을 한 후, 더 작은 작업별 데이터셋에 대해 미세 조정을 하는 방식입니다(Devlin et al., 2019). Transformer는 계산 효율성과 확장성 덕분에, 1000억 개 이상의 파라미터를 가진 모델을 훈련하는 것이 가능해졌으며(Brown et al., 2020; Lepikhin et al., 2020), 모델과 데이터셋이 커져도 성능이 포화되는 조짐은 아직 없습니다.

 

하지만 컴퓨터 비전 분야에서는 여전히 합성곱 아키텍처가 지배적입니다(LeCun et al., 1989; Krizhevsky et al., 2012; He et al., 2016). NLP의 성공에 영감을 받아, 여러 연구는 CNN과 유사한 아키텍처를 자기 주의와 결합하거나(Wang et al., 2018; Carion et al., 2020), 일부 연구는 아예 합성곱을 대체하려고 시도하고 있습니다(Ramachandran et al., 2019; Wang et al., 2020a). 후자의 모델은 이론적으로는 효율적이지만, 특수한 주의 패턴을 사용하기 때문에 현대 하드웨어 가속기에서 효과적으로 확장되지 않았습니다. 따라서 대규모 이미지 인식에서는 여전히 고전적인 ResNet과 같은 아키텍처가 최신 상태의 성능을 유지하고 있습니다(Mahajan et al., 2018; Xie et al., 2020; Kolesnikov et al., 2020).

 

NLP에서 Transformer 확장 성공에 영감을 받아, 우리는 가능한 한 적은 수정을 통해 표준 Transformer를 이미지에 직접 적용하는 실험을 진행했습니다. 이를 위해 이미지를 여러 패치로 나누고, 이 패치들의 선형 임베딩 시퀀스를 Transformer에 입력으로 제공합니다. 이미지 패치는 NLP 응용에서 토큰(단어)처럼 취급됩니다. 우리는 이 모델을 지도 학습 방식으로 이미지 분류에 대해 훈련했습니다.

 

ImageNet과 같은 중형 데이터셋에서 강력한 정규화 없이 훈련된 이 모델들은 비슷한 크기의 ResNet에 비해 몇 퍼센트 정도 낮은 정확도를 보였습니다. 이는 다소 실망스러울 수 있지만 예상 가능한 결과이기도 합니다. Transformer는 CNN이 가진 번역 불변성(translation equivariance)과 지역성(locality) 같은 귀납적 편향이 부족하므로, 데이터가 충분하지 않을 경우 일반화가 잘 이루어지지 않기 때문입니다.

 

그러나, 더 큰 데이터셋(1,400만~3억 개 이미지)에서 훈련될 경우 상황이 달라집니다. 우리는 대규모 훈련이 귀납적 편향을 능가한다는 것을 발견했습니다. 우리의 Vision Transformer(ViT)는 충분한 규모로 사전 학습된 후, 더 적은 데이터가 있는 작업에 전이되었을 때 뛰어난 결과를 얻었습니다. 공개된 ImageNet-21k 데이터셋 또는 내부 JFT-300M 데이터셋에서 사전 학습된 ViT는 여러 이미지 인식 벤치마크에서 최신 성능에 접근하거나 이를 능가했습니다. 특히, 최고 성능 모델은 ImageNet에서 88.55%, ImageNet-ReaL에서 90.72%, CIFAR-100에서 94.55%, 그리고 19개의 작업으로 구성된 VTAB에서 77.63%의 정확도를 기록했습니다.

 

2. 관련 연구

Transformer는 Vaswani et al. (2017)에 의해 기계 번역을 위해 제안되었으며, 이후 많은 자연어 처리(NLP) 작업에서 최신 방법으로 자리잡았습니다. 대형 Transformer 기반 모델은 주로 대규모 말뭉치에 대해 사전 학습된 후, 특정 작업에 맞춰 미세 조정됩니다. BERT(Devlin et al., 2019)는 노이즈 제거 방식의 자기 지도 사전 학습을 사용하고, GPT 계열의 연구는 언어 모델링을 사전 학습 작업으로 사용합니다(Radford et al., 2018; 2019; Brown et al., 2020).

 

이미지에 자기 주의를 단순히 적용하려면 각 픽셀이 모든 다른 픽셀에 주의를 기울여야 합니다. 이 경우, 픽셀 수에 따라 계산 비용이 제곱으로 증가하므로 현실적인 입력 크기에 확장되지 않습니다. 따라서 이미지 처리 문맥에서 Transformer를 적용하기 위해 여러 가지 근사 방식이 제안되었습니다. Parmar et al. (2018)은 각 쿼리 픽셀에 대해 전역이 아닌 국소 이웃에서만 자기 주의를 적용했습니다. 이러한 국소 다중 헤드 점곱(self-attention) 블록은 합성곱을 완전히 대체할 수 있습니다(Hu et al., 2019; Ramachandran et al., 2019; Zhao et al., 2020). 다른 연구에서는 Sparse Transformers(Child et al., 2019)가 이미지에 적용할 수 있도록 전역 자기 주의에 대한 확장 가능한 근사치를 사용했습니다. 또 다른 방법으로는 다양한 크기의 블록에 주의를 적용하는 방식이 있으며(Weissenborn et al., 2019), 극단적인 경우 개별 축에만 적용하기도 합니다(Ho et al., 2019; Wang et al., 2020a). 이러한 특수화된 주의 아키텍처들은 컴퓨터 비전 작업에서 유망한 결과를 보여주지만, 하드웨어 가속기에서 효율적으로 구현하려면 복잡한 엔지니어링이 필요합니다.

 

우리의 연구와 가장 관련이 깊은 모델은 Cordonnier et al. (2020)의 모델로, 입력 이미지에서 2×2 크기의 패치를 추출한 후, 그 위에 완전한 자기 주의를 적용합니다. 이 모델은 ViT와 매우 유사하지만, 우리의 연구는 대규모 사전 학습이 표준 Transformer를 최신 CNN과 경쟁할 수 있을 정도로 만들거나 더 뛰어난 성능을 발휘할 수 있음을 보여줍니다. 또한, Cordonnier et al. (2020)은 2×2 픽셀의 작은 패치 크기를 사용하므로, 모델은 낮은 해상도의 이미지에만 적용될 수 있지만, 우리는 중간 해상도의 이미지도 처리할 수 있습니다.

 

합성곱 신경망(CNN)을 자기 주의와 결합하려는 많은 시도도 있었습니다. 예를 들어, 이미지 분류를 위해 특징 맵을 보강하거나(Bello et al., 2019), 객체 탐지(Hu et al., 2018; Carion et al., 2020), 비디오 처리(Wang et al., 2018; Sun et al., 2019), 이미지 분류(Wu et al., 2020), 비지도 객체 탐지(Locatello et al., 2020), 또는 통합된 텍스트-비전 작업(Chen et al., 2020c; Lu et al., 2019; Li et al., 2019)을 위해 CNN의 출력을 추가로 처리하는 방식이 사용되었습니다.

 

또 다른 관련된 모델은 이미지 GPT(iGPT)(Chen et al., 2020a)입니다. 이는 이미지를 축소한 후 해상도와 색상 공간을 줄인 픽셀에 대해 Transformer를 적용합니다. 이 모델은 비지도 학습 방식으로 생성 모델로 훈련되며, 결과적으로 생성된 표현은 미세 조정되거나 선형 방식으로 분류 성능을 평가할 수 있습니다. 이 모델은 ImageNet에서 최대 72%의 정확도를 달성했습니다.

 

우리의 연구는 표준 ImageNet 데이터셋보다 더 큰 규모에서 이미지 인식을 탐구하는 논문들의 증가에 기여하고 있습니다. 추가적인 데이터 소스를 사용함으로써 표준 벤치마크에서 최신 성과를 달성할 수 있습니다(Mahajan et al., 2018; Touvron et al., 2019; Xie et al., 2020). 또한 Sun et al. (2017)은 데이터셋 크기에 따른 CNN 성능 확장성을 연구했으며, Kolesnikov et al. (2020)과 Djolonga et al. (2020)은 ImageNet-21k 및 JFT-300M과 같은 대규모 데이터셋으로부터의 CNN 전이 학습을 실증적으로 탐구했습니다. 우리는 이러한 두 데이터셋에 중점을 두고 있지만, 이전 연구에서 사용된 ResNet 기반 모델 대신 Transformer를 훈련하고 있습니다.

 

3. 방법

그림 1: 모델 개요
우리는 이미지를 고정된 크기의 패치로 나누고, 각 패치를 선형 임베딩한 후, 위치 임베딩을 추가하여 결과로 나온 벡터 시퀀스를 표준 Transformer 인코더에 입력합니다. 분류 작업을 수행하기 위해, 시퀀스에 학습 가능한 "분류 토큰"을 추가하는 표준 방식을 사용합니다. Transformer 인코더의 그림은 Vaswani et al. (2017)에서 영감을 받았습니다.

 

모델 설계에서는 가능한 한 원래의 Transformer(Vaswani et al., 2017)를 따릅니다. 이 의도적으로 간단한 설정의 장점은 확장 가능한 NLP Transformer 아키텍처와 그 효율적인 구현을 거의 그대로 사용할 수 있다는 점입니다.

 

3.1 Vision Transformer (ViT)
모델의 개요는 그림 1에 나와 있습니다. 표준 Transformer는 1D 토큰 임베딩 시퀀스를 입력으로 받습니다. 2D 이미지를 처리하기 위해, 이미지를 𝐱 ∈ ℝ H × W × C 형식에서 2D 패치로 변환한 시퀀스 𝐱 p ∈ ℝ N × ( P 2 ⋅ C ) 로 재구성합니다. 여기서 ( H , W ) 는 원본 이미지의 해상도, C 는 채널 수, ( P , P ) 는 각 이미지 패치의 해상도이며, N = H​W / P^2 는 결과적으로 패치의 수로, Transformer의 효과적인 입력 시퀀스 길이가 됩니다. Transformer는 모든 계층에서 일정한 잠재 벡터 크기 D 를 사용하므로, 패치를 평탄화한 후 학습 가능한 선형 투사(Eq. 1)를 통해 D 차원으로 매핑합니다. 이 투사 결과를 패치 임베딩이라고 부릅니다.

BERT의 [class] 토큰과 유사하게, 학습 가능한 임베딩을 임베딩된 패치 시퀀스 앞에 추가합니다( 𝐳^0_0 = 𝐱 class ), 이 토큰의 상태는 Transformer 인코더의 출력에서 이미지 표현 𝐲 (Eq. 4)로 사용됩니다. 사전 학습 및 미세 조정 중에, 분류 헤드가 𝐳 ^L_0 에 연결됩니다. 분류 헤드는 사전 학습 시에는 하나의 은닉 계층을 가진 MLP로 구현되고, 미세 조정 시에는 단일 선형 계층으로 구현됩니다.

위치 정보를 유지하기 위해 패치 임베딩에 위치 임베딩을 추가합니다. 우리는 표준 학습 가능한 1D 위치 임베딩을 사용하며, 더 진보된 2D 인식 위치 임베딩을 사용해도 성능 개선이 크게 없었습니다(부록 D.4). 임베딩 벡터 시퀀스는 인코더의 입력으로 사용됩니다.

Transformer 인코더(Vaswani et al., 2017)는 교차적으로 다중 헤드 자기 주의(MSA, 부록 A 참조)와 MLP 블록(Eq. 2, 3)으로 구성됩니다. 각 블록 앞에는 Layernorm(LN)이 적용되고, 각 블록 후에는 잔차 연결이 추가됩니다. (Wang et al., 2019; Baevski & Auli, 2019). MLP는 GELU 비선형성을 가진 두 개의 계층으로 이루어져 있습니다.

 

귀납적 편향
ViT는 CNN보다 이미지에 특화된 귀납적 편향이 훨씬 적다는 점을 주목할 필요가 있습니다. CNN에서는 각 계층 전체에 걸쳐 지역성, 2차원 이웃 구조, 번역 불변성이 내재되어 있습니다. 반면, ViT에서는 MLP 계층만이 지역적이고 번역 불변성이 있으며, 자기 주의 계층은 전역적입니다. 2차원 이웃 구조는 모델 초기에 이미지를 패치로 나누는 것과 미세 조정 시에 다른 해상도의 이미지에 대해 위치 임베딩을 조정할 때만 아주 제한적으로 사용됩니다. 초기화 시의 위치 임베딩은 패치의 2D 위치에 대한 정보를 전혀 포함하지 않으며, 패치 간의 모든 공간 관계는 처음부터 학습해야 합니다.

 

하이브리드 아키텍처
원시 이미지 패치 대신, 입력 시퀀스를 CNN의 특징 맵으로부터 생성할 수 있습니다(LeCun et al., 1989). 이 하이브리드 모델에서는 패치 임베딩 투사 𝐄 (Eq. 1)을 CNN 특징 맵에서 추출한 패치에 적용합니다. 특별한 경우로, 패치가 1x1 크기일 수 있으며, 이는 특징 맵의 공간 차원을 평탄화한 후 Transformer 차원으로 투사하여 입력 시퀀스를 얻는 것을 의미합니다. 분류 입력 임베딩과 위치 임베딩은 위에서 설명한 대로 추가됩니다.

 

3.2 미세 조정과 고해상도

일반적으로 우리는 ViT를 대규모 데이터셋에서 사전 학습한 후, 더 작은 하위 작업에 대해 미세 조정을 수행합니다. 이를 위해, 사전 학습된 예측 헤드를 제거하고, 크기 D × K 인 feedforward 레이어를 추가합니다. 여기서 K 는 하위 작업의 클래스 수입니다. 사전 학습 때보다 더 높은 해상도로 미세 조정하는 것이 유리할 때가 많습니다(Touvron et al., 2019; Kolesnikov et al., 2020). 더 높은 해상도의 이미지를 입력할 때, 패치 크기는 그대로 유지하므로, 더 긴 시퀀스 길이가 생성됩니다. Vision Transformer는 메모리 제약 내에서 임의의 시퀀스 길이를 처리할 수 있지만, 사전 학습된 위치 임베딩은 더 이상 유효하지 않을 수 있습니다. 따라서 원본 이미지에서의 위치에 따라 사전 학습된 위치 임베딩을 2D 보간법을 사용해 조정합니다. 이 해상도 조정과 패치 추출은 ViT에 2D 구조에 대한 귀납적 편향이 수동으로 적용되는 유일한 지점입니다.

 

4 실험

우리는 ResNet, Vision Transformer(ViT), 그리고 하이브리드 모델의 표현 학습 능력을 평가합니다. 각 모델의 데이터 요구 사항을 이해하기 위해, 다양한 크기의 데이터셋에 대해 사전 학습하고 여러 벤치마크 작업에서 평가를 진행합니다. 모델 사전 학습에 필요한 계산 비용을 고려했을 때, ViT는 사전 학습 비용이 낮으면서도 대부분의 인식 벤치마크에서 최신 성능을 달성합니다. 마지막으로, 소규모 실험을 통해 자가 지도 학습을 사용해 ViT가 미래에 유망할 가능성을 보여줍니다.

 

4.1 설정

데이터셋. 모델의 확장성을 탐구하기 위해 ILSVRC-2012 ImageNet 데이터셋(1k 클래스, 130만 이미지)을 사용하며, 이를 ImageNet이라 부릅니다. 또한 21k 클래스와 1400만 이미지를 포함하는 ImageNet-21k 및 JFT(Sun et al., 2017) 데이터셋(18k 클래스, 3억 300만 고해상도 이미지)도 사용합니다. Kolesnikov et al.(2020)의 방식에 따라, 하위 작업의 테스트 세트에 대해 사전 학습 데이터셋을 중복 제거합니다. 이 모델들은 여러 벤치마크 작업에 전이됩니다. 여기에는 원본 검증 레이블과 정제된 ReaL 레이블(Beyer et al., 2020)을 사용하는 ImageNet, CIFAR-10/100(Krizhevsky, 2009), Oxford-IIIT Pets(Parkhi et al., 2012), Oxford Flowers-102(Nilsback & Zisserman, 2008)이 포함됩니다. 이 데이터셋들에 대한 전처리는 Kolesnikov et al.(2020)의 방식을 따릅니다.

 

우리는 또한 19개 작업으로 이루어진 VTAB 분류 모음(Zhai et al., 2019b)에서 평가를 진행합니다. VTAB는 다양한 작업에 대해 1000개의 학습 예제를 사용하여 저데이터 전이를 평가합니다. 작업들은 세 그룹으로 나뉩니다: 자연 – 위에서 언급한 Pets, CIFAR 등과 같은 작업, 특수화 – 의료 및 위성 이미지, 구조화된 – 위치 추정과 같은 기하학적 이해를 요구하는 작업.

 

모델 변형. 우리는 ViT 구성을 BERT(Devlin et al., 2019)에 사용된 구성을 기반으로 하며, 이는 표 1에 요약되어 있습니다. "Base"와 "Large" 모델은 BERT에서 직접 가져왔으며, 더 큰 "Huge" 모델을 추가합니다. 여기서 우리는 모델 크기와 입력 패치 크기를 간단히 나타내기 위해 짧은 표기법을 사용합니다. 예를 들어, ViT-L/16은 "Large" 변형이며 입력 패치 크기가 16×16임을 의미합니다. Transformer의 시퀀스 길이는 패치 크기의 제곱에 반비례하므로, 패치 크기가 작은 모델일수록 계산 비용이 더 많이 듭니다.

 

기본 CNN 모델로는 ResNet(He et al., 2016)을 사용하지만, Batch Normalization(Ioffe & Szegedy, 2015) 레이어를 Group Normalization(Wu & He, 2018)으로 교체하고 표준화된 합성곱(Qiao et al., 2019)을 사용했습니다. 이러한 수정은 전이 학습 성능을 개선하며(Kolesnikov et al., 2020), 수정된 모델을 "ResNet (BiT)"라고 명명합니다. 하이브리드 모델의 경우, 중간 특징 맵을 패치 크기가 1 픽셀인 ViT에 입력합니다. 다른 시퀀스 길이를 실험하기 위해 (i) 일반 ResNet50의 4단계 출력을 가져오거나, (ii) 4단계를 제거하고 동일한 수의 계층을 3단계에 배치하여 3단계를 확장합니다. 옵션 (ii)는 4배 더 긴 시퀀스 길이를 제공하며, 더 많은 계산 비용이 소요되는 ViT 모델을 만듭니다.

 

표 1: Vision Transformer 모델 변형의 세부 사항.

 

훈련 및 미세 조정.
우리는 ResNet을 포함한 모든 모델을 Adam(Kingma & Ba, 2015)을 사용해 훈련했습니다. 이때, β ₁ = 0.9 , β ₂ = 0.999 , 배치 크기는 4096으로 설정했으며, 모든 모델의 전이 학습에 유용한 높은 가중치 감쇠 0.1 을 적용했습니다(부록 D.1에서는 일반적인 관행과 달리, 우리의 설정에서는 ResNet에 Adam이 SGD보다 약간 더 잘 작동함을 보여줍니다). 우리는 선형 학습률 워밍업 및 감쇠를 사용했으며, 자세한 내용은 부록 B.1을 참조하십시오. 미세 조정을 위해서는 모든 모델에 대해 모멘텀을 적용한 SGD와 배치 크기 512를 사용했습니다(부록 B.1.1 참조). 표 2의 ImageNet 결과를 위해 ViT-L/16은 해상도 512 , ViT-H/14는 518 에서 미세 조정을 수행했으며, Polyak & Juditsky(1992)의 평균화 기법을 0.9999 의 계수와 함께 사용했습니다(Ramachandran et al., 2019; Wang et al., 2020b).

 

평가지표.
우리는 하위 데이터셋에 대해 few-shot 학습 또는 미세 조정 정확도를 기준으로 결과를 보고합니다. 미세 조정 정확도는 각 모델을 해당 데이터셋에서 미세 조정한 후의 성능을 나타냅니다. Few-shot 정확도는 훈련 이미지의 일부 집합의 (고정된) 표현을 { − 1 , 1 } K 타겟 벡터로 매핑하는 정규화된 최소 제곱 회귀 문제를 해결하여 얻습니다. 이 공식은 폐쇄형 솔루션을 정확하게 도출할 수 있게 해줍니다. 우리는 주로 미세 조정 성능에 중점을 두지만, 미세 조정이 비용이 많이 드는 경우에는 빠른 평가를 위해 종종 선형 few-shot 정확도를 사용합니다.

 

4.2 최신 기술과의 비교

우리는 먼저 가장 큰 모델인 ViT-H/14와 ViT-L/16을 문헌에 보고된 최신 CNN들과 비교합니다. 첫 번째 비교 지점은 대형 ResNet을 사용한 지도 전이 학습을 수행하는 Big Transfer(BiT)(Kolesnikov et al., 2020)입니다. 두 번째는 Noisy Student(Xie et al., 2020)로, 이는 ImageNet과 라벨이 제거된 JFT-300M에서 반지도 학습을 사용해 훈련된 대형 EfficientNet입니다. 현재 Noisy Student는 ImageNet에서 최신 성능을 보유하고 있으며, BiT-L은 여기서 보고된 다른 데이터셋에서 최신 성능을 보이고 있습니다. 모든 모델은 TPUv3 하드웨어에서 훈련되었으며, 우리는 각 모델의 사전 학습에 걸린 TPUv3 코어-일을 보고합니다. 이는 TPU v3 코어(칩당 2개)를 사용한 훈련 시간(일)을 곱한 값입니다.

 

----------------------------------------------------------------

개인적인 이야기지만, kaggle 대회에서 noisy student는 생각보다 좋지 않았다. (사용하긴 함.)

----------------------------------------------------------------

 

표 2: 인기 있는 이미지 분류 벤치마크에서 최신 기술과의 비교
우리는 세 번의 미세 조정 실행을 통해 평균 및 표준 편차로 정확도를 보고합니다. JFT-300M 데이터셋에서 사전 학습된 Vision Transformer(ViT) 모델은 모든 데이터셋에서 ResNet 기반의 기준 모델을 능가하며, 사전 학습에 훨씬 적은 계산 자원을 사용했습니다. 더 작은 공개 ImageNet-21k 데이터셋에서 사전 학습된 ViT도 우수한 성능을 보입니다.
* Touvron et al. (2020)에서 약간 개선된 88.5% 결과가 보고되었습니다.

 

그림 2: VTAB 성능을 자연, 특수화, 구조화된 작업 그룹으로 나눈 결과

표 2는 결과를 보여줍니다. JFT-300M에서 사전 학습된 더 작은 ViT-L/16 모델은 동일한 데이터셋에서 사전 학습된 BiT-L을 모든 작업에서 능가하며, 훈련에 훨씬 적은 계산 자원을 사용했습니다. 더 큰 모델인 ViT-H/14는 특히 더 어려운 데이터셋인 ImageNet, CIFAR-100, 그리고 VTAB 모음에서 성능을 더욱 개선했습니다. 흥미롭게도, 이 모델은 이전 최신 기술보다 사전 학습에 필요한 계산 자원이 여전히 크게 적었습니다. 그러나 사전 학습 효율성은 아키텍처 선택뿐만 아니라 훈련 일정, 최적화기, 가중치 감쇠 등 다른 매개 변수에도 영향을 받을 수 있음을 주목할 필요가 있습니다. 다양한 아키텍처에 대한 성능 대 계산 자원에 대한 통제된 연구는 4.4절에서 제공합니다. 마지막으로, 공개된 ImageNet-21k 데이터셋에서 사전 학습된 ViT-L/16 모델도 대부분의 데이터셋에서 좋은 성능을 보였으며, 사전 학습에 더 적은 자원을 사용했습니다. 이 모델은 약 30일 동안 8개의 TPUv3 코어로 표준 클라우드 TPU에서 훈련할 수 있었습니다.

 

그림 2는 VTAB 작업을 각각의 그룹으로 나누고, 이 벤치마크에서 이전 최신 기술과 비교한 결과를 보여줍니다: BiT, VIVI – ImageNet과 YouTube에서 공동 학습된 ResNet(Tschannen et al., 2020), 그리고 S4L – ImageNet에서 지도 및 반지도 학습(Zhai et al., 2019a). ViT-H/14는 자연 및 구조화된 작업에서 BiT-R152x4와 다른 방법들을 능가했습니다. 특수화된 작업에서는 상위 두 모델의 성능이 유사했습니다.

 

4.3 사전 학습 데이터 요구 사항

그림 3: ImageNet으로의 전이
작은 데이터셋에서 사전 학습했을 때, 큰 ViT 모델은 BiT ResNet(음영 처리된 영역)보다 성능이 낮지만, 더 큰 데이터셋에서 사전 학습했을 때는 더 뛰어납니다. 마찬가지로, 더 큰 ViT 변형들은 데이터셋이 커질수록 작은 모델들을 능가합니다.

 

그림 4: ImageNet에서 사전 학습 크기에 따른 선형 few-shot 평가
ResNet은 작은 사전 학습 데이터셋에서 더 나은 성능을 보이지만, ViT보다 더 빨리 성능이 한계에 도달합니다. ViT-b는 모든 은닉 차원이 절반으로 줄어든 ViT-B 모델입니다.

 

Vision Transformer(ViT)는 JFT-300M이라는 대규모 데이터셋에서 사전 학습했을 때 우수한 성능을 보입니다. ViT는 ResNet보다 비전에 대한 귀납적 편향이 적은데, 그렇다면 데이터셋의 크기는 얼마나 중요한 역할을 할까요? 이를 알아보기 위해 두 가지 실험을 수행했습니다.

 

첫 번째로, ViT 모델을 점점 더 큰 데이터셋에서 사전 학습했습니다: ImageNet, ImageNet-21k, JFT-300M. 작은 데이터셋에서 성능을 향상시키기 위해, 우리는 세 가지 기본적인 정규화 매개 변수(가중치 감쇠, 드롭아웃, 라벨 스무딩)를 최적화했습니다. 그림 4는 ImageNet에 미세 조정한 후의 결과를 보여줍니다(다른 데이터셋의 결과는 표 5에서 확인할 수 있습니다). 가장 작은 데이터셋인 ImageNet에서 사전 학습한 경우, ViT-Large 모델은 ViT-Base 모델보다 성능이 낮았습니다(적당한 정규화를 적용했음에도 불구하고). 그러나 ImageNet-21k에서 사전 학습한 경우, 두 모델의 성능이 비슷해졌고, JFT-300M에서는 더 큰 모델의 이점이 드러났습니다. 그림 4는 또한 다양한 크기의 BiT 모델이 성능을 나타내는 영역을 보여줍니다. BiT CNN은 ImageNet에서 ViT보다 성능이 뛰어나지만, 더 큰 데이터셋에서는 ViT가 성능을 뛰어넘습니다.

 

두 번째로, 우리는 9M, 30M, 90M의 무작위 하위 집합과 전체 JFT-300M 데이터셋에서 모델을 훈련했습니다. 작은 하위 집합에서는 추가 정규화를 수행하지 않았으며, 모든 설정에 동일한 하이퍼 파라미터를 사용했습니다. 이를 통해 정규화의 효과가 아닌, 모델의 고유한 성질을 평가했습니다. 그러나 우리는 조기 종료를 사용했으며, 훈련 중에 달성된 최상의 검증 정확도를 보고했습니다. 계산 자원을 절약하기 위해, 우리는 전체 미세 조정 정확도 대신 few-shot 선형 정확도를 보고했습니다. 그림 4는 그 결과를 보여줍니다. Vision Transformer는 더 작은 데이터셋에서 유사한 계산 비용을 가진 ResNet보다 과적합이 더 심합니다. 예를 들어, ViT-B/32는 ResNet50보다 약간 더 빠르지만, 9M 하위 집합에서는 훨씬 성능이 떨어지고, 90M 이상의 하위 집합에서는 더 나은 성능을 보였습니다. ResNet152x2와 ViT-L/16에서도 마찬가지였습니다. 이 결과는 합성곱 신경망의 귀납적 편향이 더 작은 데이터셋에서는 유용하지만, 더 큰 데이터셋에서는 데이터로부터 직접 패턴을 학습하는 것이 충분히, 그리고 때로는 더 유익하다는 직관을 강화합니다.

 

종합적으로, ImageNet의 few-shot 결과(그림 4)와 VTAB의 저데이터 결과(표 2)는 매우 적은 데이터로의 전이에 대해 희망적인 결과를 보여줍니다. ViT의 few-shot 특성에 대한 추가 분석은 미래 연구에서 흥미로운 방향이 될 것입니다.

 

그림 5: 다양한 아키텍처에 대한 사전 학습 계산 자원 대비 성능: Vision Transformer, ResNet, 하이브리드 모델
Vision Transformer는 일반적으로 동일한 계산 예산 내에서 ResNet보다 뛰어난 성능을 보입니다. 하이브리드 모델은 작은 모델 크기에서는 순수 Transformer보다 성능을 개선하지만, 큰 모델에서는 그 차이가 사라집니다.

 

4.4 확장 연구
우리는 JFT-300M 데이터셋에서 전이 성능을 평가하여 서로 다른 모델의 확장 가능성에 대한 통제된 연구를 수행했습니다. 이 설정에서는 데이터 크기가 모델 성능의 병목을 이루지 않으며, 각 모델의 성능 대비 사전 학습 비용을 평가합니다. 연구에 포함된 모델은 다음과 같습니다: 7개의 ResNet 모델(R50x1, R50x2, R101x1, R152x1, R152x2, 7 에포크 동안 사전 학습된 모델, 그리고 14 에포크 동안 사전 학습된 R152x2와 R200x3); 6개의 Vision Transformer 모델(ViT-B/32, B/16, L/32, L/16, 7 에포크 동안 사전 학습된 모델, 그리고 14 에포크 동안 사전 학습된 L/16과 H/14); 5개의 하이브리드 모델(R50+ViT-B/32, B/16, L/32, L/16, 7 에포크 동안 사전 학습된 모델, 그리고 14 에포크 동안 사전 학습된 R50+ViT-L/16). 하이브리드 모델의 경우, 모델 이름 끝의 숫자는 패치 크기가 아니라 ResNet 백본에서의 총 다운샘플링 비율을 나타냅니다.

 

그림 5는 사전 학습 계산 자원 대비 전이 성능을 보여줍니다(계산 비용에 대한 자세한 내용은 부록 D.5를 참조하십시오). 각 모델에 대한 자세한 결과는 부록의 표 6에서 확인할 수 있습니다. 몇 가지 패턴이 관찰됩니다. 첫째, Vision Transformer는 성능/계산 비용 트레이드오프에서 ResNet을 능가합니다. ViT는 동일한 성능을 달성하는 데 약 2 − 4 배 적은 계산 자원을 사용합니다(5개의 데이터셋 평균). 둘째, 하이브리드 모델은 작은 계산 예산에서는 ViT보다 약간 더 나은 성능을 보이지만, 큰 모델에서는 그 차이가 사라집니다. 이는 다소 놀라운 결과로, 합성곱 기반의 지역적 특징 처리가 ViT의 성능을 모델 크기에 관계없이 도울 것이라고 예상할 수 있기 때문입니다. 셋째, Vision Transformer는 시도된 범위 내에서 성능이 포화되지 않는 것으로 보이며, 이는 향후 더 큰 확장 연구를 촉진할 수 있는 동기를 제공합니다.

 

4.5 Vision Transformer 분석

그림 6: 출력 토큰에서 입력 공간으로의 주의(attention)의 대표적인 예시. 자세한 내용은 부록 D.7을 참조하십시오.

Vision Transformer가 이미지 데이터를 어떻게 처리하는지 이해하기 위해, 우리는 모델의 내부 표현을 분석합니다. Vision Transformer의 첫 번째 계층은 평탄화된 패치를 선형 투사하여 저차원 공간으로 변환합니다(Eq. 1). 그림 7(왼쪽)은 학습된 임베딩 필터의 주요 구성 요소들을 보여줍니다. 이 구성 요소들은 각 패치 내에서 세부 구조의 저차원 표현을 위한 그럴듯한 기저 함수로 보입니다.

 

투사 이후, 학습된 위치 임베딩이 패치 표현에 추가됩니다. 그림 7(중앙)은 모델이 이미지 내에서 거리를 위치 임베딩의 유사성으로 인코딩하는 방법을 학습했음을 보여줍니다. 즉, 가까운 패치일수록 더 유사한 위치 임베딩을 갖는 경향이 있습니다. 또한, 행-열 구조가 나타납니다. 같은 행/열에 있는 패치들은 유사한 임베딩을 가집니다. 마지막으로, 더 큰 그리드에서는 때때로 사인파 구조가 나타납니다(부록 D 참조). 위치 임베딩이 2D 이미지의 위상을 나타내도록 학습된다는 점은, 왜 2D 인식을 고려한 수작업 임베딩 변형들이 성능을 개선하지 못하는지 설명해줍니다(부록 D.4 참조).

 

그림 7:
왼쪽: ViT-L/32의 RGB 값에 대한 초기 선형 임베딩 필터.
중앙: ViT-L/32의 위치 임베딩 유사도. 각 타일은 해당 행과 열의 패치 위치 임베딩과 다른 모든 패치의 위치 임베딩 간의 코사인 유사도를 보여줍니다.
오른쪽: 헤드별 및 네트워크 깊이에 따른 주의 영역 크기. 각 점은 한 계층에서 16개 헤드 중 하나의 평균 주의 거리를 나타냅니다. 자세한 내용은 부록 D.7을 참조하십시오.

 

자기 주의(self-attention)는 ViT가 가장 낮은 계층에서도 이미지 전체의 정보를 통합할 수 있게 해줍니다. 우리는 네트워크가 이 기능을 어느 정도 활용하는지 조사했습니다. 특히, 주의 가중치를 기반으로 정보가 통합되는 이미지 공간 내 평균 거리를 계산했습니다(그림 7, 오른쪽). 이 "주의 거리"는 CNN의 수용 영역 크기와 유사합니다. 우리는 일부 헤드가 가장 낮은 계층에서도 이미 이미지 대부분에 주의를 기울이고 있음을 발견했습니다. 이는 전역적으로 정보를 통합할 수 있는 모델의 능력이 실제로 사용되고 있음을 보여줍니다. 다른 헤드는 일관되게 낮은 계층에서 작은 주의 거리를 유지합니다. 이러한 국소화된 주의는 Transformer 앞에 ResNet을 적용하는 하이브리드 모델에서는 덜 두드러집니다(그림 7, 오른쪽). 이는 CNN의 초기 합성곱 계층과 유사한 기능을 수행할 수 있음을 시사합니다. 또한, 네트워크 깊이가 깊어짐에 따라 주의 거리가 증가합니다. 전체적으로, 모델은 분류에 의미 있는 이미지 영역에 주의를 기울입니다(그림 6 참조).

 

4.6 자가 지도 학습
Transformer는 NLP 작업에서 인상적인 성능을 보여줍니다. 그러나 이들의 성공은 뛰어난 확장성뿐만 아니라 대규모 자가 지도 학습에 기인한 경우가 많습니다(Devlin et al., 2019; Radford et al., 2018). 우리는 BERT에서 사용된 마스킹 언어 모델링 작업을 모방하여 자가 지도 학습을 위한 마스킹 패치 예측에 대한 초기 탐색을 수행했습니다. 자가 지도 사전 학습을 통해, 더 작은 ViT-B/16 모델은 ImageNet에서 79.9%의 정확도를 달성했으며, 이는 처음부터 훈련했을 때보다 2% 향상된 결과이지만, 여전히 지도 학습보다 4% 낮습니다. 자세한 내용은 부록 B.1.2에 있습니다. 대조 학습(Chen et al., 2020b; He et al., 2020; Bachman et al., 2019; Hénaff et al., 2020)에 대한 추가 탐색은 향후 연구 과제로 남겨둡니다.

 

5. 결론
우리는 Transformer를 이미지 인식에 직접 적용하는 방법을 탐구했습니다. 컴퓨터 비전에서 자기 주의를 사용하는 기존 연구들과 달리, 우리는 초기 패치 추출 단계를 제외하고는 이미지에 특화된 귀납적 편향을 아키텍처에 도입하지 않았습니다. 대신, 이미지를 패치의 시퀀스로 해석하고 이를 NLP에서 사용되는 표준 Transformer 인코더로 처리했습니다. 이 간단하지만 확장 가능한 전략은 대규모 데이터셋에서 사전 학습과 결합되었을 때 놀라울 정도로 잘 작동했습니다. 따라서 Vision Transformer는 많은 이미지 분류 데이터셋에서 최신 성능과 동등하거나 이를 초과하며, 사전 학습 비용이 상대적으로 저렴합니다.

 

이러한 초기 결과는 고무적이지만, 여전히 많은 과제가 남아 있습니다. 하나는 ViT를 감지 및 분할과 같은 다른 컴퓨터 비전 작업에 적용하는 것입니다. 우리의 결과는 Carion et al.(2020)의 연구 결과와 함께 이 접근 방식의 가능성을 시사합니다. 또 다른 과제는 자가 지도 사전 학습 방법을 계속 탐색하는 것입니다. 초기 실험은 자가 지도 사전 학습의 개선을 보여주지만, 자가 지도 학습과 대규모 지도 학습 간에는 여전히 큰 차이가 있습니다. 마지막으로, ViT의 추가 확장은 성능 향상으로 이어질 가능성이 큽니다.

 

감사의 말
이 연구는 베를린, 취리히, 암스테르담에서 수행되었습니다. Google의 많은 동료들에게 감사드리며, 특히 인프라 및 코드의 오픈 소스 공개에 중요한 도움을 준 Andreas Steiner, 대규모 학습 인프라에 도움을 준 Joan Puigcerver와 Maxim Neumann, 그리고 유익한 토론을 제공해 준 Dmitry Lepikhin, Aravindh Mahendran, Daniel Keysers, Mario Lučić, Noam Shazeer, Ashish Vaswani, Colin Raffel에게 감사드립니다.