Feature Learning in Attention Mechanisms Is More Compact and Stable Than in Convolution

We thank the reviewer for their detailed and constructive feedback, which has helped us refine our manuscript. Below, we address the specific concerns raised:

$**The key weakness**$

After carefully considering the feedback from all reviewers, we decided to focus our study on the robustness of individual architectures, specifically comparing attention-based and convolutional models. In the revised version, we propose a robustness metric that enables both theoretical and empirical evaluation of robustness. This metric enables direct comparisons between the robustness of different model architectures. While the primary focus of the previous version was to compare attention and convolutional layers, in the revised version, the main goal has shifted to introducing and validating the proposed robustness metric. The comparison between attention and convolution serves as a validation of this metric.

$**Major Weaknesses**$

1. In the revised version, we have clarified in the introduction why the proposed metric is needed and how it addresses gaps in existing methods. This discussion helps contextualize the relevance and importance of the work.

2. We provided literature review that is closely related to the topic in the revised version.

3. Thank you for pointing that out. The relationship between variance and the Wasserstein Lipschitzness indeed needs a theoretical analysis to connect them. However, in the revised version, we shifted focus entirely to Wasserstein Lipschitzness and no longer consider variance. The proposed robustness metric is now based solely on Wasserstein Lipschitzness, which is both theoretically and empirically validated.

4. We acknowledge the inconsistencies in the experimental results in the previous version, which may have stemmed from the unclear relationship between variance and Wasserstein Lipschitzness. After refocusing the paper, we provided experimental results that are consistent with the revised theoretical framework.

$**Other weaknesses**$

a. The background material has been shortened in the revised version.

b. The Wasserstein distances reported are computed on persistence diagrams generated using persistence homology. We apologize for not stating this explicitly in the original version. The revised version includes a detailed explanation of how TDA is utilized.

c. Theorem 5 from the previous version has been moved to the appendix as a lemma. The definition of the Wasserstein distance is retained in the main section, as it is essential for introducing Wasserstein Lipschitzness, which is central to our proposed metric.

d. Intrinsic dimension analysis has been removed in the revised version, as it is less relevant to the paper’s new focus.

e. The results in the appendix are now properly referenced and discussed in the main paper.

f. We have included accuracies and losses for all models in the revised manuscript (see Table 2 and Figures 7 and 8).

g. See lines 360-366 in the revised version for detailed explanation.

h. Sorry for the unclearness. Figure 2 shows the Wasserstein distance of models at epoch 1, where models had not yet converged. Since these results do not represent model behavior appropriately, we removed the epoch-1 results and only keep the results that models converge in the revised version.

i. Thank you for pointing that out.

$**Questions**$

$\bullet$ This observation was based on Figure 6 in the previous version. We apologize for not referencing it explicitly. In the revised version, we have ensured that all references to figures and results are clearly stated.

$\bullet$ The (local) intrinsic dimension is computed by calculating the local dimensionality for each point by analyzing its k nearest neighbors and averaging the results to obtain a global estimate. However, as intrinsic dimension is no longer considered in the revised version, kNN methods are not used in this work.

We hope that the revised manuscript addresses the reviewer’s concerns and clarifies the motivation, methodology, and findings of our work.

We thank the reviewer for their detailed and constructive feedback, which has helped us refine our manuscript. Below, we address the specific concerns raised:

$**A takeaway**$

After carefully considering the feedback from all reviewers, we decided to focus our study on the robustness of individual architectures, specifically comparing attention-based and convolutional models. In the revised version, we propose a robustness metric that enables both theoretical and empirical evaluation of robustness. This metric enables direct comparisons between the robustness of different model architectures. While the primary focus of the previous version was to compare attention and convolutional layers, in the revised version, the main goal has shifted to introducing and validating the proposed robustness metric. The comparison between attention and convolution serves as a validation of this metric.

$**Difference between attn -> vit as well as conv -> resnet**$

