Uma Imagem Vale Mais que 16x16 Palavras: Transformers para Reconhecimento de Imagens em Escala

Guia de leitura

Problema e motivação

As arquiteturas convolucionais (por exemplo, ResNet) têm sido a base da visão computacional moderna há muito tempo, devido a vieses indutivos integrados como equivariância à translação e localidade bidimensional. Embora os Transformers tenham se tornado o padrão no processamento de linguagem natural (NLP) devido à sua eficiência computacional e capacidade de escalonamento massivo, tentativas anteriores de trazer a autoatenção para a visão normalmente acoplavam a atenção a backbones de CNN ou utilizavam padrões especializados de atenção local/axial que eram difíceis de escalar de maneira eficiente em aceleradores de hardware. Este artigo investiga se uma arquitetura Transformer pura, aplicada diretamente a sequências de patches de imagem com modificações mínimas, pode competir ou superar CNNs no estado da arte.

Ideias centrais

• Imagens como Sequências de Palavras: Uma imagem $\mathbf{x}\in {\mathbb{R}}^{H\times W\times C}$ é dividida em patches não sobrepostos de tamanho $P\times P$. Os patches achatados são tratados de forma idêntica a tokens de NLP.
• Viés Indutivo Mínimo: Ao contrário das CNNs, onde a localidade espacial e a equivariância à translação estão presentes em cada camada, o ViT utiliza autoatenção global em todos os patches. A estrutura 2D é referenciada apenas na extração dos patches e na interpolação dos embeddings de posição para o ajuste fino em resoluções mais altas.
• Escala Supera o Viés Indutivo: Sem um viés indutivo forte, o ViT apresenta desempenho inferior ao das ResNets em conjuntos de dados pequenos a médios (como o ImageNet padrão sem regularização pesada). No entanto, quando pré-treinado em grandes conjuntos de dados (ImageNet-21k com 14M de imagens ou JFT-300M com 303M de imagens), o ViT se iguala ou supera o desempenho das CNNs consumindo substancialmente menos FLOPs de treinamento.

Método

1. Embedding de Patches e Construção da Entrada: Dado o tamanho do patch $(P,P)$, uma imagem gera $N=HW/P^2$ patches. Cada patch achatado ${\mathbf{x}}_p^i\in {\mathbb{R}}^{P^2\cdot C}$ é projetado linearmente pela matriz $\mathbf{E}\in {\mathbb{R}}^{(P^2\cdot C)\times D}$ para a dimensão $D$. Um token de classificação aprendível ${\mathbf{x}}_{\text{class}}$ é concatenado no início, e embeddings de posição aprendíveis unidimensionais ${\mathbf{E}}_{pos}\in {\mathbb{R}}^{(N+1)\times D}$ são somados: $${\mathbf{z}}_0=[{\mathbf{x}}_{\text{class}};{\mathbf{x}}_p^1\mathbf{E};{\mathbf{x}}_p^2\mathbf{E};\dots ;{\mathbf{x}}_p^N\mathbf{E}]+{\mathbf{E}}_{pos}$$
2. Camadas do Codificador Transformer: A sequência passa por $L$ blocos idênticos alternando blocos de Multihead Self-Attention (MSA) e MLP, usando pre-LayerNorm (LN) e conexões residuais: $${\mathbf{z}}_{\ell }^{\prime }=\mathrm{MSA}(\mathrm{LN}({\mathbf{z}}_{\ell -1}))+{\mathbf{z}}_{\ell -1},\ell =1\dots L$$ $${\mathbf{z}}_{\ell }=\mathrm{MLP}(\mathrm{LN}({\mathbf{z}}_{\ell }^{\prime }))+{\mathbf{z}}_{\ell }^{\prime },\ell =1\dots L$$
3. Representação e Cabeça de Classificação: A representação é obtida a partir do primeiro token de saída: $$\mathbf{y}=\mathrm{LN}({\mathbf{z}}_L^0)$$ Um pequeno MLP (com uma camada oculta) atua como a cabeça de classificação durante o pré-treinamento, sendo substituído por uma única projeção linear ($D\times K$) no ajuste fino.
4. Ajuste Fino em Resoluções Mais Altas: Ao realizar o ajuste fino com resoluções de imagem maiores (por exemplo, $384\times 384$ ou $512\times 512$), o tamanho do patch $P$ permanece fixo, o que aumenta o comprimento da sequência $N$. Os embeddings de posição unidimensionais pré-treinados são ajustados por meio de interpolação bicúbica 2D de acordo com suas posições na grade espacial original.

Resultados

• Comparação com o Estado da Arte:
  • O ViT-H/14 pré-treinado no JFT-300M alcançou 88,55% no ImageNet, 90,72% no ImageNet ReaL, 94,55% no CIFAR-100 e 77,63% na suíte de 19 tarefas do VTAB, demandando 2,5k TPUv3-core-days.
  • O ViT-L/16 pré-treinado no JFT-300M alcançou 87,76% no ImageNet com apenas 0,68k TPUv3-core-days, superando o BiT-L (ResNet152x4, 87,54%, 9,9k TPUv3-core-days) e aproximando-se do Noisy Student EfficientNet-L2 (88,4/88,5%, 12,3k TPUv3-core-days).
  • Pré-treinamento em conjuntos de dados públicos: o ViT-L/16 pré-treinado no ImageNet-21k atingiu 85,30% no ImageNet.
• Compromisso entre Computação e Desempenho: Nos experimentos de escalonamento no JFT-300M, o ViT puro dominou as ResNets na fronteira de Pareto, exigindo de $2\times$ a $4\times$ menos computação para atingir acurácia equivalente.
• Representações Internas: As camadas iniciais contêm cabeças de atenção que atendem a toda a extensão da imagem lado a lado com cabeças que têm distâncias de atenção pequenas e localizadas. A distância de atenção geralmente aumenta de forma monotônica com a profundidade.

Limitações e questões em aberto

• Necessidade de Dados (Data Hunger): O ViT generaliza mal em pequenos conjuntos de dados sem uma quantidade expressiva de dados de pré-treinamento (ImageNet-21k ou JFT-300M) ou regularização pesada.
• Lacuna no Aprendizado Autossupervisionado: Experimentos iniciais com predição de patches mascarados atingiram 79,9% no ImageNet (ViT-B/16), mantendo uma lacuna de 4% de desempenho em relação ao pré-treinamento supervisionado.
• Aplicação em Predição Densa: O trabalho inicial foca exclusivamente em classificação; adaptar transformers de patches puros para tarefas espaciais densas, como detecção de objetos e segmentação semântica, permanece como um desafio de pesquisa em aberto.