LLaMA Decoder As Vision Transformer

## Abstract

Abstract 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 causal 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 causal 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 ∼ similar-to \sim ∼ 310M and pre-training on ImageNet-21K further enhances the accuracy to 86.0%.
Extensive experiments demonstrate iLLaMA’s reliable properties: shape-texture bias, calibration, quantization compatibility, ADE20K segmentation and CIFAR transfer learning.
We hope our study can kindle fresh views to visual architectures in the wave of LLMs and inspire the development of unified multimodal models.
Pre-trained models and codes are available https://github.com/techmonsterwang/iLLaMA .

## 1 Introduction

The year 2024 saw the meteoric rise of large language models (LLMs), as well as the 4th anniversary of the Vision Transformer (ViT) . Born in 2020, ViT was influenced by the prevailing encoder-only text Transformers at the time, such as BERT , *etc*. Accordingly, ViT is allowed to borrow the encoder-only design, *i.e.*, self-attention do not use any causal mask.
As a result, advanced vision backbones ; ; and training paradigms ; have followed such convention by default.

Meanwhile, the development of text Transformers did not stand still. A series of LLMs with a decoder-only architecture (*e.g.*, LLaMA ), have sparked a new wave. Pre-trained decoder-only Transformers have demonstrated remarkable performance in diverse textual  and multimodal tasks ; ; ; .
In this context, designing unified architectures for language and vision is a promising direction. Specifically, unified models employ shared types of operators to process both visual and textual data, reducing the cost of hardware implementation. Moreover, compared to using separate models for image and text, unified models ; simplify the inference process by handling different modalities simultaneously, thereby improving practical efficiency.

Toward this goal, we took the initial step of exploring whether decoder-only Transformers can hold its effectiveness in the unimodal vision domain. In this study, we demonstrate that through straightforward supervised learning, LLaMA architecture itself can process input images with simple yet crucial modifications. We start by modifying a standard encoder-only ViT (*e.g.*, ViT-T/16), progressively adapting its components to align with those in LLaMA.
In practice, we observe an attention collapse issue, *i.e.*, the training loss fails to converge by directly adding a causal mask to the attention map. The causal mask restricts the class token from accessing the image’s global information, thereby hindering the optimization of the training loss. To this end, we propose a post-sequence class token technique, repositioning the class token to the end of image tokens (details in Section 3.3). As a result, causal mask can keep the attention score between the class token and others, allowing the model to optimize stably. We also evaluate the advantages of the causal self-attention in reducing computational complexity and enhancing the attention map rank.

Figure: Figure 1: Left: iLLaMA architecture. Right: our design roadmap. Colored and gray bars represent the results of the tiny and base regimes, with the red line depicting the training loss of the tiny regime. iLLaMA strives to process visual tokens using standard LLaMa components, *e.g.*, causal self-attention. The proposed PS [cls] and soft mask strategy help overcome training challenges. Block details of ViT , VisionLLaMA , and our iLLaMA is compared in Figure 5 in Appendix A.
Refer to caption: x1.png

Moreover, we explore several training techniques for the proposed causal Transformer.
When observing things, humans start by broadly catching global connections, then narrow down to focus on specifics.
Motivated by this, we develop a soft mask approach – bi-directional self-attention degenerates to a causal self-attention at the onset of training – to further boost the network performance. Soft mask does not alter the causal self-attention during inference but improves the initial training behavior of the network. (details in Section 3.6). We illustrate different types of masks in Figure 3.

Equipped with these modifications, we propose a decoder-only vision Transformer with causal self-attention inside, dubbed image LLaMA (iLLaMA), as shown in Figure 1. We conduct a thorough evaluation of iLLaMA’s properties, including ImageNet-1K classification , calibration, shape-texture bias, quantization compatibility, ADE20K semantic segmentation , and CIFAR transfer learning . Experimental results show that iLLaMA delivers favorable and reliable performance to its encoder-only counterparts (*i.e.*, ViT, VisionLLaMA), while maintaining a pure decoder design.
More importantly, a spectral analysis on the attention map empirically shows that compared to bi-directional counterparts, causal self-attention has a higher rank (see Figure 4), which allows for learning complex image representation. The contribution of our work can be summarized as follows:

- •
We explore the adaption of the LLaMA decoder to visual tasks and identify the attention collapse issue caused by causal self-attention. To address this, we introduce the PS [CLS] method.
- •
We propose a soft mask strategy to optimize causal self-attention and analyzed the improvement in the rank of the causal attention map.
- •
We develop a series of iLLaMA models and empirically validate its performance on ImageNet, along with practical properties such as quantization compatibility, calibration, shape-texture bias.

We hope our work to inspire a re-evaluation of vision backbone design in the era of LLMs and provide fresh insights for their architectural unification.

## 2 Preliminaries and Motivation

Encoder and decoder. We briefly summarize the encoder and decoder in Transformer . Both of them basically consist of attention module and a MLP module, each followed by a residual connection. The key difference between them is the mask scheme in their self-attention. Encoders use bi-directional self-attention, and decoders employ causal self-attention and cross-attention. However, the latter is typically omitted in decoder-only
LLMs ; , we thus focus on comparing causal and bi-directional self-attention as follows, in terms of the mask setting. Denote ${\bf X}\in\mathbb{R}^{N\times d},{\bf O}\in\mathbb{R}^{N\times d}$ as the input and output sequences, where $N$ is the number of tokens and $d$ is the embedding dimension. $W_{\bf q},W_{\bf k},W_{\bf v}\in\mathbb{R}^{d\times d}$ denotes the linear mapping of query, key and value, respectively. Generally, self-attention module can be formulated as (set the head number and batch size as $1$ for simplicity):

$$ $\small{\bf A}=\frac{1}{\sqrt{d}}(W_{\bf q}({\bf X})\cdot W_{\bf k}({\bf X})^{ \top}),\leavevmode\nobreak\ \leavevmode\nobreak\ \leavevmode\nobreak\ {\bf O}= \operatorname{Softmax}({\bf A}+{\bf M})\cdot W_{\bf v}({\bf X}),\leavevmode \nobreak\ \leavevmode\nobreak\ \leavevmode\nobreak\ {\bf P}_{i,j}=0, \leavevmode\nobreak\ \leavevmode\nobreak\ \leavevmode\nobreak\ {\bf Q}_{i,j}= \left\{\begin{aligned} 0&,i\geq j\\ -\infty&,i<j\end{aligned}\right.$ (1) $$

where $i,j\in[1,N]$, ${\bf A}\in\mathbb{R}^{N\times N}$, ${\bf M}\in\mathbb{R}^{N\times N}$ denote the attention map and mask. ${\bf P}\in\mathbb{R}^{N\times N}$, ${\bf Q}\in\mathbb{R}^{N\times N}$ are masks in the encoder and decoder, respectively.
For a causal self-attention, we have ${\bf M}={\bf Q}$.
Such design allows subsequent tokens only attend to the preceding ones, but not vice versa. For a bi-directional self-attention, we have ${\bf M}={\bf P}$, ensuring mutual
visibility for each token.

Recent LLMs-related image models. Recent image models ; ; are trained with an autoregressive objective, targeting at solving visual tasks. Pang et al.  add a text pre-trained frozen LLM block to a ViT encoder to facilitate the performance. Our work, on the other hand, is motivated to explore in-depth how the decoder design in LLMs can be adapted to image models using simple supervised learning to achieve an architectural alignment. A concurrent work VisionLLaMA  proposes vision models for recognition and generation tasks based on the LLaMA components. Differently, we: 1) introduce causal self-attention from LLaMA, addressing the associated attention collapse issue, while VisionLLaMA retains an encoder architecture; 2) develop a soft mask technique to assist training the decoder; 3) expand the dataset to the larger ImageNet-21K to demonstrate scalability, achieving 86.0% ImageNet accuracy that outperforms VisionLLaMA’s best results.

## 3 A Roadmap: Solving Attention Collapse and Optimization Improvement

This section introduces the design roadmap of iLLaMA. As we aim to adapt LLMs to vision, we choose LLaMA  and ViT  as language and vision baselines due to their successful practices.
The trajectory can be divided into two dimensions, *i.e.*, architecture (Section 3.1-3.4) and training techniques (Section 3.5-3.6). First, we focus on block designs including 1) feed foward network, 2) normalization layer, 3) self-attention, 4) positional embedding, illustrated in Figure 1. Next, we study training techniques and develop a soft mask strategy to facilitate optimization.
Finally, we provide an analysis in terms of efficiency and attention map rank (Section 3.7).
We start with ViT-T/16 and ViT-B/16 with around 5.7M and 86.4M parameters, respectively, and gradually replace the corresponding components with those from LLaMA. We conduct experiments on ImageNet-1K , following the training recipe adopted from  (details in Table 12 of Appendix C.1). Considering the differences between visual perception and text generation tasks, we maintain ViT’s non-autoregressive manner in our network. Each step change and the corresponding results are reported in Table 17 of Appendix D.

### 3.1 Feed Forward Network (FFN)