Yes, the difference is the exclusion of residual connections and normalization layers. In the previous version, we concluded that ViTs have higher Wasserstein Lipschitz numbers than resnet18. However, we only considered resnet18 in the previous version, which uses basic blocks. In larger models such as resnet50 and resnet101, bottleneck blocks are typically used. In the revised version, we consider bottleneck blocks (for better generalization) and resnet50&101, finding that the empirical observations also align with the theoretical results for real-world models (see Section 4.2). It is important to note that the theoretical comparisons in this work are conducted under simplified settings (e.g., most entries follow $\mathcal{N}(0,\sigma^2)$ except for convolutional layers). While these settings provide valuable insights, they may not generalize to certain extreme scenarios. However, an important takeaway from the Lipschitzness bounds is that the Lipschitzness bounds show us the parameter dependencies of the smoothness of models (see l. 299-303).

$**Experimental results**$

We have added the training details in the revised version. The accuracies and losses of the models are also included in the revised version (see Table 2 and Figure 7 and 8).

$**Low or high variance**$

Both low and high variance in the feature space have advantages depending on the context: low variance indicates a more robust model, less sensitive to noise or perturbations, while high variance indicates greater expressiveness, capturing a wide range of input features.

However, in the revised version, we no longer emphasize variance analysis. Instead, we focus on introducing and validating the proposed robustness metric (based on Wasserstein Lipschitness), which better aligns with the paper's main goal.

$**Data visualization**$ Given the shift in the paper’s focus, data visualizations are less central to validating the results, and therefore, we have removed them in the revised version. Nonetheless, we greatly appreciate the reviewer’s suggestions.

We hope that these clarifications and the revisions made to the manuscript address the reviewer’s concerns.

We thank the reviewer for their valuable feedback and comments.

$**Exploration of hybrid models**$

After carefully considering the feedback from all reviewers, we decided to focus our study on the robustness of individual architectures, specifically comparing attention-based and convolutional models. In the revised version, we propose a robustness metric and analyze the differences between these two architectures. Since the primary goal of our work is to demonstrate the robustness properties of attention and convolutional layers independently and validate the proposed metric, exploring hybrid models falls outside the scope of this study. Nonetheless, we appreciate the suggestion and recognize the potential of hybrid models for future investigations.

$**Training on CIFAR-10 may not yield strong performance for ViTs**$

In the revised manuscript, we have included the accuracy and loss of each model (Table 2 and Figure 7,8). ViTs achieve validation accuracies ranging from 77.8% to 87.0%, which may indicate a limitation due to the relatively small dataset size. However, since their performance is reasonably close to ResNets (around 91%), we believe that the ViTs are not significantly under-trained, and their performance suffices for our comparative analysis. Nonetheless, we acknowledge that larger datasets might yield even stronger results for ViTs.

We thank the reviewer for their detailed and constructive feedback, which has helped us refine our manuscript. Below, we address the specific concerns raised:

$**Usefulness of the theory and the significance of the theory**$

After carefully considering the feedback from all reviewers, we decided to focus our study on the robustness of individual architectures, specifically comparing attention-based and convolutional models. In the revised version, we propose a robustness metric that enables both theoretical and empirical evaluation of robustness. This metric facilitates comparisons between different model architectures, offering insights into their robustness properties.

While the primary focus of the previous version was to compare attention and convolutional layers, in the revised version, the main goal has shifted to introducing and validating the proposed robustness metric. The comparison between attention and convolution serves as a validation of this metric. Additionally, we appreciate the reviewer’s suggestion to consider Differentially Private (DP) SGD training and related work. While this is beyond the scope of the current study, we acknowledge its importance and plan to explore such connections in future work.

$**Toy practical applications**$

We train models on CIFAR-10 and conduct robustness experiments on CIFAR-10C. The experimental aligns with the theoretical results, and it validates our method. We appreciate the reviewer’s suggestion regarding practical applications

$**How were the models trained?**$

All models are trained from scratch on CIFAR-10. In our settings, even the smallest ViT (6 heads, 384 embedding dimensions) converges.

