Structured Initialization for Attention in Vision Transformers

Authors: Jianqiao Zheng, Xueqian Li, Simon Lucey
Primary area: learning theory
Keywords: Transformer, Learning theory, Initialization, ConvMixer, Attention map

Abstract:
The application of Vision Transformers (ViTs) to new domains where an inductive bias is known but only small datasets are available to train upon is a growing area of interest.
However, training ViT networks on small-scale datasets poses a significant challenge. 
In contrast, Convolutional Neural Networks (CNNs) have an architectural inductive bias enabling them to perform well on such problems. 
In this paper, we propose that the architectural bias inherent to CNNs can be reinterpreted as an initialization bias within ViT. 
Specifically, based on our theoretical findings that the convolutional structures of CNNs allow random impulse filters to achieve performance comparable to their learned counterparts, we design a ``structured initialization'' for ViT with optimization.
Unlike conventional initialization methods for ViTs, which typically (1) rely on empirical results such as attention weights in pretrained models, (2) focus on the distribution of the attention weights, resulting in unstructured attention maps, our approach is grounded in a solid theoretical analysis, and builds structured attention maps.
This key difference in the attention map empowers ViTs to perform equally well on small-scale problems while preserving their structural flexibility for large-scale applications.
We show that our method achieves significant performance improvements over conventional ViT initialization methods across numerous small-scale benchmarks including CIFAR-10, CIFAR-100, and SVHN, while maintaining on-par if not better performance on large-scale datasets such as ImageNet-1K.

Decision: reject
Reviewer scores: 5, 5, 3

Reviews and discussion:
<thread>
creation date:
1729529714868

thread_type:
official_review

<thread_object>
creator: reviewer_tf2y

<content>
summary:
This paper introduces a novel method for improving the performance of ViTs when trained on small-scale datasets by incorporating structured initialization. The authors identify that ViTs struggle with small datasets compared to CNNs, which benefit from inherent inductive biases. By reinterpreting the architectural bias in CNNs as an initialization bias for ViTs, the authors propose a "structured initialization" method that results in structured attention maps for ViTs. The key contribution lies in the use of random convolutional impulse filters to guide the initialization process. The method is theoretically justified and empirically validated across serveral benchmarks. The paper demonstrates that structured initialization yields performance improvements on small datasets without compromising ViT’s flexibility on larger datasets.

strengths:
Structured architecture achieves good performance across both small and large datasets, which demonstrates its scalability and flexibility.

weaknesses:
1. The core argument of the method is that the convolutional structure can be transferred to the attention mechanism in transformers by initializing the attention maps with random impulse filters. However, this analogy between convolutional layers in CNNs and the attention mechanism in ViTs may be overly simplistic. CNNs' convolutional filters are spatially local and fixed in structure, while attention in ViTs is meant to capture long-range dependencies and is more flexible. This difference is crucial, and the method does not seem to fully address how imposing a rigid, convolution-like structure at initialization aligns with the flexibility that the attention mechanism needs. The convolution structure might limit the model's ability to learn long-range dependencies that are essential to the transformer. The claim that random impulse filters can replace learned convolutional filters is somewhat true for CNNs under certain conditions (like ConvMixers), but applying this to ViTs is more challenging. The attention mechanism is a more complex and dynamic operation compared to convolutions, and it’s unclear if the same approximation can hold. In practice, imposing a convolution-like structure might hinder the attention mechanism's ability to adapt during training.

2. The paper proposes that impulse filters, combined with the softmax operation, can initialize the attention maps. The softmax function ensures that all outputs are non-negative, which is a crucial difference from convolutions, which can have both positive and negative values. Random convolutional filters may contain both positive and negative values, while softmax output does not. This inconsistency could cause issues. The authors acknowledge this (stating the filters must be positive), but they do not provide a deep exploration of how this might affect the quality or flexibility of the learned attention maps. Relying on impulse filters could reduce the model's expressivity, especially if the initialized filters are too rigid and only positive-valued patterns are learned initially.

3. The authors propose an iterative optimization process to solve for the initial values of $Q_{init}$ and $K_{init}$ such that the resulting attention maps resemble impulse convolution filters. The optimization is based on a pseudo-input, which is generated from positional encodings rather than actual data. This could introduce an unwanted bias into the model's initial learning process. While using positional encoding as pseudo-input is an interesting idea, the paper does not adequately explore how different choices of pseudo-inputs affect the results or whether using actual training data for initialization would be a better alternative.

4. The optimization process described is more computationally expensive (up to 10,000 iterations using Adam) compared with traditional initialization methods. This added complexity raises the question of whether the benefits of structured initialization outweigh the cost, especially given that the improvements on large datasets are marginal. There is no discussion of the computational cost vs. benefit of this method compared to standard initialization techniques.