FFN in Transformer are implemented differently in ViT and LLaMa, *i.e.*, multi-layer perceptron (MLP) and SwiGLU . MLP consists of two sequential linear mappings, with a GELU  function inserted. Meanwhile, SwiGLU combines three linear mappings, integrating a SiLU ; ; function. This structure allows for the modulation of high-dimensional features through a gating mechanism.
We substitute the Transformer’s MLPs with SwiGLUs, while maintaining comparable computational cost. As shown in Figure 1, this improves performance from $73.8\%$ to $74.3\%$, and from $81.3\%$ to $82.0\%$ for the ViT-T/16 and ViT-B/16 regime.
This highlights SwiGLU’s effectiveness not only in language models but also in vision, inspiring further exploration of other components.

We will now use SwiGLU to substitute MLP in each block.

### 3.2 Normalization Layer

Transformers need normalization layer for stable training, *i.e.*, layer normalization (LN)  in ViT and root mean square layer normalization (RMSNorm)  in LLaMA, respectively. We replace all LNs with RMSNorms in our network and empirically observed that the accuracy of the ViT-T/16 regime increased from $74.3\%$ to $74.5\%$. However, similar improvements in precision were not observed in the ViT-B/16 regime (from $82.0\%$ to $81.7\%$). Nonetheless, compared to LN, RMSNorm removes the shift term computation, bringing simplicity to the network ; ; ; .

We will use RMSNorm instead of LN as the normalization layer in each block.

Figure: Figure 2: (a) mask in causal self-attention. (b) mask in causal self-attention with our post-sequence class token (PS [cls]) method. (c) modified causal mask. Their ablation results are shown in Table 1.
Refer to caption: x2.png

### 3.3 Causal Self-Attention Leads to Attention Collapse

Attention collapse issue. As a common practice for Transformer decoders, the key component for causal self-attention is the causal mask, *i.e.*, a lower triangular mask matrix, illustrated in Eq. 1 and Figure 2(a). With such, each token can get the attention score of all its previous ones.
We add the causal mask to our network via a non-autoregressive way.
The reason is that visual perception tasks, unlike text generation, require only inference once.
As a result, we observe that the training loss fails to converge in both ViT-T/16 and ViT-B/16 regimes (line 1 in Table 1). We posit that such issue stems from the influence of the lower triangular matrix, which prevents the class token from "seeing" other image tokens. As illustrated in Figure 2(a), when the class token is positioned at the start of the patch embedding, its attention score for all other image tokens gets zero due to a causal mask. We term this occurrence as the attention collapse issue, which leads to a loss of connection between the class token and other image patches, thereby hindering the optimization of the network.

**Table 1: Results of PS [cls] and the modified causal mask. Training converges in both settings.**
| Model | Tiny | Train Loss | Base | Train Loss |
| --- | --- | --- | --- | --- |
| None | 0.1 | Failed | 0.1 | Failed |
| PS [cls] | 71.9 | 3.599 | 80.6 | 2.869 |
| Modified | 72.5 | 3.550 | 80.4 | 2.857 |

Post-sequence class token (PS [cls]).
The attention collapse issue stems from the inappropriate placement of the token. To this end, we suggest a PS [cls] strategy, by placing it at the end of the token sequence, without changing the causal mask, as shown in Figure 2(b) and Figure 1. Such modification ensures that the class token can achieve global information about all image tokens, while maintaining a causal self-attention property.
As a result, we observe that the attention collapse issue is eliminated and the training process starts to stabilize, leading the network performance to $71.9\%$ for ViT-T/16 and $80.6\%$ for ViT-B/16 regime, respectively (line 2 in Table 1).

To test our hypothesis about the reason of the attention collapse issue, we also explore a mask setting in Figure 2(c). In this setting, we do not change the position of the class token. Instead, we unmask the first row of the mask (*i.e.*, attention score of the class token) on the basis of the causal self-attention, termed as "modified causal mask". Ablation results (line 3 in Table 1) shows that both settings can solve the attention collapse issue as expected, and the "modified causal mask" leads to a better $72.5\%$ accuracy for ViT-T/16 regime, validating our hypothesis about the reason.
Although the results do not surpass the performance of bi-directional counterpart, they demonstrate the potential for optimizing causal self-attention in a decoder-only image model.
We also observe that the PS [cls] method yields higher accuracy with a slightly larger training loss for ViT-B/16 regime, suggesting lower overfitting.

We will employ causal self-attention with the proposed PS [cls] method in each block.

### 3.4 Positional Embedding

A standard ViT use learnable positional embedding (LPE) to preserve positional information, typically adding it directly to the patch embedding. Meanwhile, rotary positional embedding (RoPE)  is widely employed in LLMs ; , which functions in the attention of each block.
We first use RoPE alone, which boosts the accuracy of ViT-T/16 and ViT-B/16 regimes to $72.6\%$ and $81.2\%$, from $71.9\%$ and $80.6\%$, respectively.
The encouraging results illustrate that the concepts of "position" in image and text do not exist an inherent gap. Since LPE functions only once before all Transformer blocks, keeping it does not disrupt the alignment with LLaMA within each block. Thus, we reintroduce the LPE, which improves the accuracy of ViT-T/16 regime from $72.6\%$ to $73.2\%$, suggesting that the two positional embeddings are not redundant but rather synergistic, enhancing network performance.

We will use both LPE and RoPE for positional embedding. So far, we have investigated each block component, and thus fix the final architecture dubbed iLLaMA. Next, we explore training strategies.

### 3.5 Data Augmentation

Mixup  and cutmix  that we used to train our iLLaMA (0.8 and 1.0), are borrowed from DeiT ’s recipe. Unlike the bi-directional self-attention used in DeiT, causal self-attention affects the connection between image tokens. Meanwhile, these two hyper-parameters affect the content of the input image, which further influences the subsequent embedding. Thus, we reevaluate their impact on iLLaMA optimization. Specifically, we discover that a combination of 0.1 mixup and 0.1 cutmix improves the performance of the iLLaMA-T/16 from $73.2\%$ to $74.3\%$, whereas a combination of 0.95 and 1.0 leads the iLLaMA-B/16 to a $81.3\%$ accuracy. Other ablations are detailed in Section 4.1.

### 3.6 Soft Mask Strategy: Optimization Improvement

When observing objects, humans tend to perceive broad connections, then focus on specific details. Motivated by this, we propose a soft mask technique to improve the model’s optimization: starting with bi-directional self-attentions in the early training epochs and gradually shifting completely to causal self-attentions as the optimization goes. Specifically, self-attention can be formulated as:

$$ $\displaystyle{\bf A}$ $\displaystyle=\frac{1}{\sqrt{d}}(W_{\bf q}({\bf X})\cdot W_{\bf k}({\bf X})^{ \top}),\leavevmode\nobreak\ \leavevmode\nobreak\ \leavevmode\nobreak\ {\bf O}= (\operatorname{Softmax}({\bf A})\odot{\bf S})\cdot W_{\bf v}({\bf X}),$ (2) $\displaystyle{\bf S}$ $\displaystyle=\alpha{\bf B}+(1-\alpha){\bf C},\leavevmode\nobreak\ \leavevmode \nobreak\ \leavevmode\nobreak\ {\bf B}_{i,j}=1,\leavevmode\nobreak\ \leavevmode\nobreak\ \leavevmode\nobreak\ {\bf C}_{i,j}=\left\{\begin{aligned} 1&,i\geq j\\ 0&,i<j\end{aligned}\right.$ $$

Figure: Figure 3: (a) Soft mask gradually transitions from a bi-directional mask into a causal mask during training through a constant or linear schedule. (b) Ablation results of training loss and test accuracy.
Refer to caption: x5.png

where $i,j\in[1,N]$, ${\bf S}\in\mathbb{R}^{N\times N}$ denotes the soft mask, which is defined as a linear combination of a bi-directional mask ${\bf B}$ and a causal mask ${\bf C}$. $\alpha$ is the hyper-parameter controlling the mask configuration, *i.e.*, soft mask degenerates into ${\bf B}$ or ${\bf C}$ when $\alpha=1$ or $\alpha=0$, respectively. As illustrated in Figure 3(a), $\alpha$ involves three related hyper-parameters: 1) scheme: how $\alpha$ drops from 1 to 0: we try a linear or a constant scheme. 2) cutoff epochs: when will $\alpha$ drops to 0. 3) learning rate (lr) warmup ; : this hyper-parameter overlaps with the duration of soft mask. We initially set the lr warmup epochs at 50, consistent with previous settings. When using a linear scheme with 50 and 25 cutoff epochs, we observe an improvement in performance for both iLLaMA-T/16 and iLLaMA-B/16 models, reaching $74.9\%$ and $81.6\%$ from $74.3\%$ and $81.3\%$, respectively. Ablations results are detailed in Section. 4.1.
To intuitively observe the impact of soft mask, we plot the training curve of the iLLaMA-T/16 in Figure 3(b), using a constant scheme with 50 cutoff epochs. When soft mask ends, we observe that although there was a sharp drop in accuracy, the model ends up achieving better performance. Similar case of the iLLaMA-B/16 are shown in Figure 6 of Appendix F. Additionally, we discover that a lower learning rate warmup helps iLLaMA-T/16 achieve $75.0\%$ top-1 accuracy, by using a constant scheme with 50 cutoff epochs. Therefore, we use this warmup method for iLLaMA-T/16. Notably, the final training loss within both iLLaMA-T/16 and iLLaMA-B/16 decreases when using soft masks, suggesting a mitigation of potential underfitting concerns.

### 3.7 Analysis of causal Self-attention