Training details: All models are trained on the training dataset of CIFAR-10 (Attns for 100 epochs, Convs for 200 epochs, ResNets for 100 epochs, and ViTs for 200 epochs). After convergence, we switch the models to evaluation mode to freeze their parameters. Then, we input the test dataset and collect the outputs from all layers and compute their persistence diagrams. We have added the training details in the revised version. The accuracies and losses of the models are also included in the revised version (see Table 2 and Figure 7 and 8).

We hope this revised version addresses the reviewer’s concerns and clarifies the motivation, methodology, and significance of our work.

We thank the reviewer for their detailed and constructive feedback, which has helped us refine our manuscript. Below, we address the specific concerns raised:

$**Writing Problems**$

After carefully considering the feedback from all reviewers, we decided to focus our study on the robustness of individual architectures, specifically comparing attention-based and convolutional models. In the revised version, we propose a robustness metric that enables both theoretical and empirical evaluation of robustness. This metric facilitates comparisons between different model architectures, offering insights into their robustness properties. While the primary focus of the previous version was to compare attention and convolutional layers, in the revised version, the main goal has shifted to introducing and validating the proposed robustness metric. The comparison between attention and convolution serves as a validation of this metric.

In this revision, we have made the following key changes:

1. Mask attention: We no longer consider masked attention and instead focus on standard attention mechanisms. Sorry for the misunderstanding.

2. Metric choice: We now use the $W_1$ distance instead of $W_2$, using the fact that for two measures $\mu$ and $\nu$, $W_1(\mu,\nu) \leq W_2(\mu,\nu)$. This allows the computed Lip$^{W_1}$ to directly extend to Lip$^{W_2}$.

3. Background material: The background material has been shortened and integrated into relevant sections. For instance, the discussion of persistent homology has been combined with the introduction of the proposed method.

$**Theoretical contributions**$

$\bullet$ In previous works, the Lipschitz constant of attention includes exponential terms. In our revised version, we compute a tighter bound without exponential terms. See $**Proof of Theorem 1**$ for details.

$\bullet$ We have provided concrete results for ResNets and Transformers in the revised version. See Section 4.2.

$**Experimental contributions**$

$\bullet$ We have streamlined the experimental results to align with the revised manuscript's focus. Specifically, we computed the change rate of the Wasserstein distance of persistence diagrams, which validates the proposed robustness metric. The results indicate that attention-based models are more robust than convolutional models, aligning with our theoretical predictions. Then we conduct robustness experiments using CIFAR-10C to verify the findings.

$\bullet$ Activation variance is computed as described in l. 367-370 of the revised manuscript.

$\bullet$ The practice of stacking attention layers without residual connections is to compare $*TopoLip*$ (the metric proposed in the revised version, which stands for the maximum change rate of the Wasserstein distance). If including residual connection or layer normalization, it will be difficult to measure the Lipschitzness of attention (also the reason for the previous version).

$**Questions**$

1. See the above answer.

2. See Table 2 and Figure 7 and 8 for accuracies and losses of models. Yes, the measure formulation of attention is valid for a finite number of tokens, but the number of tokens should be sufficiently large

3. Yes. However, we only considered resnet18 in the previous version, which uses basic blocks. In larger models such as resnet50 and resnet101, bottleneck blocks are typically used. In the revised version, we consider bottleneck blocks (for better generalization) and resnet50&101, finding that the empirical observations also align with the theoretical results for real-world models (see Section 4.2). It is important to note that the theoretical comparisons in this work are conducted under simplified settings (e.g., most entries follow $\mathcal{N}(0,\sigma^2)$ except for convolutional layers). While these settings provide valuable insights, they may not generalize to certain extreme scenarios. However, an important takeaway from the Lipschitzness bounds is that the Lipschitzness bounds show us the parameter dependencies of the smoothness of models (see l. 299-303).

We hope that these clarifications and the revisions made to the manuscript address the reviewer’s concerns.