[Transformer] Patch n’ Pack: NaViT, a Vision Transformer for any Aspect Ratio and Resolution
Overview
얼마전 Google DeepMind에서 "Patch n’ Pack: NaViT, a Vision Transformer for any Aspect Ratio and Resolution" 라는 논문이 나왔습니다. 기존 컴퓨터 비전 모델에서는 이미지를 처리하기 전에 고정된 해상도로 이미지 크기를 설정하게 되고, ViT(Vision Transformer)와 같은 모델은 flexible sequence-based modeling을 하기 때문에 다양한 input sequence length를 제공하게 됩니다.
본 논문에서는 임의의 해상도와 종횡비(aspect ratio)를 처리하기 위해 학습 중에 sequence packing을 사용하는 NaViT(Native Resolution ViT)를 제안합니다. 모델을 유연하게 사용할 수 있고, large-scale supervised learning 및 contrastive image-text pretraining 학습을 효율적으로 가능하게 합니다.
NaViT는 image and video classification, object detection, semantic segmentation 과 같은 분야에 적용 할 수 있다고 합니다. inference time에 input resolution flexibility를 사용하여 test 시 cost-performance trade-off를 해결 할 수 있다고 하네요. 이는 CNN이 설계한 입력 및 모델링 파이프라인 표준에서 벗어나 ViT이 더 발전될 수 있는 방향을 제시한다고 합니다.
1. Introduction
ViT 모델은 단순하고 확장 가능한 특성이 있기 때문에 기존의 CNN을 대체할 수 있게 되었는데요, 이는 이미지를 패치로 분할하고 각 패치는 토큰에 선형적으로 반영이 됩니다. 이 때 일반적으로 이미지는 고정된 정사각형 종횡비로 크기가 조정된 다음 고정된 수의 패치로 분할되게 됩니다.
최근 FlexiViT(Beyer et al., 2023)에서는 하나의 아키텍처 내에서 다중 패치 크기를 지원하여 sequence length를 원활하게 변형할 수 있도록 하는데요, 이는 각 학습 단계에서 패치 크기의 random sampling과 initial convolution embedding이 multiple patch size를 지원 할 수 있도록 resizing algorithm에 의해 수행됩니다.
또한 Pix2Struct(Lee et al., 2022)에서는 차트나 문서 이해와 같은 작업에 대해 유용한 종횡비를 유지하는 alternative patching approach를 제안했습니다.
본 논문에서는 NaViT를 제안합니다. 서로 다른 이미지의 여러 패치가 하나의 시퀀스(Patch n' Pack)로 압축되어 종횡비를 유지하면서 다양한 해상도를 사용할 수 있게 합니다. 이는 자연어 처리 분야의 example packing에서 영감을 얻었다고 하네요. 가변 길이 입력에 대해 효율적으로 학습하기 위해 여러 example가 single sequence로 압축됩니다. 그래서 아래 3가지를 입증합니다.
1. 학습 시간에 해상도를 random sampling 하면 학습 비용이 크게 줄어듦
2. NaViT는 다양한 해상도에서 고성능을 제공, 추론 시 cost-performance trade-off 가능, 새로운 작업시 더 적은 비용으로 접근 가능
3. example packing으로 가능해진 fixed batch shape은 aspect-ratio preserving resolution-sampling, variable token dropping rate, adaptive computation 등의 새로운 아이디어로 발전 가능
이러한 fixed computational budget은 ViT 성능을 능가할 수 있습니다. 예를 들면 NaViT가 ViT의 성능 까지 비슷한 수준으로 학습하는데에는 4배 더 적은 컴퓨팅으로 가능하게 합니다. 대단하지 않나요?......
또한 pre-training 및 fine-tuning 중에 NaViT를 여러 해상도에 노출 시키게 되므로 single model이 다양한 해상도에서 평가할 때 우수한 성능을 보여 추론 비용 측면에서 NaViT를 사용하는게 더 유리해집니다.
2. Method
학습 할 때 신경망은 일반적으로 input의 batch로 학습되고, 여지껏 컨볼루션 신경망과 관련된 구조적 한계 때문에 이미지 크기를 조정하거나 고정된 크기로 패딩하는 관행이 있습니다. 본 논문에서 이미지 크기를 조정하는 것은 성능에 해를 끼치고, 고정된 크기로 패딩하는 것은 비효율적이라고 지적합니다. 아래 그림을 보시면 각 ImageNet, LVIS, WebLI 데이터세트에서 종횡비가 정사각형이 아닙니다.
언어 모델링에서는 example packing을 통해 fixed sequence length의 제한을 우회하는 것이 일반적인데, 여러 개별 example의 token이 하나의 sequence로 결합되어 언어 모델의 학습을 크게 가속화 할 수 있습니다. 이미지 분야에서도 언어 모델링과 마찬가지로 이미지를 일련의 패치(token)으로 취급함으로써 ViT 구조도 이러한 Patch n’ Pack의 패러다임의 이점을 누릴 수 있습니다. 그래서 본 논문에서는 이러한 기술을 사용하여 Native resolution으로 이미지에 대해 학습하는 방법인 NaViT를 제안합니다.
2.1 Architectural changes
NaViT는 원래 ViT를 기반으로 구축되었지만, 원칙적으로 일련의 패치에서 작동하는 모든 ViT의 변형을 사용할 수 있습니다. Patch n' Pack을 활성화 하기 위해 다음과 같이 아키텍처를 수정합니다.
Masked self attention and masked pooling
example이 서로 어긋나는 것을 방지하기 위해 additional self-attention mask가 도입됩니다. 마찬가지로 encoder 상단의 maskted pooling은 각 example 내에서 token representation을 pooling 하여 sequence의 example 당 single vector representation을 생성하는 것을 목표로 합니다. 그림 2에서 masking을 통해 attention receptive filed가 제어되는 방식을 보여줍니다.
Factorized & fractional positional embeddings
임의의 해상도와 종횡비를 처리하기 위해 postion embedding을 살펴보면, 해상도 RxR의 정사각형 이미지가 주어지면 패치 크기가 P인 vanilla ViT는 길이 (R/P)^2 의 1D positinal embedding을 학습합니다. 이러한 embedding을 linearly interpolation 하는 것은 더 높은 해상도 R에서 학습하거나 평가하는데 필요합니다. Pix2struct(Lee et al., 2022)에서는 학습된 2D absolute positional embedding을 도입하여 [maxLen, manLen] 크기의 positional embedding을 학습하고, 각 패치의 좌표 (x, y)로 인덱싱 합니다. 이를 통해 최대 R = P · maxLen 해상도로 다양한 종횡비를 사용할 수 있게 됩니다. 하지만 학습 중에 (x, y) 좌표의 모든 조합을 확인해야한다고 하네요.
가변 종횡비를 지원하고, unseen resolution를 위해 factorized positional embedding을 도입하여, x, y 좌표의 separate embedding ϕx and ϕy로 분해합니다. 그 다음 이들은 다시 합산된다고 하네요. 여기서 2가지 스키마를 고려합니다.
1. absolute embedding (여기서 ϕ(p) : [0, maxLen] → R^D 는 absolute patch index의 함수)
2. fractional embedding (여기서 ϕ(r) : [0, 1] → R^D는 r의 함수, r = p/side-length, 이미지에 따른 Relative distance)
후자는 이미지 크기와 무관한 positional embedding parameter를 제공하지만, 원래 종횡비를 부분적으로 난독화(obfuscates)하여 패치 수에 암시적으로만 나타낸다고 하네요.
본 논문에서는 simple learned embeddingϕ, sinusoidal embeddings, earned Fourier positional embedding used by NeRF을 고려한다고 합니다.
2.2 Training changes
Patch n' pack을 사용하면 NaViT 학습 중에 새로운 기술을 사용할 수 있다고 합니다.
Continuous Token dropping
Token dropping(학습 중 입력 패치의 random omission(생략)) 기법은 학습을 가속화 시키기 위해 개발되었다고 합니다. 그러나 일반적으로 모든 example에서 동일한 비율의 토큰이 삭제됩니다. packing은 지속적인 token drop을 가능하게 하므로 token drop rate는 이미지마다 다를 수 있다고 합니다. 이를 통해 이미지를 계속 보며 drop 시키기 때문에 train/inference의 불일치를 줄임으로써 더 빠른 처리량의 이점을 얻을 수 있다고 하네요. 그리고 packing을 사용하면 미리 정의된 schedule에 따라 drop distribution이 학습 전반에 걸쳐 달라질 수 있다고 합니다.
Resolution sampling
NaViT는 각 이미지의 원래 해상도를 사용하여 학습할 수 있습니다. 또는 종횡비를 유지하면서 총 픽셀 수를 다시 샘플링 할 수 있다고 하네요. vanilla ViT에서는 더 큰 처리량(작은 이미지에 대한 학습)과 더 높은 성능(평가 시 고해상도를 지원하기 위해 더 큰 이미지에 대한 학습)사이에 tension이 있는데, 종종 모델은 더 작은 해상도에서 pre-training 되고 더 높은 해상도에서 fine-tuning 된다고 하네요.
NaViT는 훨씬 더 유연하고, 각 이미지의 종횡비를 유지하면서 이미지 크기 분포에서 샘플리앟여 mixed-resolution training을 가능하게 합니다. 이를 통해 더 높은 처리량과 큰 이미지가 더 노출 되게끔 학습하는게 가능해져서 동등한 ViT에 대해 성능이 크게 향상됩니다.
2.3 Efficiency of NaViT
Patch n’ Pack 방식의 효율성에 대해 알아봅시다.
Self attention cost
attention의 O(n^2) cost는 여러 이미지를 더 긴 시퀀스로 packing 할 때의 자연스러운 문제인데, 많은 연구에서 2차 scaling을 제거하는 것을 목표로 하지만 trasnformer hidden dimension이 scaling 됨에 따라 MLP 포함 전체 cost에서 attention이 차지하는 비율이 더 작아진다고 하네요. 아래 그림은 model dimension이 늘어남에따라 packing을 더 많이 할 수록 overhead에 효과적이라는 것을 보여줍니다.
Packing, and sequence-level padding
여러 example을 포함하는 final sequence length는 고정되어야 합니다. 본 논문에서는 부록 A.3에 있는 greedy packing approach를 사용한다고 합니다. 일반적으로 fixed length에 정확히 추가되는 example의 완벽한 조합은 없기 때문에 padding token을 사용해야합니다. 예를 들어 나머지 token에 정확히 맞도록 sequence에서 마지막 example의 해상도 또는 token의 삭제 비율을 동적으로 선택할 수 있다고 합니다. 그러나 일반적으로 token의 2% 미만이 padding token이기 때문에 간단하게 해결해도 된다고 하네요.
Padding examples and the contrastive loss
token 당 loss는 packing 된 sequence로 구현하기 쉽습니다. 하지만 많은 비전 모델에서 일반적으로 pooled representation에 적용되는 example-level loss로 학습됩니다. 먼저 packing을 고려하여 일반적인 pooling head를 수정해야하고, 그 다음 각 sequence에서 multiple pooling representation을 추출해야합니다. Fixed batch shape은 B sequence batch에서 최대 B x E_max를 추출한다는 가정이 필요합니다. 즉 항상 배치 사이즈에 맞게 시퀀스를 추출해야합니다. sequence에 E_max 이미지 보다 많은 이미지가 포함되어있으면 추가적인 이미지들이 삭제되어 cost가 낭비되는데, 반대로 적은 경우에는 fake padding representation을 처리하게 됩니다.
이는 contrastive learning의 문제이며, 이를 방지하기 위해 chunked contrastive loss를 사용했다고 하네요. local device subset에 대한 계산을 수행하고, global softmax normalization에 필요한 통계를 효율적으로 축적하여 softmax에 대한 모든 data point를 수집할 필요가 없다고 하네요. 이는 loss로 인한 bottleneck 없이 E_max의 높은 값을 가능하게 합니다.
3. Experiments
NaViT에 사용하는 기본 아키텍처는 vanilla ViT를 따르며, JFT-4B (Zhai et al., 2022) 및 contrastive language-image training (Radford et al., 2021) on WebLI (Chen et al., 2022c)에 나온 설정으로 NaViT를 pre-train 합니다. 일반적으로 JFT의 경우 inception crop이 pre-train에 적용되며, 두 경우 모두 이미지 크기가 정사각형으로 조정됩니다. 별도로 지정하지 않으면 모든 NaViT는 이러한 작업 없이 사전 학습되며 종횡비를 유지하게 됩니다. NaViT는 FLAX 라이브러리를 사용하여 JAX에서 구현되고, Scenic 내에 구축된다고 하네요.
Classification pretraining
JFT-4B(Zhai et al., 2022)에서 (Dehghani et al., 2023)의 설정에 따라 sigmoid cross-entropy loss을 사용하여 supervised classification 목표로 NaViT를 pre-train 합니다.
Contrastive pre-training
이미지 모델과 함께 contrastive image-text loss를 사용하여 학습됩니다. packing은 text sequence 가 고정된 길이로 채워질 필요가 없기 때문에 text-tower의 효율성을 향상시킵니다. 이는 일반적인 설정입니다. contrastive model은 zero-shot ImageNet classification 및 COCO image-text retrieval에서 평가됩니다.
NaViT는 아래와 같은 장점을 가집니다.
1. Improved training efficiency and performance
2. Benefits of variable resolution
3. Benefits of variable token dropping
아래 그림은 fixed computational cost에서 더 낮은 해상도를 샘플링하면 처리량이 증가하고, 성능이 향상되며 다양한 해상도에서 모델을 더 잘 사용할 수 있다는 것을 보여줍니다. 가변해상도와 고정해상도로 학습된 NaViT-B/16 모델은 mixed resolution의 이점을 보여줍니다.
ImageNet에서 다양한 해상도로 fine-tune 된 JFT B/16 model에서 모든 설정에서 성능이 전반적으로 능가하는 것을 볼 수 있습니다.
모든 설정(파란색, 빨간색)에서 전반적으로 NaViT는 ViT(주황색)를 능가합니다. 왼쪽에서는 variable resolution(빨간색)로 fine-tune된 단일 NaViT는 하나의 해상도(파란색)로만 조정된 모델만큼 우수합니다. 오른쪽에서는 Mixed-resolution pretraining은 낮은 해상도(파란색 곡선의 왼쪽 끝)에서 미세 조정할 때 높은 해상도에서 잘 수행됩니다.
아래 그림 7과 같이 더 낮은 해상도 쪽으로 편향된 side length를 직접 샘플링하면 고정된 computational budget에서 최고의 성능을 어등ㄹ 수 있다고 하네요.
아래 그림 8은 NaViT-B/16의 JFT pre-train 중에 token drop rate를 줄임으로써 추가 개선이 가능함을 보여줍니다.
sequence packing을 통한 continuous token dropping 전략은 성능을 향상 시키게 됩니다.
또한 아래 그림을 보시면 factorized 접근법이 baseline ViT와 Pix2struct 방법보다 좋다는 것을 볼 수 있습니다. Factorized embedding은 더하기 방식으로 결합하는 것이 가장 좋네요.
아래 사진은 NaViT가 ImageNet에서 성능이 더 좋음을 나타냅니다. 특히 ImageNet-A 에는 종횡비가 극단적이고, center crop 외부에 중요한 정보가 많은 이미지들이 있는데 여기서 더 나은 성능을 자랑합니다. 하지만 ObjectNet에서는 ViT가 성능을 따라잡게 됩니다. 두 데이터 세트 모두 종횡비가 극단적인 이미지를 포함하고 있는데도 불구하고 따라잡는다는 것은 흥미롭습니다. 이는 ImageNet-A 및 ObjectNet을 위해 이미지에 적용하는, 종횡비를 보존하기 위한 center crop 때문이라고 생각한다고 하네요. 이는 ObjectNet에는 유용하지만 ImageNet-A에는 유용하지 않은 것 같다고 합니다.
NaViT는 다른 task에도 적용이 가능합니다. 아래는 object detection 분야에 적용한 결과입니다.
4. Conclusions and future work
vision transformer에 sequence packing을 간단하게 적용한 Patch n’ Pack이 training efficiency을 크게 향상시킨다는 것을 입증했습니다. NaViT 모델은 추론 시 많은 해상도에 적용할 수 있으며 new task에 저렴하게 적용할 수 있습니다. Patch n' Pack은 train 및 inference efficiency을 개선하기 위한 adaptive computation 및 새로운 알고리즘을 포함하여 이전에는 fixed batch shape의 필요성으로 인해 방해를 받았던 다양한 연구를 가능하게 합니다.