Next, we analyze the advantages of using causal self-attention in iLLaMA, in terms of computational efficiency and expressive ability of visual representation through the lens of attention map rank.

Computational Complexity. We compare the efficiency of causal self-attention and bi-directional baseline. For a self-attention with a sequence length $N$ and embedding dimension $D$, FLOPs are reported in Table 2 (RoPE is not involved as only attention computations are considered). causal self-attention, in light of the lower triangular property of its attention map, slightly reduces the FLOPs compared to the bi-directional baseline — the degree of reduction grows as sequence length increases.

**Table 2: Computational complexity results. causal mask slightly reduces FLOPs required in the self-attention.**
| Type | Bi-directional | causal |
| --- | --- | --- |
| FLOPs | $4ND^{2}+2N^{2}D$ | $4ND^{2}+N^{2}D+(\lfloor N^{2}/2\rfloor+1)D$ |

Attention map rank. We examine the representation learning power of causal attention through a spectrum analysis. Following ; , we perform singular value decomposition on the attention maps of the pre-trained ViT-T/16 and iLLaMA-T/16 models.
Next, we sort the singular values and plot a curve illustrating the relationship between the cumulative normalized singular values and matrix indices. The results are conducted using $30$ images randomly selected from the ImageNet-1K validation set. As shown in Figure 4, the curve of ViT exhibits concave function characteristics, while the curve of iLLaMA is close to a linear function, indicating a more uniform distribution of singular values in iLLaMA’s attention map. Approximating the matrix rank by the index at which the cumulative normalized singular value reaches 0.8, we observe that the index value of iLLaMA is about 48 higher than that of ViT ($\sim$129-th v.s. $\sim$81-th). Under such premise, compared to ViT, the attention map of iLLaMA can be approximated with a certain error by a higher-rank matrix. Accordingly, the rank of the attention map may affect the expressive capabilities of the learned representations , suggesting that the causal self-attention in iLLaMA has the potential to learn complex visual representations, as empirically demonstrated in Section 4.2. Detailed results for different layers are provided in Figure 11 of Appendix E.

Figure: Figure 4: Rank analysis of the attention map in head 1, layer 1 of the pretrained ViT-T and iLLaMA-T with $N=197$. Difference between them is about 48.
Refer to caption: x7.png

Closing remarks.
So far, we have finished the design roadmap of iLLaMA through architectural and training strategy modification.
iLLaMA, a decoder-only Transformer, shows advantages in computational complexity and attention map rank through its causal self-attention mechanism. Notably, while all components of iLLaMA are essentially derived from LLaMA, relying only on them is insufficient for an effective weight optimization, as demonstrated in Section 3.3. In fact, the proposed PS [cls] and soft mask strategy effectively address this issue and assist in iLLaMA training. However, to achieve a comprehensive understanding of iLLaMA’s properties, some useful evaluation should be conducted:
1) Scalability for large model capacities (>300M parameters) and dataset sizes (>10M training images, *e.g.*, ImageNet-21K). 2) Other practical evaluation dimensions, such as model calibration, shape-texture bias, downstream task performance, quantization compatibility, discussed below.

## 4 Experiments

This section provide a comprehensive evaluation of iLLaMA. We first report ablation results, *e.g.*, the effectiveness of data augmentation and different soft mask strategies. Next, we compare iLLaMA with other strong baselines on ImageNet classification. Beyond ImageNet accuracy, we also examine its efficacy on calibration, shape-texture bias, and evaluate its compatibility with quantization-aware training and downstream task performance.

### 4.1 Ablation Study

Influence of data augmentation. Base on the observation in Section 3.5, we examined multiple sets of cutmix and mixup settings, as reported in Table 5. We empirically observe that the smaller iLLaMA-T/16 are more sensitive to two data augmentation strategies and perform better with lower hyper-parameters, whereas the larger iLLaMA-B/16 are suited to higher ones. This may be related to the architectural differences between LLaMA’s Transformer decoder and ViT’s encoder type.

Influence of soft mask scheduling strategies and epochs. As mentioned in Section 3.6, the proposed soft mask technique includes three hyper-parameters, *i.e.*, schedule, cutoff epochs and lr warmup epochs. Here we evaluate the robustness of soft mask to hyper-parameter settings, with results detailed in Table 4. Beyond the linear schedule, inspired by , we also implemented a constant option. Additionally, we fixed the learning rate warm-up epochs at 50 and experimented with different cutoff epochs. The results reveal that the soft mask facilitates the optimization of iLLaMA under both linear and constant scheduling, suitable for models of both tiny and base sizes. Moreover, setting the cutoff epochs to span a wide range from 25 to 100 is advantageous. Notably, the soft mask can be easily integrated into existing code frameworks (*e.g.*, timm ) with negligible additional training costs, thereby facilitating its effortless application on future related architectures.

**Table 3: Soft mask scheduling for iLLaMA-T/16 and iLLaMA-B/16 on ImageNet-1K.**
| Schedule | Cutoff Epochs | Tiny | Base |
| --- | --- | --- | --- |
| no softmask | - | 74.3 | 81.3 |
| linear | 25 | 74.8 | 81.6 |
| linear | 50 | 74.9 | 81.5 |
| linear | 100 | 74.9 | 81.5 |
| constant | 25 | 74.7 | 81.5 |
| constant | 50 | 74.8 | 81.5 |

**Table 5: Mixup and cutmix ablation results.**
| Mixup | Cutmix | Tiny | Mixup | Cutmix | Base |
| --- | --- | --- | --- | --- | --- |
| 0.8 | 1.0 | 73.2 | 0.8 | 1.0 | 81.2 |
| 0.5 | 0.4 | 73.8 | 0.9 | 0.9 | 81.2 |
| 0.3 | 0.3 | 73.9 | 0.9 | 1.0 | 81.2 |
| 0.2 | 0.2 | 74.3 | 1.0 | 1.0 | 81.2 |
| 0.1 | 0.1 | 74.3 | 0.95 | 1.0 | 81.3 |

Influence of soft mask for training and testing loss.
Deep neural networks often face underfitting, marked by difficulty in continuously reducing training loss and resulting in poor test accuracy .
We compare the training and testing losses of the iLLaMA-T/16 and iLLaMA-B/16 models with and without the use of the soft mask strategy. As shown in Table 4, soft mask can reduce training loss in both regimes, mitigating potential underfitting issue and reducing testing loss.

### 4.2 Comparison with Recent Architectures on ImageNet-1K Classification

We conducted experiments on the ImageNet-1K  benchmark with different model sizes (*i.e.*, iLLaMA-T/S/B/L). Detailed architecture configurations are shown in Table 11 of Appendix A. Our ImageNet-1K/21K (pre-)training and ImageNet-1K fine-tuning recipes are shown in Table 13 and Table 14 of Appendix C.
We also study the use of LLaMA2-7B pre-trained weights for iLLaMA initialization, and the paradigm and results are detailed in Figure 9 and Table 18 in Appendix I.

ImageNet-1K training. We train iLLaMA-T/S/B on ImageNet-1K for 300 epochs with AdamW optimizer  and a batch size of 4096. The ImageNet-1K trained iLLaMA-T/B models are, in fact, the outcome of the explorations completed in Section 3.6. For the settings of soft mask schedule, cutoff epochs, and learning rate warmup epochs, we tune slightly for the iLLaMA-S model.

ImageNet-21K pre-training. We use the ‘Winter21 variant of ImageNet-21K-P’ (refered to as ImageNet-21K) dataset  (^1^11downloaded from: https://www.image-net.org/download-images.php) for the large-scale pre-training of our iLLaMA, which contains 11,060,223 training images and 522,500 testing images from 10,450 classes. Only the training set was used. We pre-train iLLaMA-B/L on ImageNet-21K for 90 epochs using a constant soft mask schedule, with cutoff epochs and learning rate warmup epochs set to 30 and 5, respectively.

ImageNet-1K fine-tuning. For iLLaMA-B model trained on ImageNet-1K, we fine-tune at a resolution of 384$\times$384. Similarly, for the iLLaMA-B/L model trained on ImageNet-21K, we fine-tune at resolutions of 224$\times$224 and 384$\times$384, respectively. All fine-tuning was conducted for 30 epochs using the AdamW optimizer. We follow DeiT  for interpolating positional embeddings to allow our iLLaMA to handle inputs at a higher resolution.

Results. Table 6 shows a comparison between iLLaMA and other strong visual baselines, including ConvNets (ConvNeXt , ConvNeXt-V2 ), vision Transformers (ViT , Swin Transformer ), MLPs (PoolFormer , VanillaNet ), and recent language model inspired models (AIM , VisionLLaMA ). We present three observations: 1) The performance-parameter trade-off of iLLaMA surpasses other LM-inspired models such as AIM and VisionLLaMA, presumably due to its use of causal attention and soft mask training techniques. 2) iLLaMA exhibits a superior accuracy-throughput trade-off compared to strong hierarchical baselines such as ConvNeXt-V2-N/T/B and Swin-S/B. We attribute this to iLLaMA’s isotropic design (each intermediate block has the same feature resolution), which benefits from a straightforward and efficient architecture, enhancing inference speed. 3) Scalability of model capacity and dataset size: After comprehensive pre-training on the expanded ImageNet-21K dataset, the iLLaMA-B model achieves more than $85.0\%$ accuracy on ImageNet-1K with under 100M parameters, significantly outperforming ViT-B’s $84.0\%$. Upon scaling up to the larger iLLaMA-L, accuracy reaches $86.0\%$, exceeding that of ViT-L pre-trained on ImageNet-21K and the AIM-7B pre-trained on the DFN-2B+ dataset. To our knowledge, this showcases SOTA performance for LLaMA-type architectures.

