[Paper Review] VisionLLaMA: A Unified LLaMA Backbone for Vision Tasks

Xiangxiang Chu, Jianlin Su+2 more|arXiv (Cornell University)|Mar 1, 2024

Advanced Image and Video Retrieval Techniques4 citations

TL;DR

VisionLLaMA introduces a unified vision transformer architecture inspired by the LLaMA language model, enabling efficient and effective performance across diverse vision tasks such as image classification, detection, segmentation, and diffusion-based generation. By adapting LLaMA's transformer design with a novel 2D rotary positional encoding (AS2DRoPE), it achieves faster convergence and state-of-the-art performance without architectural bells and whistles.

ABSTRACT

Large language models are built on top of a transformer-based architecture to process textual inputs. For example, the LLaMA stands out among many open-source implementations. Can the same transformer be used to process 2D images? In this paper, we answer this question by unveiling a LLaMA-like vision transformer in plain and pyramid forms, termed VisionLLaMA, which is tailored for this purpose. VisionLLaMA is a unified and generic modelling framework for solving most vision tasks. We extensively evaluate its effectiveness using typical pre-training paradigms in a good portion of downstream tasks of image perception and especially image generation. In many cases, VisionLLaMA have exhibited substantial gains over the previous state-of-the-art vision transformers. We believe that VisionLLaMA can serve as a strong new baseline model for vision generation and understanding. Our code is released at https://github.com/Meituan-AutoML/VisionLLaMA.

Motivation & Objective

• To bridge the architectural gap between vision and language models by adapting the LLaMA transformer architecture for 2D image inputs.
• To develop a generic, unified framework for solving multiple vision tasks using a single backbone, reducing modality-specific design complexity.
• To improve training efficiency and performance in vision tasks by leveraging LLaMA’s architectural strengths, such as RMSNorm, SwiGLU, and RoPE.
• To enable zero-shot generalization across diverse image resolutions through a novel positional encoding scheme (AS2DRoPE).

Proposed method

• Proposes VisionLLaMA, a vision transformer with a LLaMA-like architecture, including RMSNorm, SwiGLU, and rotary positional embeddings (RoPE), adapted for 2D image data.
• Introduces AS2DRoPE (Auto-scaled 2D RoPE), a learnable 2D rotary positional encoding that supports arbitrary input resolutions via interpolation scaling.
• Employs both plain and pyramid-based backbone designs to evaluate performance across different vision task paradigms (e.g., ViT vs. Swin Transformer-style).
• Applies VisionLLaMA to supervised and self-supervised pre-training, including MAE and contrastive learning, on ImageNet and LAION datasets.
• Integrates VisionLLaMA into diffusion models (DiT and SiT) for text-to-image generation, using latent diffusion with CLIP-encoders.
• Uses a unified training pipeline with LLaMA-style optimization: self-supervised pre-training followed by supervised fine-tuning with RLHF-like adaptation.

Experimental results

Research questions

• RQ1 Can the LLaMA transformer architecture be effectively adapted for 2D vision tasks, despite differences in modality structure (1D vs. 2D, plain vs. pyramid)?
• RQ2 How does AS2DRoPE compare to standard 2D sin-cos positional encoding and NTK-based RoPE in terms of generalization across resolutions and convergence speed?
• RQ3 Does VisionLLaMA outperform existing vision transformers (e.g., ViT, DeiT, Swin) in both supervised and self-supervised pre-training settings across diverse vision tasks?
• RQ4 To what extent can VisionLLaMA serve as a strong backbone for diffusion-based image generation, especially in terms of fidelity and training efficiency?
• RQ5 What is the theoretical and empirical basis for VisionLLaMA’s faster convergence and improved performance compared to standard ViT?

Key findings

• VisionLLaMA achieves a top-1 accuracy of 81.6% on ImageNet-1K using the pyramid small model, outperforming ViT-Base and DeiT3-Large under the same training setup.
• In image generation using DiT-LLaMA-XL, VisionLLaMA generates high-fidelity 256×256 images with a classifier-free guidance (CFG) ratio of 4.0, demonstrating strong generative capability.
• VisionLLaMA converges faster than ViT and DeiT3-Large, with SiT-LLaMA reaching better performance at 300k steps than ViT at 400k steps in diffusion training.
• The AS2DRoPE positional encoding enables zero-shot evaluation at arbitrary resolutions (e.g., 224×224, 384×384, 512×512), achieving 79.5% top-1 accuracy on 512×512 resolution without fine-tuning.
• In MAE pre-training, VisionLLaMA maintains lower training loss than ViT-Base throughout 800 epochs, indicating more stable and efficient optimization.
• Theoretical analysis shows that RoPE-based attention mechanisms in VisionLLaMA reduce content-position interaction compared to additive positional encoding, enabling faster and more effective learning.