Adaptación del Decodificador LLaMA al Transformer de Visión

Adapting LLaMA Decoder to Vision Transformer

April 10, 2024

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

cs.AI

Resumen

Este trabajo examina si los Transformers de solo decodificador, como LLaMA, originalmente diseñados para modelos de lenguaje extenso (LLMs), pueden adaptarse al campo de la visión por computadora. Primero "LLaMAficamos" un ViT estándar paso a paso para alinearlo con la arquitectura de LLaMA, y descubrimos que aplicar directamente una máscara causal a la autoatención provoca un colapso de atención, lo que resulta en el fallo del entrenamiento de la red. Proponemos reposicionar el token de clase detrás de los tokens de imagen mediante una técnica de token de clase pos-secuencia para superar este desafío, permitiendo que la autoatención causal capture eficientemente la información completa de la imagen. Además, desarrollamos una estrategia de máscara suave que introduce gradualmente una máscara causal en la autoatención al inicio del entrenamiento para facilitar el comportamiento de optimización. El modelo adaptado, denominado image LLaMA (iLLaMA), es similar en arquitectura a LLaMA y permite el aprendizaje supervisado directo. Su autoatención causal mejora la eficiencia computacional y aprende representaciones complejas al elevar los rangos de los mapas de atención. iLLaMA rivaliza en rendimiento con sus contrapartes de solo codificador, alcanzando un 75.1% de precisión top-1 en ImageNet con solo 5.7M de parámetros. Escalar el modelo a ~310M y preentrenarlo en ImageNet-21K mejora aún más la precisión al 86.0%. Experimentos extensos demuestran las propiedades confiables de iLLaMA: calibración, sesgo de forma-textura, compatibilidad con cuantización, segmentación en ADE20K y transferencia de aprendizaje en CIFAR. Esperamos que nuestro estudio pueda inspirar nuevas perspectivas en el diseño de modelos visuales en la ola de los LLMs. Los modelos preentrenados y los códigos están disponibles aquí.

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.