**Table 6: ImageNet-1K accuracy. Throughput (images/s) are tested on Nvidia A100 GPU with a batch size of 1024. Hie.: Hierarchical, Iso.: Isotropic, Sup.: Supervised (pre-)training, AR.: Autoregressive pre-training. $\spadesuit$ ConvNet, $\blacksquare$ Vision Transformer, $\clubsuit$ MLP, $\maltese$ LM-inspired visual model, $\bigstar$ LLaMA.**
| Model | Dataset Used | Objective | Type | Image Size | Params | MACs | Throughput | Acc |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| $\spadesuit$ ConvNeXt-S | IN-1K | Sup. | Hie. | 224$\times$224 | 50M | 8.7G | 1185 | 83.1 |
| $\spadesuit$ ConvNeXt-B | IN-1K | Sup. | Hie. | 224$\times$224 | 89M | 15.4G | 877 | 83.8 |
| $\spadesuit$ ConvNeXt-L | IN-1K | Sup. | Hie. | 224$\times$224 | 198M | 34.4G | 543 | 84.3 |
| $\spadesuit$ ConvNeXtV2-N | IN-1K | Sup. | Hie. | 224$\times$224 | 15.6M | 2.45G | 2120 | 81.2 |
| $\spadesuit$ ConvNeXtV2-T | IN-1K | Sup. | Hie. | 224$\times$224 | 28.6M | 4.47G | 1362 | 82.5 |
| $\spadesuit$ ConvNeXtV2-B | IN-1K | Sup. | Hie. | 224$\times$224 | 88.7M | 15.4G | 645 | 84.3 |
| $\blacksquare$ Swin-S | IN-1K | Sup. | Hie. | 224$\times$224 | 50M | 8.7G | 934 | 83.0 |
| $\blacksquare$ Swin-B | IN-1K | Sup. | Hie. | 224$\times$224 | 88M | 15.4G | 710 | 83.5 |
| $\blacksquare$ DeiT-Ti | IN-1K | Sup. | Iso. | 224$\times$224 | 5.7M | 1.3G | 6051 | 72.2 |
| $\blacksquare$ DeiT-S | IN-1K | Sup. | Iso. | 224$\times$224 | 22.1M | 4.6G | 3080 | 79.8 |
| $\blacksquare$ DeiT-B | IN-1K | Sup. | Iso. | 224$\times$224 | 86.4M | 17.6G | 1348 | 81.8 |
| $\blacksquare$ ViT-B/16 | IN-21K, IN-1K | Sup., Sup. | Iso. | 384$\times$384 | 86.4M | 55.5G | 349 | 84.0 |
| $\blacksquare$ ViT-L/16 | IN-21K, IN-1K | Sup., Sup. | Iso. | 384$\times$384 | 304.1M | 191.2G | 124 | 85.2 |
| $\clubsuit$ PoolFormer-S12 | IN-1K | Sup. | Hie. | 224$\times$224 | 12M | 1.8G | 4354 | 77.2 |
| $\clubsuit$ PoolFormer-M48 | IN-1K | Sup. | Hie. | 224$\times$224 | 73M | 11.6G | 768 | 82.5 |
| $\clubsuit$ VanillaNet-5 | IN-1K | Sup. | Hie. | 224$\times$224 | 15.5M | 5.2G | - | 72.5 |
| $\clubsuit$ VanillaNet-13-1.5$\times$ | IN-1K | Sup. | Hie. | 224$\times$224 | 127.8M | 26.5G | - | 82.5 |
| $\maltese$ AIM-0.6B | DFN-2B+, IN-1K | AR., Sup. | Iso. | 224$\times$224 | 0.6B | - | - | 78.5 |
| $\maltese$ AIM-3B | DFN-2B+, IN-1K | AR., Sup. | Iso. | 224$\times$224 | 3B | - | - | 82.2 |
| $\maltese$ AIM-7B | DFN-2B+, IN-1K | AR., Sup. | Iso. | 224$\times$224 | 7B | - | - | 82.4 |
| $\maltese$ P-VisionLLaMA-S | IN-1K | Sup. | Hie. | 224$\times$224 | 24M | - | - | 81.6 |
| $\maltese$ P-VisionLLaMA-B | IN-1K | Sup. | Hie. | 224$\times$224 | 56M | - | - | 83.2 |
| $\maltese$ P-VisionLLaMA-L | IN-1K | Sup. | Hie. | 224$\times$224 | 99M | - | - | 83.6 |
| $\maltese$ VisionLLaMA-L | IN-1K, IN-1K | Sup., Sup. | Iso. | 224$\times$224 | 310M | - | - | 84.6 |
| $\bigstar$ iLLaMA-T | IN-1K | Sup. | Iso. | 224$\times$224 | 5.7M | 1.3G | 6958 | 75.0 |
| $\bigstar$ iLLaMA-S | IN-1K | Sup. | Iso. | 224$\times$224 | 21.9M | 4.6G | 3222 | 79.9 |
| $\bigstar$ iLLaMA-B | IN-1K | Sup. | Iso. | 224$\times$224 | 86.3M | 17.6G | 1345 | 81.6 |
| $\bigstar$ iLLaMA-B | IN-1K | Sup. | Iso. | 384$\times$384 | 86.3M | 55.5G | 332 | 83.0 |
| $\bigstar$ iLLaMA-B | IN-21K, IN-1K | Sup., Sup. | Iso. | 224$\times$224 | 86.3M | 17.6G | 1345 | 83.6 |
| $\bigstar$ iLLaMA-B | IN-21K, IN-1K | Sup., Sup. | Iso. | 384$\times$384 | 86.3M | 55.5G | 332 | 85.0 |
| $\bigstar$ iLLaMA-L | IN-21K, IN-1K | Sup., Sup. | Iso. | 224$\times$224 | 310.2M | 62.8G | 456 | 84.8 |
| $\bigstar$ iLLaMA-L | IN-21K, IN-1K | Sup., Sup. | Iso. | 384$\times$384 | 310.2M | 194.7G | 116 | 86.0 |

### 4.3 Model Calibration and Shape-Texture Bias

Beyond ImageNet accuracy, we also examined iLLaMA’s calibration properties and shape-texture bias for a more detailed evaluation, following . Besides iLLaMA, we also explore two prevalent architectures, *i.e.*, ConvNeXt  and DeiT3 , representing ConvNets and Transformers, respectively. We apply ImageNet-21K pre-trained and ImageNet-1K fine-tuned models in this section.

Model calibration.
Model calibration represents the relationship between a model’s precision and confidence across samples of varying difficulty, *i.e.*, poor-calibrated models tend to produce overly confident yet incorrect predictions, whereas well-calibrated models demonstrate a strong correlation between confidence and accuracy . Calibration is commonly measured using the Expected Calibration Error (ECE), where a lower ECE is favorable. ECE results for different models on ImageNet-1K are presented in Table 8. The calibration of iLLaMA is lower than that of DeiT3, suggesting a more reliable output confidence. We also plot the reliability diagrams  to intuitively compare the calibration of different models, as shown in Figure 7 of Appendix G.

Shape-texture bias. Shape-texture bias measures the extent to which the model relies on the shape or texture of the image when performing recognition . We generally prefer models to mimic human eye behavior, relying more on shape rather than texture ; . We calculate the shape ratio for all models on cue-conflict images and report the results in Table 8, following . Our iLLaMA shows the largest shape ratio of $41.45\%$ among the three compared baselines, suggesting the potential of the LLM architecture for vision. Detailed results are provided in Figure 8 of Appendix H.

### 4.4 Compatibility with Quantization

Since a practical goal for neural networks is deployment on low-bit hardware chips, we further examine iLLaMA’s compatibility with quantization. We basically follow Q-ViT  to apply quantization-aware training (QAT) to iLLaMA, with weights and activations of all blocks’ FFN and causal self-attention layers to 8 bits. Quantization recipes and results are shown in Table 15 of Appendix C.4 and Table 8. Different sizes of low-bit iLLaMA maintain accuracy well, and 8-bit iLLaMA-T is even compete favorably with the full-precision DeiT-T  ($72.4\%$ v.s. $72.2\%$).

**Table 7: Quantization results. $\#$Bits ($\mathrm{w}$-$\mathrm{a}$): $\mathrm{w}$ bit weights, $\mathrm{a}$ bit activations. 8-bit iLLaMA-T matches 32-bit DeiT-T.**
| Model | $\#$Bits | Tiny | Small |
| --- | --- | --- | --- |
| DeiT | 32-32 | 72.2 | 79.8 |
| iLLaMA | 32-32 | 75.0 | 79.9 |
| iLLaMA | 8-8 | 72.4 | 77.4 |

### 4.5 Transferability on Downstream Tasks

