Linearly Decomposing and Recomposing Vision Transformers for Diverse-Scale Models

To W1 & Q1: Details of decomposition

Thank you for your question. Due to the parameter gap, the 12-layer model with only one component $A_1$ in the first stage does not perform as well as a typical 12-layer model. However, as more components are trained, the performance becomes comparable to the model with the same parameters.

As shown in Fig.C of the PDF, using DeiT-Base as an example, ${W_l}$ are the layer parameters of the model and are constructed by trainable parameters ${A_n}$ (include all modules in a ViT layer like MSA, MLP). In the 1-st stage, each layer $W_l$ initially contains one component, $A_1$, shared across all layers, with coefficients calculated using a predefined polynomial (Eq.3). The number of learnable parameters in the 12-layer model is approximately equal to $A_1$, achieving 76.6% performance, as in Fig.4(a).

In the 2-nd stage, a new component, $A_2$, is added to each layer, and $A_1$ and $A_2$ are combined using their respective coefficients to initialize each layer for forward propagation as Eq.1. During training, the parameters of $A_1$ remain fixed, and only the newly added component $A_2$ is updated. In this way, the learnable parameters still equals to the parameters of one component. This process is repeated until all components are fully trained.

As the number of added components increases, the performance of the 12-layer model improves, reaching 83.70% with 12 components, which is comparable to DeiT-Base trained as a whole, with a performance of 83.51% under the same training conditions and parameters. This confirms our assumption that each ViT layer can be linearly decomposed into a shared component {$\{A_1, \ldots, A_n\}$}, and that a model composed of components {$\{A_1, \ldots, A_k\}$} can perform comparably to the full model $W$ under the same parameters as in Fig.6. We will include these details in the revision for clarity. Should you have any further suggestions, we would be happy to address them.

To W2: Details of recomposition

After all components are trained, the components can be directly used to initialize model without training in the recomposition. In Fig.D of the PDF, we present two examples. Regardless of the training method, different model structures are initialized by Eq.9. Under the "without constraint" training method, Eq.9 is only used to initialize $\{W_l\}$; the trainable parameters of each layer ($\{W_l\}$) are updated independently during training, the same as the normal ViT fine-tuning. The "with constraint" training method decomposed the $\{W_l\}$ to trainable parameters $\{A_k\}$ while keeping the polynomial coefficients fixed using Eq.9. This means that, during the gradient back-propagation, the gradient is backed to $\{A_k\}$ instead of $\{W_l\}$. We hope this explanation clarifies the reconstruction, and we will include these details in the revision.

To W3 & Q3: Empirical design of components & latency improvement

Thank you for your questions. In the decomposition and recomposition, all sub-modules within a ViT layer are treated uniformly. Therefore, in our paper, each component $A_n$ includes all sub-modules within a ViT layer $W$, such as MSA, MLP, FFN, and others. As in Fig.C of PDF, in the different stages of decomposition, the parameters of each layer are obtained by Eq.1. In the recomposition, as in Fig.D, ViT models of varying layer can be initialized with different numbers of components, with each layer’s parameters initialized by Eq.9.

Regarding the latency improvements for different scales, in recomposition, models of different sizes can be initialized with components by Eq.9, without the need for repeated training as required in knowledge distillation, which contributes to our latency improvement. In Sec.4.1.2 "Computing Resources", we compare the computational savings of our method against other model compression methods, highlighting the efficiency of our approach.

To W4: Literature review

We appreciate the suggestion to expand the literature review. Currently our literature review includes model initialization and model compression (Micro-architecture design, Model pruning, Quantization, and Knowledge distillation). In the revision, we will also provide a detailed discussion of other parameter-efficient paradigms, including Low-Rank Factorization (LRF) and Parameter Sharing (PS). For LRK and PS, generating models of different sizes typically requires redesigning the factorization or sharing scheme for each target model size, which can be complex and resource-intensive. In contrast, our approach allows flexible and economical generation of models with different sizes without repeated training or redesign.

To W5: Writing

Thank you for the feedback. We'll improve the revision by enhancing writing quality, correcting typos, and refining formatting for a more polished presentation.

To Q2: Details of Fig.2

Thank you for your question. In Fig.2, the x-axis is the layer index of ViT, while the y-axis is the values of the parameters in each layer after performing PCA, reduced to one dimension (details are in the Suppl Material Sec.A.1). The figure shows a linear relationship between these two variables. This observation suggests that the parameters across layers in a well-trained transformer are not independent variables, but may be linearly related in a higher dimensional space, possibly correlated with the layer index. This implies that the parameters can be reconstructed using the form given in Eq.1, enabling flexible initialization for models of different scales. We will clarify this point in the revision to enhance understanding.

To Q4: Experimental design of decomposition training

Thank you. You've raised an important point; ideally, components would be added and trained one at a time, but that's time consuming. To balance efficiency and performance, we trained multiple components in stages. We will clarify this point in the revision to ensure better understanding.

Thank you for your insightful comments and for pointing out the relevance of the two seminal works. We agree that these works are indeed crucial in the context of parameter-efficient learning. E^2VPT introduces a method for efficient fine-tuning of large-scale transformer models through learnable prompts and pruning techniques, reducing the number of tunable parameters while maintaining performance. Facing the Elephant in the Room analyzes when and why Visual Prompt Tuning (VPT) is effective, offering insights into its applicability depending on task objectives and data distributions. Both are closely aligned with our approach to parameter-efficient learning, and we will include discussions of these works in our revision to provide a more comprehensive overview of existing approaches in the field of parameter-efficient learning.

To 1: Components sharing approach