5. The use of random impulse convolution filters assumes that locality is always important, but this assumption may not hold in tasks where global context is critical. In CNNs, locality is useful because of the hierarchical structure of learned features. However, in transformers, the attention mechanism is specifically designed to handle long-range dependencies. By forcing the model to start with local dependencies (via impulse filters), the authors may inadvertently restrict the model's ability to learn global features early in training, leading to potential issues in tasks where global context is key from the beginning.

6. The method relies on several hyperparameters, including filter size (3x3 or 5x5) and the number of iterations for optimization. However, the choice of these hyperparameters is not adequately justified or explored. The method's performance is likely sensitive to these parameters, but there is no thorough analysis of how variations in filter size or optimization parameters affect results. Given the complexity of the initialization process, these aspects should have been investigated in detail to ensure the method’s robustness.

questions:
1. Why Impulse Filters? The reasoning behind choosing impulse filters (instead of other structured filters) for initializing attention maps could be explained in more detail. Are impulse filters the best possible choice, or could other filter types provide better generalization?

2. How does this structured initialization affect the deeper layers of ViTs after finetuning? Figure 3 is quite interesting. Do the constraints imposed by impulse filters affect the deeper layers’ ability to fine-tune long-range dependencies? The paper does not discuss the long-term effects of this initialization on the network’s convergence.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730642647866

thread_type:
official_review

<thread_object>
creator: reviewer_awu8

<content>
summary:
The paper presents an approach to initialize ViT through structured initialization of attention maps. By incorporating CNN-like inductive biases during initialization, it aims to combine the local spatial processing capabilities with the global relationship learning of attention mechanisms and take advantage of CNNs' inductive bias. Experimental results on several small-scale datasets validate its effectiveness.

strengths:
- The paper is easy to follow.
- As much work has been trying to introduce convolutional design into the ViT model, this paper provides an interesting viewpoint that initializing the attention map as CNNs can also help to introduce the inductive bias and subsequentially improve the performance of trained ViT on small-scale datasets.
- A theoretical explanation is provided to show the connection between the structural initialization in ViT and inductive bias in CNNs.
- Some special designs like more heads and various initialization conv kernel sizes are adopted.

weaknesses:
- The fundamental approach of forcing attention maps to mimic convolutional kernels seems to contradict the core advantage of attention mechanisms, as their advantage is to learn flexible, dynamic global relationships. It would be better to justify why structured initialization is preferred over simply incorporating convolutional blocks into the architecture, which would be a more straightforward solution. 
- It would be better to provide more analysis of why this approach is better compared to well-established solutions: 
  - Transfer learning from large-scale pre-trained ViTs
  - Hybrid architectures combining convolution and attention
- The optimization process required for initializing attention maps introduces additional computational overhead during training, and one needs to further choose the optimizer for initialization and conv kernel size (as well as other hyperparameters for different model sizes), which makes it impractical.

questions:
-  Does the structured initialization limit the model's ability to learn better representation? It would be better to provide some representation-level analysis using metrics like Centered Kernel Alignment (CKA) similarity [1].

[1] Kornblith, Simon et al. “Similarity of Neural Network Representations Revisited.” ICML 2019.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730742848304

thread_type:
official_review

<thread_object>
creator: reviewer_p9fk

<content>
summary:
This paper addresses the challenge of applying Vision Transformers (ViTs) to new domains with small datasets, where Convolutional Neural Networks (CNNs) typically excel due to their inherent architectural inductive bias. The authors propose a novel approach that reinterprets CNN's architectural bias as an initialization bias for ViTs, termed "structured initialization." Unlike traditional ViT initialization methods that rely on empirical results or attention weight distributions, this method is theoretically grounded and constructs structured attention maps. The paper demonstrates that this structured initialization enables ViTs to achieve performance comparable to CNNs on small-scale datasets while retaining the flexibility to perform well on larger-scale applications. The proposed method shows significant improvements over conventional ViT initialization across several small-scale benchmarks, including CIFAR-10, CIFAR-100, and SVHN, and maintains competitive performance on large-scale datasets like ImageNet-1K.

strengths:
1.Theoretical Foundation: The structured initialization method is based on solid theoretical analysis rather than just empirical results, providing a strong rationale for its effectiveness.

2.Performance Improvements: The method consistently shows significant performance improvements over conventional ViT initialization methods in small-scale datasets, which is a notable achievement.

weaknesses:
1. In terms of innovation, the Transformer architecture was initially designed to minimize inductive bias. The author's attempt to incorporate structural biases from CNNs into the Transformer seems to go against the original intent of the Transformer design, which could be seen as a step backward for the evolution of Transformer models.

2. The variety of experimental backbones is somewhat limited. It would be beneficial to conduct experiments with DeiT or Swin-Transformer to compare results. Furthermore, aside from classification tasks, it would be interesting to test the method on detection or segmentation tasks to further evaluate its versatility and effectiveness.