CIFAR transfer learning. We fine-tune ViT-T and iLLaMA-T on the CIFAR datasets , including an ablation of the soft mask on iLLaMA. Detailed recipes are shown in Table 16 of Appendix C.5. iLLaMA’s performance on CIFAR datasets is essentially on par with ViT, assuring that iLLaMA can be confidently applied in the transfer learning field as a practical alternative to ViT. Additionally, soft mask is helpful in the relatively complicated CIFAR100, demonstrating its generalizability.

ADE20K semantic segmentation. We fine-tune our ImageNet-1K pre-trained iLLaMA and ViT models on ADE20K  dataset using UperNet  to perform semantic segmentation task. For both iLLaMA and ViT, we set the learning rate as 6e-5 and weight decay as 0.01. Table 10 presents the results. iLLaMA’s performance is marginally lower than ViT’s, which we attribute to the potential impact of the masking mechanism in iLLaMA’s causal attention on high-resolution dense prediction tasks. This suggests there is still space for optimization, a subject for future investigation.

**Table 9: Soft mask for CIFAR transfer learning. Soft mask improves iLLaMA performance without changing the inference architecture.**
| Model | CIFAR10 | CIFAR100 |
| --- | --- | --- |
| ViT-T | 98.0 | 85.5 |
| iLLaMA-T | 97.9 | 84.8 |
| + soft mask | 97.9 | 85.5 |

## 5 Conclusions

In this paper, we systematically studies whether Transformer decoder, an architecture that has shown amazing potential in LLMs, can also take root in learning visual representation through straightforward supervised training. The key component – causal self-attention we used – is not novel and is inherited from existing LLM architectures, but we propose pivotal techniques, *i.e.*, PS [cls] and soft mask strategies, to effectively adapt them to visual tasks. The proposed iLLaMA outperforms many ConvNets, ViTs, and MLPs on imagenet, and demonstrates robust quantization compatibility, calibration, and shape-texture bias, thereby showing its practicality. We hope that this work will inspire a rethinking of generic yet practical architecture that can fully unify both vision and text.

## Appendix A Network Configuration

In Table 11, we provide detailed architecture configurations for iLLaMA models of various capacities. Our approach to scaling up the model size, from small to large, is similar to that of the ViT. Thus, akin to ViT, iLLaMA benefits from the simplicity of an isotropic architecture and high throughput, with its internal features remaining unchanged in resolution and number of channel as the depth increases.

**Table 11: Detailed iLLaMA architecture configurations.**
|  | Tiny (T) | Small (S) | Base (B) | Large (L) |
| --- | --- | --- | --- | --- |
| depth | 12 | 12 | 12 | 24 |
| embedding dim | 192 | 384 | 768 | 1024 |
| number of heads | 3 | 6 | 12 | 16 |
| #param. (M) | 5.7 | 21.9 | 86.3 | 310.2 |
| MACs (G) | 1.3 | 4.6 | 17.6 | 62.8 |

We provide a block-level comparison between iLLaMA and ViT model in Figure 5.
VisionLLaMA uses SwiGLU, and AS2D RoPE to build LLaMA-style architecture. Differently, we further uses RMSNorm, modified causal self-attention and 1D RoPE from LLaMA to replace layer normalization, bi-directional self-attention, and proposes two pivotal strategies, *i.e.*, PS [cls] and soft mask to help the optimization of our iLLaMA. We also keep the learnable positional embedding as ViT does.

Figure: Figure 5: Comparison between ViT , VisionLLaMA , and iLLaMA blocks.
Refer to caption: x8.png

## Appendix B PyTorch-like Code of iLLaMA Causal Self-attention