Thank you for your insightful question. The component sharing in our model is applied within each individual layer, unlike the cross-layer sharing seen in MiniViT. Additionally, the forward process follows the typical ViT architecture, such as DeiT, during training. In the revision, we will include more detailed explanations to clarify these points.

To 2: Training approaches in recomposition

Thank you for your question. For the two training approaches, the initialization methods used are the same. The key difference lies in how the parameters are optimized during training. In the "without constraint" approach, Eq.9 is only used to provide an initialization of layer parameters $W_l$ for ViTs. During training, the parameters of each layer $W_l$ are then updated independently, the same as the normal ViT fine-tuning, without being constrained by Eq.9. In contrast, the "with constraint" approach optimizes the component parameters directly, rather than the individual layer parameters $W_l$. This means that during gradient back-propagation, the gradient is applied to ${A_k}$ instead of ${W_l}$. Therefore, the parameters of each layer remain constrained by Eq.9, with components shared across all layers. We will clarify this further in the revision.

To 1: Swin Transformers and MobileViT

Thank you for your question, which inspired us to explore the applicability of our method to Swin Transformers and MobileViT. Swin Transformers consist of four stages, where different stages have transformer layers with different resolutions, resulting in different components decomposed from layers across different stages. Due to time constraints, we focused on decomposing the layers in stage 3 because it has a larger number of layers. We visualized the relationship between the layer parameters and the corresponding layer index in stage 3. In Fig.A(a) and Fig.A(b) of PDF, we present the results for Swin-T and Swin-B, respectively. We observed a linear relationship between the parameters and the layer index, similar to that shown in Fig.2 in paper for plain transformers. Therefore, we adapted our method for Swin transformers by treating stage 3 as a plain transformer during decomposition. We have successfully decomposed one component from stage 3 of the Swin-Tiny model, achieving a result of 77.61%. This is close to the Swin-Tiny's performance on ImageNet-1k, which is 81.2%, and the initial results are promising. We will include the full experimental design and results in the revision.

For MobileViT, in Fig.A(c) we also visualize the transformer layers across different ViT blocks. We observed that within each block, there is a linear relationship between the layer parameters and the layer index. However, since MobileViT has relatively fewer parameters, it may not be economical and efficient to decompose it. Given the linear relationship, it is still possible that the transformer layers in MobileViT could be initialized with pre-trained components. We will perform further experiments to validate this possibility and incorporate these results in revision.

To 2: More comparative experiments

Thank you for your question. For decomposition, as in Fig.4(a), our decomposed models outperform the original DeiT-B model with the same number of parameters under the same training conditions. Fig.4(b) further shows that our decomposition approach leads to better performing models in different parameter regimes compared to the SN-Net family.

For the recomposition, we compared our recomposed models to different-sized DeiT models and other compression methods such as Mini-DeiT-B and Efficient-ViT. The parameters for these models were downloaded from their official sources and fine-tuned under the same training epochs as ours over 9 downstream tasks. The results in Fig.6 demonstrate the effectiveness of our method. In addition, we provide detailed numerical performance in the supplementary material (Tab.2).

Furthermore, we conducted additional comparison experiments with the 12-layer ViT, based on the DeiT-Base structure, initialized with 4 components and trained using the "with constraint" method (28.7M parameters), compared to the pre-trained Swin-Tiny (28M parameters), as shown in Tab.A of the PDF. The results indicate that the model initialized with components performs comparably to the pre-trained Swin-Tiny on downstream tasks with a similar number of parameters. This further demonstrates the feasibility of using components to initialize ViTs of different scales. We will provide the full experimental results in the revision.

To 1: Details of recomposition

Thank you for your feedback. In the revision, we will provide a more detailed explanation of the recomposition process and update Fig.3 in the paper to illustrate the initialization of different model architectures and the various training methods. In Fig.D of the PDF, we provide two examples. At the top of Fig.D, a 2-layer ViT is initialized with 3 components, and at the bottom, a 3-layer ViT is initialized with 4 components. Regardless of the training method, different model structures are initialized using Eq.9. As shown in Fig.D, for the "without constraint" training, Eq.9 is used solely for initialization. During training, the parameters of each layer ($W_l$) are updated independently, similar to a typical ViT, without adhering to Eq.9. In contrast, the "with constraint" training method updates the parameters of the component $A_n$ during training and maintains the polynomial coefficients unchanged, meaning the components remain shared across all layers, and each layer's parameters continue to satisfy Eq.9.

To 2: More linear relationship examples for other models

Thank you for the suggestion. The supplementary material provides results for MoCo and DINO models using the same process as described in the paper. Additionally, in Fig.B, we show the relationship between each layer's parameters and the corresponding layer index for MAE-B[1], BEITv2-B[2], and RoBERTa-B[3] models. It is evident that these models also exhibit a linear relationship between the parameters and the layer index. In the revision, we will include these results to further demonstrate the generalizability of this pattern.

To 3: Chebyshev polynomials

Thank you for your observation. In Tab.1 of paper, our experimental results demonstrate that Chebyshev polynomials outperform other polynomials in recomposition. This may be because Chebyshev polynomials are known to provide the best uniform approximation of continuous functions [4]. And each layer of the ViT model can be viewed as a function, the parameter matrix of each layer can be decomposed into components and represented as a weighted sum using Chebyshev polynomials. In the revision, we will expand on this theoretical motivation and explicitly highlight the desirable mathematical properties of Chebyshev polynomials.

[1]Masked Autoencoders Are Scalable Vision Learners.

[2]BEIT: BERT Pre-Training of Image Transformers

[3]RoBERTa: A Robustly Optimized BERT Pretraining Approach

[4]Trefethen, L. N. (2013). Approximation Theory and Approximation Practice. SIAM.