Adapter le décodeur LLaMA au Vision Transformer

Adapting LLaMA Decoder to Vision Transformer

April 10, 2024

Auteurs: Jiahao Wang, Wenqi Shao, Mengzhao Chen, Chengyue Wu, Yong Liu, Kaipeng Zhang, Songyang Zhang, Kai Chen, Ping Luo

cs.AI

Résumé

Ce travail examine si les Transformers à décodeur uniquement, tels que LLaMA, initialement conçus pour les grands modèles de langage (LLMs), peuvent être adaptés au domaine de la vision par ordinateur. Nous commençons par "LLaMAfier" un ViT standard étape par étape pour l'aligner sur l'architecture de LLaMA, et constatons que l'application directe d'un masque causal à l'auto-attention entraîne un problème d'effondrement de l'attention, ce qui fait échouer l'entraînement du réseau. Nous proposons de repositionner le token de classe derrière les tokens d'image grâce à une technique de token de classe post-séquence pour surmonter ce défi, permettant à l'auto-attention causale de capturer efficacement l'ensemble des informations de l'image. De plus, nous développons une stratégie de masque doux qui introduit progressivement un masque causal à l'auto-attention au début de l'entraînement pour faciliter le comportement d'optimisation. Le modèle adapté, baptisé image LLaMA (iLLaMA), est similaire à LLaMA en architecture et permet un apprentissage supervisé direct. Son auto-attention causale améliore l'efficacité computationnelle et apprend des représentations complexes en augmentant les rangs des cartes d'attention. iLLaMA rivalise avec ses homologues à encodeur uniquement, atteignant une précision top-1 de 75,1% sur ImageNet avec seulement 5,7M de paramètres. Le passage à une échelle de ~310M de paramètres et un pré-entraînement sur ImageNet-21K améliore encore la précision à 86,0%. Des expériences approfondies démontrent les propriétés fiables d'iLLaMA : calibration, biais forme-texture, compatibilité avec la quantification, segmentation ADE20K et apprentissage par transfert sur CIFAR. Nous espérons que notre étude pourra inspirer de nouvelles perspectives dans la conception de modèles visuels à l'ère des LLMs. Les modèles pré-entraînés et les codes sont disponibles ici.

English

This work examines whether decoder-only Transformers such as LLaMA, which were originally designed for large language models (LLMs), can be adapted to the computer vision field. We first "LLaMAfy" a standard ViT step-by-step to align with LLaMA's architecture, and find that directly applying a casual mask to the self-attention brings an attention collapse issue, resulting in the failure to the network training. We suggest to reposition the class token behind the image tokens with a post-sequence class token technique to overcome this challenge, enabling causal self-attention to efficiently capture the entire image's information. Additionally, we develop a soft mask strategy that gradually introduces a casual mask to the self-attention at the onset of training to facilitate the optimization behavior. The tailored model, dubbed as image LLaMA (iLLaMA), is akin to LLaMA in architecture and enables direct supervised learning. Its causal self-attention boosts computational efficiency and learns complex representation by elevating attention map ranks. iLLaMA rivals the performance with its encoder-only counterparts, achieving 75.1% ImageNet top-1 accuracy with only 5.7M parameters. Scaling the model to ~310M and pre-training on ImageNet-21K further enhances the accuracy to 86.0%. Extensive experiments demonstrate iLLaMA's reliable properties: calibration, shape-texture bias, quantization compatibility, ADE20K segmentation and CIFAR transfer learning. We hope our study can kindle fresh views to visual model design in the wave of LLMs. Pre-trained models and codes are available here.