The PyTorch-like implementation of our iLLaMA causal self-attention is shown as Algorithm 1.
The iLLaMA code exhibits a high degree of similarity in structure and composition to the official LLaMA code (^2^22[https://github.com/meta-llama/llama](https://github.com/meta-llama/llama)) released by Meta, potentially offering considerable coding cost savings in developing a unified vision and language network with such architecture.

Figure: Algorithm 1 PyTorch code of iLLaMA causal self-attention

## Appendix C Experimental Settings

### C.1 Training Recipe in Section 3

Our training recipe for training the tiny and base models during the “designing iLLaMA: a roadmap” (Section 3) is primarily adapted from ConvNeXt ; , summarized in Table 12.

Basically, both regimes use the same experimental setup, with the only difference being the stochastic depth rate at 0.0 and 0.4, respectively. Notably, for the ViT baseline, our experimental results are $73.8\%$ and $81.5\%$, as shown in Table 17, which slightly differ from the results of $73.9\%$ and $81.6\%$ reported in .

Utilizing only the basic training recipe with architectural modifications, the performance of iLLaMA’s tiny and base models achieves $73.2\%$ and $81.2\%$, as shown in Table 17, yet remains below the ViT baseline. We attribute this to the impairing effect of causal self-attention on the information mixing among tokens. Thus, we enhance the training recipe, detailed next.

**Table 12: Our training recipe for Section 3 in the main paper, adapted from .**
| Training Configuration | iLLaMA-T/B |
| --- | --- |
| Initialization: |  |
| weight init | trunc. normal (0.2) |
| Training recipe: |  |
| optimizer | AdamW |
| optimizer momentum | $\beta_{1},\beta_{2}{=}0.9,0.999$ |
| Learning hyper-parameters: |  |
| base learning rate | 4e-3 |
| learning rate schedule | cosine decay |
| weight decay | 0.05 |
| batch size | 4096 |
| training epochs | 300 |
| lr warmup epochs | 50 |
| warmup schedule | linear |
| gradient clip | None |
| exp. mov. avg. (EMA) | None |
| Dropout: |  |
| dropout rate | 0.0 |
| stochastic depth rate | 0.0/0.4 |
| Data augmentation: |  |
| input resolution | $224^{2}$ |
| randAugment | (9, 0.5) |
| random erasing | 0.25 |
| label smoothing | 0.1 |
| mixup | 0.8 |
| cutmix | 1.0 |

### C.2 ImageNet (Pre-)training Recipe

As illustrated in Table 13, we provide the detailed ImageNet-1K training hyper-parameters and ImageNet-21K pre-training hyper-parameters for the experimental results in Table 6.

For the iLLaMA-T/S/B models, we train directly on ImageNet-1K and discover that models of different sizes are suited to different soft mask settings. For instance, the soft mask schedules are set to constant/linear/linear, respectively, with cutoff epochs designated as 50/50/25. We train the iLLaMA-T/S/B models using 8 A100 GPUs.

We pre-trained the iLLaMA-B/L models on ImageNet-21K for 90 epochs, adhering to the practices in . We set the cutoff epochs to 30, indicating that the iLLaMA models’ self-attention fully transitions to causal self-attention after 30 epochs. We pre-train the iLLaMA-B/L models using 8 A100 GPUs.

**Table 13: Our (pre-)training settings for iLLaMa model on ImageNet-1K/ImageNet-21K, respectively, adapted from . Some key training techniques are highlighted.**
|  | iLLaMA-T/S/B | iLLaMA-B/L |
| --- | --- | --- |
| (Pre-)Training Configuration | ImageNet-1K | ImageNet-21K |
| Initialization: |  |  |
| weight init | trunc. normal (0.2) | trunc. normal (0.2) |
| Training recipe: |  |  |
| optimizer | AdamW | AdamW |
| optimizer momentum | $\beta_{1},\beta_{2}{=}0.9,0.999$ | $\beta_{1},\beta_{2}{=}0.9,0.999$ |
| Learning hyper-parameters: |  |  |
| base learning rate | 4e-3 | 1e-3 |
| learning rate schedule | cosine decay | cosine decay |
| weight decay | 0.05 | 0.01 |
| batch size | 4096 | 4096 |
| training epochs | 300 | 90 |
| warmup schedule | linear | linear |
| gradient clip | None | None |
| exp. mov. avg. (EMA) | None | None |
| Dropout: |  |  |
| dropout rate | 0.0 | 0.0 |
| stochastic depth rate | 0.0/0.1/0.4 | 0.1 |
| Data augmentation: |  |  |
| input resolution | $224^{2}$ | $224^{2}$ |
| randAugment | (9, 0.5) | (9, 0.5) |
| random erasing | 0.25 | 0.25 |
| label smoothing | 0.1 | 0.1 |
| mixup | 0.1/0.5/0.95 | 0.8 |
| cutmix | 0.1/0.5/1.0 | 1.0 |
| Soft mask: |  |  |
| soft mask schedule | constant/linear/linear | constant |
| cutoff epochs | 50/50/25 | 30 |
| lr warmup epochs | 5/5/50 | 5 |

### C.3 ImageNet Fine-tuning Recipe

We present the results of fine-tuning models pre-trained on ImageNet-1K at a resolution of $384\times 384$, as well as the outcomes of fine-tuning models pre-trained on ImageNet-21K at resolutions of $224\times 224$ and $384\times 384$, as shown in Table 14. All ImageNet-1K fine-tuning experiments were conducted for 30 epochs, following the convention in .

For the iLLaMA-B model pre-trained on ImageNet-1K, we used a relatively higher stochastic depth rate of 0.8. For the iLLaMA-B/L models pre-trained on ImageNet-21K, we employed relatively lower stochastic depth rates of 0.2 and 0.3, respectively.

Additionally, we standardized the cutoff epoch at 0 for our ImageNet-1K fine-tuning experiments, ensuring the application of a causal mask in self-attention to align with the LLaMA architecture. We also opted not to use learning rate warmup. We fine-tune the models using 8 A100 GPUs.

**Table 14: Our fine-tuning settings for iLLaMa model on ImageNet-1K, adapted from . Some key training techniques are highlighted.**
|  | iLLaMA-B | iLLaMA-B/L | iLLaMA-B/L |
| --- | --- | --- | --- |
| (Pre-)Training Configuration | ImageNet-1K | ImageNet-21K | ImageNet-21K |
| 224^2 | 224^2 | 224^2 |  |
| Fine-Tuning Configuration | ImageNet-1K | ImageNet-1K | ImageNet-1K |
| Initialization: |  |  |  |
| weight init | trunc. normal (0.2) | trunc. normal (0.2) | trunc. normal (0.2) |
| Training recipe: |  |  |  |
| optimizer | AdamW | AdamW | AdamW |
| optimizer momentum | $\beta_{1},\beta_{2}{=}0.9,0.999$ | $\beta_{1},\beta_{2}{=}0.9,0.999$ | $\beta_{1},\beta_{2}{=}0.9,0.999$ |
| Learning hyper-parameters: |  |  |  |
| base learning rate | 8e-5 | 8e-5/6e-5 | 1.1e-4/3.5e-5 |
| learning rate schedule | cosine decay | cosine decay | cosine decay |
| weight decay | 1e-8 | 1e-8 | 1e-8 |
| batch size | 512 | 512 | 512 |
| training epochs | 30 | 30 | 30 |
| warmup schedule | linear | linear | linear |
| gradient clip | None | None | None |
| exp. mov. avg. (EMA) | None | None | None |
| Dropout: |  |  |  |
| dropout rate | 0.0 | 0.0 | 0.0 |
| stochastic depth rate | 0.8 | 0.2/0.3 | 0.2/0.3 |
| Data augmentation: |  |  |  |
| input resolution | $384^{2}$ | $224^{2}$ | $384^{2}$ |
| randAugment | (9, 0.5) | (9, 0.5) | (9, 0.5) |
| random erasing | 0.25 | 0.25 | 0.25 |
| label smoothing | 0.1 | 0.1 | 0.1 |
| mixup | 0 | 0 | 0 |
| cutmix | 0 | 0 | 0 |
| Soft mask: |  |  |  |
| soft mask schedule | constant | constant | constant |
| cutoff epochs | 0 | 0 | 0 |
| lr warmup epochs | 0 | 0 | 0 |

### C.4 Quantization-aware Training Recipe

We provide our quantization-aware training recipe for iLLaMA in Table 15. Basically we follow the Q-ViT method proposed in , with only weights and activations in each basic block’s causal self-attention and FFN module are quantized to 8 bit width.

**Table 15: Our quantization-aware training settings for iLLaMa model on ImageNet-1K, adapted from ; . Some key training techniques are highlighted.**
|  | iLLaMA-T/S |
| --- | --- |
| (Pre-)Training Configuration | ImageNet-1K |
| Initialization: |  |
| weight init | trunc. normal (0.2) |
| Training recipe: |  |
| optimizer | AdamW |
| optimizer momentum | $\beta_{1},\beta_{2}{=}0.9,0.999$ |
| Learning hyper-parameters: |  |
| base learning rate | 3e-3/4e-3 |
| learning rate schedule | cosine decay |
| weight decay | 0.05 |
| batch size | 4096 |
| training epochs | 300 |
| warmup schedule | linear |
| gradient clip | None |
| exp. mov. avg. (EMA) | None |
| Dropout: |  |
| dropout rate | 0.0 |
| stochastic depth rate | 0.0/0.1 |
| Data augmentation: |  |
| input resolution | $224^{2}$ |
| randAugment | (9, 0.5) |
| random erasing | 0.25 |
| label smoothing | 0.1 |
| mixup | 0.1/0.5 |
| cutmix | 0.1/0.5 |
| Soft mask: |  |
| soft mask schedule | constant/linear |
| cutoff epochs | 50/50 |
| lr warmup epochs | 5/5 |

### C.5 CIFAR Transfer Learning Recipe

We further provide our training recipe for transfer learning on the CIFAR10 and CIFAR100 datasets, as shown in Table 16.

In our transfer learning experiments, we consistently apply a linear soft mask schedule. However, for the CIFAR10 and CIFAR100 datasets, we use cutoff epochs of 25 and 50, respectively.

**Table 16: Our transfer learning settings for ViT-T and iLLaMa-T model on CIFAR10/100, respectively, adapted from . Some key training techniques are highlighted.**
| Transfer Learning Configuration | CIFAR10 | CIFAR100 |
| --- | --- | --- |
| For both ViT-T and iLLaMA-T: |  |  |
| base learning rate | 2e-3 | 2e-3 |
| batch size | 1024 | 1024 |
| training epochs | 300 | 300 |
| stochastic depth rate | 0.0 | 0.0 |
| lr warmup epochs | 50 | 50 |
| For iLLaMA-T only: |  |  |
| soft mask schedule | linear | linear |
| cutoff epochs | 25 | 50 |

Figure: Figure 6: Training curves for iLLaMA-B/16 regime w/ and w/o soft mask. When soft mask ends, the model experiences a similar pattern to the training curve of iLLaMA-T/16 regime, with eventually a lower test loss observed.
Refer to caption: x9.png

## Appendix D Designing iLLaMA: detailed results

We present the comprehensive experimental results of our exploration journey of iLLaMA in Table 17. This table not only delineates the stepwise accuracy of both the tiny and base models, as depicted in Figure 1, but also outlines the training loss at each step. The general trend observed is that as the training loss of the models decreases, their accuracy increases.

Overall, the trend in changes for the base model is broadly similar to that of the tiny model. However, in contrast to the tiny model, the implementation of RoPE coupled with subsequent integration of LPE does not affect the base model’s performance. This lack of impact, we theorize, stems from the base regime’s reduced susceptibility to underfitting compared to the tiny regime, hence the addition of extra learnable parameters offers less benefit to its performance.

Notably, vanilla causal self-attention mechanism proves inadequate for model optimization, an issue effectively addressed by implementing the PS [CLS] method. Additionally, the application of the soft mask technique significantly contributes to the training efficacy of both model sizes.

**Table 17: ImageNet-1K classification accuracy via gradually replacing components in ViT-T/16 and ViT-B/16 with counterparts in LLaMA, better or worse than the ViT baseline results with our basic training recipe. Components from or modified from LLaMA are highlighted. P.E.: positional embedding, Bd.: bi-directional self-attention, Cs.: causal self-attention.**
| Ablation | FFN | Norm | Attention | P.E. | Tiny | Train Loss | Base | Train Loss |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| ViT | MLP | LN | Bd. | LPE | 72.2 | - | 81.8 | - |
| results with our basic training recipe |  |  |  |  |  |  |  |  |
| ViT | MLP | LN | Bd. | LPE | 73.8 | 3.451 | 81.5 | 2.828 |
| + LLaMa FFN | SwiGLU | LN | Bd. | LPE | 74.3 | 3.407 | 82.0 | 2.724 |
| + LLaMa Norm | SwiGLU | RMS | Bd. | LPE | 74.5 | 3.406 | 81.7 | 2.721 |
| + LLaMa S.A. | SwiGLU | RMS | Cs. | LPE | 0.1 | Failed | 0.1 | Failed |
| + LLaMa S.A. | SwiGLU | RMS | Cs. + PS [CLS] | LPE | 71.9 | 3.599 | 80.6 | 2.869 |
| + LLaMa P.E. | SwiGLU | RMS | Cs. + PS [CLS] | RoPE | 72.6 | 3.618 | 81.2 | 2.861 |
| + LPE P.E. | SwiGLU | RMS | Cs. + PS [CLS] | RoPE + LPE | 73.2 | 3.531 | 81.2 | 2.839 |
| modify the training techniques |  |  |  |  |  |  |  |  |
| + data aug. | SwiGLU | RMS | Cs. + PS [CLS] | RoPE + LPE | 74.3 | 2.990 | 81.3 | 2.868 |
| + soft mask | SwiGLU | RMS | Cs. + PS [CLS] | RoPE + LPE | 75.0 | 2.955 | 81.6 | 2.828 |

## Appendix E Rank Analysis of causal Self-attention

#### Detailed visualization results.

We provide rank analysis results of all 3 heads in layer 1, 4, 8, 12 of ViT-T/16 and iLLaMA-T/16 in Figure 11. Besides the observation in Section 3.7, We make four observations:
1) Not each head in each layer of iLLaMA’s self-attention shows stronger concavity, suggesting that not every attention matrix of iLLaMA has a higher rank than its ViT counterpart.
2) In most cases, particularly in the shallow layers, the distribution of singular values in iLLaMA appears more uniform than in ViT.
3) In certain attention maps (*e.g.*, layer 8, head 2, and layer 8, head 3), the rank of ViT’s attention matrix is low, resulting in an skewed distribution of information. In contrast, such extreme cases were not observed in our iLLaMA.
4) The distribution of singular values in ViT varies significantly across different layers and heads (e.g., layer 1, head 1; layer 4, head 1; layer 8, head 1; layer 8, head 2), whereas iLLaMA’s distribution appears relatively more stable.

## Appendix F Analysis for Soft Mask Method

