VisionLLaMA: Unified Transformer for Vision

• VisionLLaMA is a family of pure-transformer backbones that repurpose LLaMA’s architecture using 2D rotary positional encodings (AS2DRoPE) to effectively manage visual data.
• The approach includes both plain (ViT-style) and pyramid (Twins-style) variants, enabling robust performance across image classification, segmentation, detection, and generative modeling tasks.
• Empirical evaluations reveal that VisionLLaMA consistently matches or outperforms established baselines like DeiT3, Swin, Twins, DiT, and SiT with enhanced efficiency and convergence speed.

VisionLLaMA is a family of pure-transformer model backbones for vision tasks designed to retain maximal architectural and codebase compatibility with the LLaMA LLM while extending transformer processing power to two-dimensional image domains. Introduced in "VisionLLaMA: A Unified LLaMA Backbone for Vision Tasks" (Chu et al., 1 Mar 2024), VisionLLaMA encompasses both plain (ViT-style) and pyramid (Twins-style) variants, enabling unified treatment of a diverse set of visual recognition and generative tasks. This approach applies the language-model transformer paradigm directly to vision, introducing key innovations such as 2D rotary positional encodings (AS2DRoPE) and efficient minimal architectural adaptations, yielding state-of-the-art (SOTA) or strong baseline performance across classification, self-supervised representation learning, semantic segmentation, dense detection, and diffusion-based generative modeling.

1. Architectural Adaptations for Vision

1.1 Plain VisionLLaMA (ViT-style Backbone)

VisionLLaMA processes an input image of size $H \times W$ by dividing it into non-overlapping $P \times P$ patches, each mapped via a trainable projection to a $d$-dimensional embedding. For an image, $N = HW/P^2$ and $X \in \mathbb{R}^{N \times (P^2 \cdot 3)}$ are the set of raw flattened patches. A learnable class token $x_{\mathrm{cls}}$ is prepended to form the embedding sequence:

$$E = [x_{\mathrm{cls}}; X W_e + b_e] \in \mathbb{R}^{(N+1) \times d}$$

where $W_e \in \mathbb{R}^{(P^2 \cdot 3) \times d}$ and $b_e \in \mathbb{R}^d$. Unlike classical ViT, VisionLLaMA does not utilize additive or absolute positional embeddings at the input. Instead, each transformer block applies 2D rotary position embeddings (AS2DRoPE) to the input at every layer.

A VisionLLaMA block operates:

$$\tilde{z}_{ij}^{\,\ell} = \mathrm{MHSA}(\mathrm{AS2DRoPE}(\mathrm{LN}(z_{ij}^{\,\ell-1}))) + z_{ij}^{\,\ell-1}, \ z_{ij}^{\,\ell} = \mathrm{SwiGLU}(\mathrm{LN}(\tilde{z}_{ij}^{\,\ell})) + \tilde{z}_{ij}^{\,\ell}$$

where $(i,j)$ indexes the 2D spatial position, LayerNorm (LN), Multi-Head Self-Attention (MHSA), and SwiGLU feed-forward blocks follow the LLaMA conventions, with all steps extended to accept 2D token arrays.

The AS2DRoPE rotational matrix for location $(i,j)$ is:

$$R_{i,j} = \mathrm{diag}\left(\{ \mathsf{Rot}_{\theta_m}(i) \}_{m=0,4,...,d-4},\{\mathsf{Rot}_{\theta_m}(j)\} \right)$$

with

$$\mathsf{Rot}_\theta(x) = \begin{pmatrix} \cos(x\theta) & -\sin(x\theta) \ \sin(x\theta) & \cos(x\theta) \end{pmatrix} ,\quad \theta_m = 10000^{-m/d}$$

This formulation enables AS2DRoPE to be shared across heads and interpolated at arbitrary spatial resolutions, with negligible computational overhead.

1.2 Pyramid VisionLLaMA (Twins-SVT-style Backbone)

For dense prediction tasks, VisionLLaMA employs a Twins-SVT inspired feature pyramid. The backbone comprises patch-embedding convolutional stages (stride $P_s$, channels $C_s$), interleaved blocks with local (LSA) and global (GSA) attention as:

$$\hat{z}^{\,\ell} = \mathrm{LSA}(\mathrm{AS2DRoPE}(\mathrm{LN}(z^{\,\ell-1}))) + z^{\,\ell-1} \ z^{\,\ell} = \mathrm{SwiGLU}(\mathrm{LN}(\hat{z}^{\,\ell}))+\hat{z}^{\,\ell}$$

GSA replaces LSA for long-range contextualization. The final feature maps are globally averaged pooled before application to task-specific heads.

2. Model Objectives and Training Paradigms

2.1 Discriminative Objectives

• Image Classification (ImageNet–1K): Utilizes categorical cross-entropy:

$$\mathcal{L}_{\rm CE} = -\sum_{c=1}^K y_c \log p_c, \qquad p = \mathrm{softmax}(W_{\rm head}\,z_{\rm cls})$$

where $z_{\rm cls}$ is the class token output.

• Masked Image Modeling (MAE-style):

$$\mathcal{L}_{\rm MIM} = \mathbb{E}_{x,M} \| x_M - \mathcal{D}(\mathcal{E}(x_{\bar M})) \|^2_2$$

with $M$ masking 75% of patches; encoder $\mathcal{E}$ and decoder $\mathcal{D}$ share only AS2DRoPE for positional input.

2.2 Generative Objectives

• Diffusion Models (DiT/SiT frameworks): Training is performed to predict additive Gaussian noise $\epsilon$:

$$\mathcal{L}_{\rm DDPM} = \mathbb{E}_{x,\epsilon,t} \| \epsilon - \epsilon_\theta(z_t, t) \|^2_2$$