questions:
Why not apply structured initialization to the value (V) component of self-attention? Additionally, how are the feed-forward network (FFN) layers, normalization layers, and projection layers initialized in the proposed method?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732053321484

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
A common response to all reviewers

comment:
We thank all reviewers for the thoughtful comments and questions. We would like to make a common response to clarify the advantage and innovation of our method.

**Q: Does introducing CNN structural biases to Transformer contradict the advantage of its structure?**

**A:** We would like to clarify that our method focuses on introducing the **architectural inductive bias** from CNNs to initialize the attention map **without changing the Transformer architectures**. We have emphasized this argument in the abstract and introduction, stating that unlike previous arts [c1, c2, c3] that directly introduce convolutions into attention, potentially damaging the structural advantages of Transformers, our method only introduces architectural inductive bias in attention map initialization, **maintains the inherent structural flexibility in ViTs**. Therefore, our method **preserves the Transformer architectures** and still allows the ViTs to learn flexible, dynamic global relationships. This is also one of the central innovations of our method. In addition, we have also provided experiments on large-scale applications to validate the stable performance of our method that preserves the architectural flexibility of ViTs.

References (also cited in our paper):

[c1] Yuan, Kun, et al. "Incorporating convolution designs into visual transformers." Proceedings of the IEEE/CVF international conference on computer vision. 2021.

[c2] Li, Kunchang, et al. "Uniformer: Unifying convolution and self-attention for visual recognition." IEEE Transactions on Pattern Analysis and Machine Intelligence 45.10 (2023): 12581-12600.

[c3] Dai, Zihang, et al. "Coatnet: Marrying convolution and attention for all data sizes." Advances in neural information processing systems 34 (2021): 3965-3977.


**Q: Does pre-optimization add computation overhead?**

**A:** We kindly refer the reviewer to Section 4.2 Line 289-291, where we stated that our pre-optimization is **not a pretrained step since no real data is involved**. Also, our optimization algorithm is very simple, and it serves as a surrogate for the SVD solver, converging in just a few seconds, *i.e.*, **around 5 seconds**. In addition, it is worth noting that for each embedding dimension and kernel size, the initialization only needs to be calculated once. As a simple optimization for the SVD solver, the optimization will converge, and no careful choice of hyperparameters such as different optimizers or kernel size is needed.


We have also provided individual responses below. If you have further questions, please give us comments and we are happy to discuss.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732437221725

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
A revised version of the paper is submitted

comment:
We have submitted a revised version with all changes marked in red. Specifically, we have emphasized our arguments in the introduction and contribution sections. Additional figures of attention maps before and after training are added to Appendix Figure 6 for further clarification.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733109410441

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
An implementation of our method shared with all reviewers

comment:
Dear reviewers, we have shared an anonymous Colab link with an implementation of our method here: [https://colab.research.google.com/drive/1Oic07RhIbqKUcVoCQnd2A2c73LBtdSq6?usp=sharing](https://colab.research.google.com/drive/1Oic07RhIbqKUcVoCQnd2A2c73LBtdSq6?usp=sharing). 

In this notebook, we have provided all training details for the three main initialization methods that are compared in our paper: **random** ([timm default](https://github.com/huggingface/pytorch-image-models)) initialization, **mimetic** initialization ([Trockman & Kolte](https://proceedings.mlr.press/v202/trockman23a/trockman23a.pdf)), and our **structured initialization** ([ours](https://openreview.net/pdf?id=z9UBpl4pv5)) method.

We only trained for **20 epochs** to save time. Note that in our paper (real training), we train for **200 epochs** with **all the other settings (hyperparameters) *the same as* the settings defined in this notebook**.

We have also included a summary of (1) initialization (pre-optimization) time, (2) pre-training attention maps, and (3) post-training attention maps. We hope it will further clarify the concerns that reviewers have raised.

Please also note that the initialization time is directly computed using the **Colab T4 GPU**, therefore, it is different from what we have reported in our paper (**$\sim$5 seconds** for our initialization method). Nonetheless, the pre-optimization time (in seconds) is neglectable compared to the long training time which can last for hours or even a few days.

Please let us know if you have further questions or concerns.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes structured initialization to improve the performance of Vision Transformers (ViTs) on small-scale datasets. The approach reinterprets architectural bias from CNN as an initialization bias for ViTs. This method enables ViTs to achieve competitive performance with CNNs on small datasets, such as CIFAR-10, CIFAR-100, and SVHN, while maintaining flexibility and strong performance on larger datasets like ImageNet-1K. The paper provides both theoretical justification and empirical validation, demonstrating improvements over conventional ViT initialization methods.

After the review process, all reviewers unanimously voted for rejection, citing several concerns, including limited experiments, marginal improvements on large datasets, and the generalizability of the approach to other datasets.