In this section, we plot the training results for iLLaMA-B/16 with and without the use of the soft mask technique in Figure 6. We can observe that the results display a similar pattern to those of iLLaMA-T/16 (Figure 3(b)).

We set the cutoff epochs to 50 and used a constant schedule. When the soft mask ends, there is a sudden increase in training loss and a steep decline in model accuracy. However, the final accuracy surpasses the baseline, and the training loss is also optimized to a lower value. Such phenomenon shows the versatility of the soft mask technique across models of varying capacities, and shows that causal self-attention can still effectively model even when a portion of the attention map is masked.

## Appendix G Model Calibration

Figure: Figure 7: Calibration results of (a) ConvNeXt-B (b) DeiT3-B and (c) iLLaMA-B pretrained on ImageNet-21K and fine-tuned on ImageNet-1K.
Refer to caption: x10.png

To qualitatively present the calibration property, we plot the reliability diagrams of ConvNeXt-B, DeiT3-B and the proposed iLLaMA-B using ImageNet-1K in Figure 7, following . For well-calibrated models, the direction of accuracy in their reliability diagrams show a roughly diagonal pattern, *i.e.*, the difference between accuracy and confidence is small. Intuitively, the confidence of the early bins of the iLLaMA presents results below the accuracy level, indicating that iLLaMA tends to be underconfident. This observation, akin to that observed in the DeiT3, may be a common characteristic of Transformer-based architectures and was also noted in .

## Appendix H Shape-Texture Bias

We visualize the shape-texture bias results on cue-conflict images of ConvNeXt-B, DeiT3-B and the proposed iLLaMA-B in Figure 8, following . The three dashed lines of different colors represent the average shape decision of the three models over all categories. Generally, a more leftward average shape ratio indicates that the model relies more on global shape information for recognition tasks. iLLaMA shows higher shape scores relative to ConvNeXt and DeiT3.

Figure: Figure 8: Shape-texture bias results of ConvNeXt-B, DeiT3-B and iLLaMA-B pre-trained on ImageNet-21K and fine-tuned on ImageNet-1K. sup: supervised learning paradigm.
Refer to caption: x13.png

## Appendix I Initializating iLLaMA Using Pre-trained LLaMA

Previous studies  have demonstrated that data unrelated to the image modality can be used to improve the performance of visual models. In fact, the pre-training dataset of LLaMA, which is entirely text, is irrelevant to the visual tasks that iLLaMA addresses. More important, the architectural components of iLLaMA are aligned with those of LLaMA. This alignment facilitates our exploration of using LLaMA’s parameters to initialize iLLaMA, allowing us to fully exploit the potential of the weights of pre-trained LLMs.

**Table 18: Ablation results of weight selection of iLLaMA using LLaMA2-7B pre-trained weights.**
| Model | Initialization | Tiny | Small | Base |
| --- | --- | --- | --- | --- |
| iLLaMA | w/ weight selection | 74.5 | 79.9 | 81.4 |
| iLLaMA | w/o weight selection | 75.0 | 79.9 | 81.6 |

We use the pre-trained LLaMA2-7B  to initialize our iLLaMA, instead of training from scratch. To match the dimensions of the weights, we employ the weight selection  method to initialize iLLaMA-T/S/B using a subset of the LLaMA2-7B pre-trained weights. Next, we train and evaluate the iLLaMA models, which are initialized using LLaMA2-7B, on the ImageNet-1K dataset. Other hyperparameter settings are consistent with Section 4.2. The results are shown in Table 18.
Using LLaMA2 to initialize iLLaMA does not yield significant performance improvements. We attribute this to two main reasons:
1) The size difference between the two models is considerably large (LLaMA-2-7B’s 7B parameters vs. iLLaMA-T’s 5.7M parameters), resulting in a limited proportion of selected weights compared to meaningful pre-trained weights.
2) The training strategy was not adequately optimized. We believe that when using parameter inheritance, the corresponding training strategy should also be adjusted. However, we continued to use the training recipe designed for training from scratch.

Figure: Figure 9: iLLaMA initialization by pre-trained LLaMA2-7B using weight selection .
Refer to caption: x14.png

## Appendix J Loss Landscape

Figure: Figure 10: Loss landscape illustration of (a) ViT-T/16 and (b) iLLaMA-T/16.
Refer to caption: x15.png

As shown in Figure 10, we visualized the loss landscape  of the iLLaMA-T/16 and ViT-T/16. The loss landscape of ViT and iLLaMA exhibits similar patterns, with the steepness and bumps observed in ViT seeming to be softened.

Figure: (a) layer 1, head 1
Refer to caption: x17.png

## Appendix K Limitations

We have shown that the LLaMA architecture, enhanced by the developed post-sequence [cls] and soft mask techniques, is adept at adapting to tasks such as visual recognition and semantic segmentation. However, iLLaMA’s application remains predominantly within the realm of perception. In fact, such decoder-only architecture, favored by LLMs in the NLP field, can do more complex tasks, such as reasoning and generation. This may be due to their massive training data and the next sentence prediction training paradigm, which is not explored by iLLaMA. Thus, a critical validation step of aligning the architectures of text and visual models would be to construct a multi-modal large language model that fully leverages LLaMA components. In this envisioned model, both visual and textual feature extractors would be realized through the LLaMA architecture.
Futhermore, we strongly argue that iLLaMA’s successful attempts at basic supervised training strategies and classification tasks provide a foundation for more complex tasks, such as object detection and depth estimation. This represents a compelling avenue for future research.

## Appendix L Societal Impact

After the ChatGPT milestone in 2022, open-source architectures like LLaMA began to shine in the text domain. In the real world, images and text are the two main mediums of information and data types. For neural networks, having a unified architecture for language and vision models allows people to process these two distinct types of information using the same structure, which aids in the specialization of hardware implementation. This paper transfers the architecture widely used in language models to vision models, facilitating the achievement of this goal. The pretrained models and code provided in this paper can be used in a plug-and-play manner to serve this objective.

## Peer Reviews

confidence: 4
contribution: 1
initial_rating: 3
presentation: 3
questions: - Why should we design the decoder-only architecture, unlike the one with encoder architecture like VisionLLaMA?
   - If the objective is to align vision and language modalities within a single architecture, as stated in line 35, the authors should have demonstrated either cross-modal transferability or developed a unified model capable of handling both modalities.
soundness: 3
strengths: - The paper is well-written and clearly structured, making it easy to follow.
- Through extensive ablation studies, the authors effectively demonstrate the impact of each component.
summary: The paper investigates the potential of using a decoder-only Transformer architecture, specifically LLaMa, as a vision Transformer (ViT) classifier for image processing. The authors propose a model called image LLaMa (iLLaMA), by transforming ViT into LLaMA-like architecture through step-by-step, component-wise modification. During the modification of attention mechanism, the authors discovered the a key issue called "attention collapse" that hampers network training. To mitigate this, they introduce a post-sequence class token and a soft mask strategy to facilitate improved model optimization and performance. The extensive experiments conducted illustrate iLLaMA's effectiveness in various tasks, suggesting it could contribute to standardization in AI models for text and image processing.
weaknesses: - Unfortunately for the authors, the existence of VisionLLaMA, which shares similar objectives, significantly diminishes the novelty and contribution of this work. While it should be taken into account that these studies were conducted contemporaneously, the paper still shows considerable weaknesses in comparison. Even taking into account the differences outlined in Section 2.2, it is difficult to argue that this work creates additional value. The claimed distinctions do not appear substantial enough to establish its unique contribution to the field.

- Questionable Design Choices:
    - The adoption of causal mode attention appears primarily motivated by mimicking LLaMA's architecture rather than addressing specific technical needs.
    - The computational benefits of causal attention seems minimal, as evidenced in Table 2.
    - The requirement for a soft mask strategy (which uses bidirectional attention) to optimize causal attention's effectiveness seems contradictory to the original design intent.
    - The paper's emphasis on decoder-only architecture lacks convincing justification, especially given VisionLLaMA's successful use of an encoder architecture.
- Performance and Efficiency Concerns:
    - The model's claimed superiority relies heavily on larger datasets, IN-21K, making direct comparisons to VisionLLaMA problematic.
    - Under equivalent training datasets and model scale, VisionLLaMA demonstrates superior performance.
    - Moreover, VisionLLaMA achieves better efficiency with fewer parameters than iLLaMA.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 4
questions: 1. Will you plan to release the code ?
soundness: 3
strengths: 1. The paper innovatively adapts the decoder-only Transformer architecture from large language models (LLMs) to the computer vision domain, demonstrating the potential for unified architectures across modalities.
2. It introduces a post-sequence class token (PS [cls]) technique to address the attention collapse issue in causal self-attention, which is crucial for the effective training of the proposed iLLaMA model.
3. The soft mask strategy developed in the paper gradually introduces causal masking during training, which improves optimization behavior and mitigates underfitting concerns.
4. iLLaMA achieves competitive performance compared to encoder-only models on ImageNet with significantly fewer parameters, showcasing efficiency and effectiveness.
5. The paper provides a comprehensive evaluation of iLLaMA, including not only accuracy but also calibration, shape-texture bias, quantization compatibility, and transfer learning performance, demonstrating its robustness across various practical metrics.
6. The proposed model scales well with increased model capacity and dataset size, achieving state-of-the-art performance on ImageNet with minimal parameters, highlighting its scalability.
summary: The paper explores the adaptation of decoder-only Transformers, specifically LLaMA, originally designed for large language models (LLMs), to the field of computer vision.
It proposes post-sequence class token to overcome the attention collapse issue when applying causal mode attention, and introduces a soft mask strategy to improve model training behavior.
The paper presents a novel approach to adapting large language model architectures for computer vision tasks, achieving promising results and opening avenues for future research in unified vision and language models.
weaknesses: 1. While the paper discusses the potential of unified architectures for language and vision models, it does not explore multimodal tasks that involve both text and images. This is a significant limitation since multimodal capabilities are crucial for assessing the true potential of unified models.