where $z_t$ is the latent at diffusion timestep $t$. At sampling, classifier-free guidance is supported up to a guidance scale of 4.0.

2.3 Pre-training and Optimization

VisionLLaMA is pre-trained using both supervised and self-supervised paradigms on ImageNet-1K. MAE pre-training is conducted with AdamW ($\beta_1=0.9, \beta_2=0.95$), batch size 4096, learning rate $1.5\times10^{-4}$, 800 or 1600 epochs (with linear warm-up), and strong data augmentation. Supervised training uses LAMB or AdamW optimizers and extensive schedules. Generative diffusion (DiT/SiT) models use AdamW, batch size 256, over 400K steps with standard linear $\beta$ diffusion schedules.

3. Evaluation Across Vision Tasks

3.1 Image Classification

On ImageNet-1K, plain VisionLLaMA-L (310M params) achieves 84.6% top-1 accuracy in two-stage training, matching or exceeding DeiT3-Large (84.5%). Pyramid VisionLLaMA-B (56M params) obtains 83.2%, outperforming Swin-T (81.3%) and Twins-S (81.7%).

MAE-based self-supervised pretraining yields VisionLLaMA-Base 84.0% (800ep) and 84.3% (1600ep), surpassing ViT-Base-MAE (83.2%) and MaskFeat (84.0%). Linear-probe performance is improved (VisionLLaMA-Base 69.7% vs ViT-Base 65.1%).

3.2 Semantic Segmentation

Coupled with UPerNet, pyramid VisionLLaMA attains 49.1% mIoU (B) and 50.0% (L) on ADE20K, compared to Swin-S 47.6%/Swin-B 48.1% and Twins-B/L 47.7%/48.8%, both for supervised and self-supervised paradigms.

3.3 Object Detection

In COCO detection, Mask-RCNN with pyramid VisionLLaMA-B yields AP/mAP of 49.1/43.8, outperforming Swin-S (47.6/42.8) and Twins-B (48.0/43.0). For ViTDet-style detection (plain MAE), VisionLLaMA-Base achieves 52.2/46.3 AP with 36 epochs, surpassing ViT-Base-MAE’s 51.6/45.7 with 100 epochs.

3.4 Image Generation (Diffusion)

• DiT: VisionLLaMA-XL/2 achieves FID 9.8 (2352K steps) on 256$^2$ generation, improved to 2.4 under CFG guidance, outperforming all ViT/DiT baselines (e.g., DiT-XL/2 FID=10.7 (w/o guidance: 18.7 vs 22.5)).
• SiT: VisionLLaMA-L/2 scores FID 14.3 (Euler SDE), outperforming DiT/SiT competitive models.

A summary table for reference:

Task	VisionLLaMA Result	Baseline for Comparison
ImageNet-1K	84.6% (L, 310M, 2-stage)	DeiT3-L 84.5%
ADE20K mIoU	50.0% (L, Pyramid)	Swin-B 48.1%, Twins-L 48.8%
COCO AP (det)	49.1/43.8 (B, Pyramid)	Swin-S 47.6/42.8, Twins-B 48.0/43
DiT FID (XL/2,CFG)	2.4 (VisionLLaMA)	3.2 (DiT XL/2)

4. Model Complexity and Computational Efficiency

Plain VisionLLaMA models (S/B/L: 22M/86M/310M params) feature 4.6/15.4/30 GFLOPs per forward pass and, in the small variant, process 817 images/s. Pyramid VisionLLaMA (S/B/L: 24M/56M/99M) are computationally lighter (2.9/8.6/15.1 GFLOPs), supporting faster throughputs up to 1059 images/s (small). Generative diffusion models range from 5.6 to 118 GFLOPs with parameter counts from 130M to 675M.

A typical trade-off is that plain VisionLLaMA is more effective for global-image tasks (classification, generative modeling), while pyramid VisionLLaMA excels at dense spatial tasks (segmentation, detection) with improved FLOP efficiency.

5. Key Empirical Insights and Conclusions

VisionLLaMA yields systematic performance improvements or ties across a broad spectrum of image understanding and generation tasks when compared to strong tuned baselines, including DeiT3, Swin, Twins, DiT, and SiT. The introduction of AS2DRoPE as a sole means of positional encoding is critical: it is lightweight, interpolates gracefully to test-time resolutions, and fosters faster convergence as observed in both generative and discriminative pretraining.

Training curves reveal 2–5× acceleration in generative models' convergence (FID vs. step), and earlier accuracy gains for both self-supervised learning and supervised tasks. Qualitative evidence demonstrates that diffusion-based VisionLLaMA variants produce "crisp, semantically coherent" images, supporting the model’s capacity for high-fidelity visual synthesis (Chu et al., 1 Mar 2024).

By reusing LLaMA’s core transformer mechanisms and only minorly altering input embedding and block composition for visual input, VisionLLaMA establishes a unified transformer baseline adaptable to virtually all modern vision tasks, bridging the architecture-standardization gap between LLMs and scene-level visual modeling.

6. Relationship to Other LLaMA-based Vision Architectures

The principle underlying VisionLLaMA—using LLaMA’s transformer structure beyond text—is directly extended in works such as Vista-LLaMA, which apply LLaMA decoders to video-language understanding with specialized attention modifications (EDVT-Attention, sequential visual projection, etc.) (Ma et al., 2023). While VisionLLaMA focuses on pure-vision representation, Vista-LLaMA and similar efforts investigate cross-modal capabilities and hallucination mitigation, suggesting the extensibility of VisionLLaMA’s methodology across the modalities.

A plausible implication is that the architectural unity offered by VisionLLaMA, combined with precise control over positional information (e.g., AS2DRoPE), confers benefits not only for high-precision 2D vision tasks but also for cross-modal video-language modeling and large-scale generative tasks.