[논문 리뷰] Scaling Vision Transformers to 22 Billion Parameters

Mostafa Dehghani, Josip Djolonga외 40명|arXiv (Cornell University)|2023. 02. 10.

Multimodal Machine Learning Applications인용 수 118

한 줄 요약

본 논문은 ViT-22B를 제시합니다. 안정성과 효율성을 위한 아키텍처 변경을 갖춘 22B 매개변수 비전 트랜스포머인 ViT-22B를 제시하며, 분류, 제로샷, 조밀 예측, 비디오, 공정성/강건성 벤치마크 전반에서 최첨단 또는 경쟁력 있는 성능을 달성합니다.

ABSTRACT

The scaling of Transformers has driven breakthrough capabilities for language models. At present, the largest large language models (LLMs) contain upwards of 100B parameters. Vision Transformers (ViT) have introduced the same architecture to image and video modelling, but these have not yet been successfully scaled to nearly the same degree; the largest dense ViT contains 4B parameters (Chen et al., 2022). We present a recipe for highly efficient and stable training of a 22B-parameter ViT (ViT-22B) and perform a wide variety of experiments on the resulting model. When evaluated on downstream tasks (often with a lightweight linear model on frozen features), ViT-22B demonstrates increasing performance with scale. We further observe other interesting benefits of scale, including an improved tradeoff between fairness and performance, state-of-the-art alignment to human visual perception in terms of shape/texture bias, and improved robustness. ViT-22B demonstrates the potential for "LLM-like" scaling in vision, and provides key steps towards getting there.

연구 동기 및 목표

• 22B-매개변수 Vision Transformer(ViT-22B)를 가능하게 하는 확장 가능한 학습 기술을 입증합니다.
• 다양한 작업에서 ViT-22B를 평가합니다: 이미지 분류, 제로샷 전이, 조밀 예측 및 비디오.
• 모델 규모가 증가함에 따른 공정성, 강건성, 보정 및 인간 정합성 효과를 분석합니다.
• 대규모 ViT가 작은 백본으로의 증류를 통해 효과적인 교사 역할을 할 수 있음을 보여줍니다.

제안 방법

• 스케일링에서의 안정성과 학습 가속화를 위해 병렬 레이어, QK 정규화 및 바이어스 생략을 도입합니다.
• 처리량 극대화를 위해 TPUv4 위의 2D 메시 샤딩으로 비동기식 모델-병렬 접근법을 사용합니다(1.15k 토큰/초/코어).
• 큰 모델과 배치 크기에 맞게 모델 매개변수와 활성화를 샤딩하면서 계산과 통신을 겹치게 합니다.
• 이미지 4B인 JFT-파생 데이터셋에서 이미지당 256 토큰과 177k-스텝 일정으로 사전 학습합니다.
• 다양한 다운스트림 작업에서 선형 프로빙, 고정 이미지 튜닝 및 엔드-투-엔드 파인 튜닝을 사용하여 평가합니다.

실험 결과

연구 질문

• RQ1 아키텍처 변경이 포함된 ViT-스케일 학습이 ViT-22B에 대해 안정적이고 효율적인 학습을 낳을 수 있는가?
• RQ2 이전 ViT 및 LiT 기반 방법과 비교하여 표준 및 분포 외 이미지 분류 작업에서 ViT-22B의 성능은 어떠한가?
• RQ3 ViT의 확대가 제로샷, 전이 및 교차 도메인 강건성, 공정성 및 인간 정합성 지표를 향상시키는가?
• RQ4 고정된 백본으로 사용될 때 ViT-22B가 조밀 예측 및 비디오 작업에 강력한 특징 표현을 제공할 수 있는가?

주요 결과

• ViT-22B는 동결된 특징 추출기로서 강력한 ImageNet 성능(89.5% 정확도)과 매칭된 텍스트 타워를 갖춘 제로샷 ImageNet에서 85.9%를 달성합니다.
• ViT-22B를 ViT-B/16 및 ViT-L/16로 증류하면 해당 소형 모델들에 대해 최첨단 ImageNet 정확도(각각 88.6%, 89.6%)를 얻습니다.
• ObjectNet에서 모델 크기에 따라 제로샷 결과가 향상되며, 이 도전적인 데이터셋에서 ViT-22B의 새로운 SOTA를 설정합니다.
• ViT-22B는 인간에 대한 형태 편향 정렬(87% 형태 편향)을 개선하고 하위그룹 간의 공정성/강건성 트레이드오프 및 보정 지표에서 더 나은 성능을 보여줍니다.
• 조밀 예측 전이(ADE20k 소수샷) 및 단안 깊이 추정이 ViT-22B 특성으로부터 이점을 얻어 ViT-L 및 ViT-G 기준선보다 우수합니다.
• 고정된 ViT-22B 백본을 사용한 비디오 평가는 이전의 4B 매개변수 모델에 비해 경쟁력 있는 결과를 보여주며, 전체 파인튜닝의 이득 여지가 남아 있습니다.