---

confidence: 5
contribution: 1
initial_rating: 3
presentation: 2
questions: 1. How does iLLaMA’s causal self-attention compare to traditional bi-directional attention in fine-grained object detection tasks, where spatial dependencies might be critical?

2. Would incorporating an explicit multimodal training objective enhance iLLaMA’s ability to handle both image and language tasks in a unified model?

3. Why should we incur additional costs to adapt a decoder-only architecture for image encoding when encoder-only and decoder-only architectures share the same attention method? This clarification could help readers understand the trade-offs of aligning image models with language-model structures, particularly if the benefits of causal attention are marginal compared to bi-directional approaches in vision tasks.
soundness: 2
strengths: 1. Architectural Alignment: Adapting LLaMA’s decoder architecture to image classification, the authors propose a novel integration that demonstrates the feasibility of using language-model structures for visual tasks.

2. Training Techniques: The introduction of PS [cls] and a soft mask to address the attention collapse issue in causal attention is well-motivated and effective in stabilizing training.

3.Extensive Validation: The model is evaluated on multiple properties, including calibration and shape-texture bias, showing robust performance across different visual tasks.
summary: The paper introduces iLLaMA, an adaptation of the LLaMA decoder model to function as a Vision Transformer (ViT) for image classification. iLLaMA addresses the inherent challenges of causal self-attention in image tasks by introducing the post-sequence class token (PS [cls]) and a soft mask strategy. These techniques enhance iLLaMA’s stability and performance, yielding competitive results on ImageNet with relatively few parameters. The paper validates iLLaMA’s capabilities across various tasks, such as calibration, shape-texture bias, and quantization compatibility, positioning it as a streamlined alternative to encoder-only models for visual processing.
weaknesses: 1. Limited Scope in Multimodal Tasks: Although aligned with LLaMA’s architecture, iLLaMA primarily focuses on vision tasks, and the paper does not explore its application to more complex multimodal scenarios.

2. Limited Performance Improvement Over ViT: Although iLLaMA introduces novel adaptations for causal attention in visual tasks, the model does not show a substantial performance enhancement compared to ViT on key benchmarks. This raises questions about the practical gains of adapting a decoder-only architecture for image classification.

3.Computational Demands in Soft Masking: While effective, the soft mask strategy could still present some overhead in training time for larger models, which may limit scalability.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 4
questions: Despite the paper's significant novelty and the innovative concept of utilizing Large Language Models (LLMs) as pre-trained vision learners, there are several areas that would benefit from further clarification and additional experimentation:

Clarification on Attention Collapse: The paper addresses the issue of attention collapse within vision transformers, but it would be beneficial to elaborate on this phenomenon. Specifically, is attention collapse analogous to mode collapse in statistical machine learning, where a model generates limited diversity in outputs despite diverse inputs? A detailed explanation would help readers, especially those new to this area, understand how attention collapse affects model performance and how the proposed Patch Selection [CLS] (PS [CLS]) technique mitigates this issue.

Transparency of Experimental Results: Regarding Table 6, it is not entirely clear whether all the experimental results were obtained directly by the authors or if some results were sourced from existing literature. Clarifying which results are original and which are replicated or cited from other works would enhance the credibility of the empirical evaluation. This transparency is crucial for readers to accurately assess the contributions and validity of the experimental comparisons.

Visualization of Attention Maps: Although the authors demonstrate significant improvements over previous baselines like Vision Transformers (ViT) and ResNet, the attention maps presented in Figure 12 of the supplementary material appear faded in the case of iLLaMA. This raises questions about the model's ability to focus on salient image regions. Providing clearer visualizations or explaining why the attention maps are less pronounced would help in understanding the effectiveness of the PS [CLS] technique in enhancing the attention mechanism within the model.

While the paper excels in presenting architectural innovations and is well-articulated with self-explanatory diagrams, addressing these questions would substantially strengthen its contributions. Additionally, incorporating empirical experiments with newer architectures such as Extended LSTMs, MAMBA, or QK-Normalized architectures could provide a more comprehensive evaluation. These comparisons would help in assessing the true strengths and limitations of iLLaMA relative to cutting-edge models in the field.

By providing detailed answers to these queries and expanding the empirical analysis, the paper would improve in both clarity and scientific rigor. The novelty of repurposing LLMs as pre-trained vision learners is a significant leap forward, and addressing these points would further solidify the paper's impact. I encourage the authors to consider these suggestions, as they would enhance the overall quality and reception of the work.
soundness: 3
strengths: The paper presents a novel methodology for integrating Large Language Model (LLM) architectures into the visual modality, addressing significant challenges in the process. The authors introduce the Patch Selection [CLS] (PS [CLS]) technique, which effectively mitigates the issue of attention collapse commonly observed in vision transformers. This problem arises when the attention mechanism disproportionately focuses on certain regions, leading to suboptimal feature representation. The PS [CLS] technique enhances the model's ability to attend to diverse and informative regions of an image by refining the attention distribution across patches.

To empirically validate their approach, the authors adapt the LLaMA model—a powerful LLM known for its capabilities in language understanding—for image processing tasks, resulting in the iLLaMA model. They pre-train iLLaMA on the extensive ImageNet dataset to ensure robust feature learning. The model is then evaluated across multiple benchmark datasets, where it demonstrates superior performance compared to existing state-of-the-art models. The experiments cover various aspects of computer vision tasks, including image classification, object detection, and segmentation, showcasing the versatility and effectiveness of the proposed approach.

The paper's architectural novelty lies in the seamless integration of LLM architectures with vision transformers through the PS [CLS] technique. This integration allows the model to leverage the strengths of LLMs—such as advanced sequence modeling and contextual understanding—in processing visual data. The authors provide a comprehensive analysis of the model's architecture, detailing how the PS [CLS] technique operates within the transformer framework to enhance attention mechanisms.
summary: In this paper, the authors present a novel approach for adapting Large Language Model (LLM) architectures to the visual modality. They introduce the Patch Selection [CLS] (PS [CLS]) technique to address the issue of attention collapse commonly observed in vision transformers. This technique enhances the model's ability to focus on salient image regions by refining the attention mechanism within the transformer architecture. To validate their methodology, the authors pre-train iLLaMA, an adaptation of the LLaMA model tailored for image processing, on the ImageNet dataset. They conduct extensive experiments to benchmark iLLaMA's performance across various datasets. The results demonstrate that iLLaMA, equipped with the PS [CLS] technique, achieves superior performance compared to existing models, thereby showcasing the effectiveness of their proposed approach in advancing computer vision tasks.
weaknesses: Despite the paper's substantial novelty and the thorough explanation of the architectural innovations, there is a noticeable gap in the empirical evaluation section. Specifically, the paper does not include comparative experiments with newer methods such as Extended LSTMs, MAMBA, or QK-normalized architectures. These models represent significant advancements in the field and are known for their effectiveness in handling complex visual tasks.

The omission of these comparisons makes it challenging to fully assess the strength and generalizability of the proposed iLLaMA model. Without empirical evidence demonstrating how iLLaMA performs relative to these cutting-edge architectures, it's difficult to quantify its advantages or identify potential areas where it may fall short. Including experiments that benchmark iLLaMA against these models would provide a more comprehensive understanding of its performance and robustness.

i feel there should be substantial Experiment of how the behavior of method changes if one to choose the like Mixture of Exparts or Mixtures of base.

Meta-review: This paper presents an interesting idea that makes use of the architecture of the LLaMA model (a language model) as a vision transformer. For this purpose, the LLaMA architecture is step-by-step adjusted to fit the vision data (see Figure 1, the right-hand side). While the final architecture obtains acceptable results on ImageNet-1K, the accuracy is still inferior to vision architectures (e.g. Swin transformer). Adding ImageNet-21K for pre-training improves accuracy, but the numbers are not fairly compared against recent models. In summary, the adapted model does not show a clear benefit in visual understanding. This paper also has other drawbacks, such as (1) the performance on downstream tasks is weak, (2) the idea is closely related to VisionLLaMA, a published ECCV'24 paper, and (3) the fitness of vision and language data -- note that the architecture still needs extensive visual pre-training, and loading language pre-trained models does not help.

The reviewers gave diverse scores (3/5/8/8). While The AC reads the paper and all discussions carefully, and tends to concur with the major concerns of the reviewers. The method is inspired by the need for unifying vision and language understanding, but it is doubtful whether the architecture needs to be unified in this way. Also, as Reviewer **6zNt** pointed out, some statements made by the authors are not convincing enough. All of the above makes the final conclusion of the paper (in particular, the necessity of designing and training such a model) doubtful. Summarizing the above, although the paper is interesting, the clear weakness makes it difficult for the AC to recommend acceptance.

Decision: Reject