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

Baiyuan Chen

Keywords: Feature Learning, Attention, Convolution, Transformer, ResNet, Lipschitz Continuity, Wasserstein Distance, Topological Data Analysis

TL;DR: We theoretically and experimentally prove that attention mechanisms process data in a more compact and stable way than convolutional layers, meaning the outputs are closer together and more robust to different data distributions.

Abstract: Robustness is a crucial attribute of machine learning models, A robust model ensures consistent performance under input corruptions, adversarial attacks, and out-of-distribution data. While the Wasserstein distance is widely used for assessing robustness by quantifying geometric discrepancies between distributions, its application to layer-wise analysis is limited since computing the Wasserstein distance usually involves dimensionality reduction, which is not suitable for models like CNNs that have layers with diverse output dimensions. To address this, we propose $\textit{TopoLip}$, a novel metric that facilitates layer-wise robustness analysis. TopoLip enables theoretical and empirical evaluation of robustness, providing insights into how model parameters influence performance. By comparing Transformers and ResNets, we demonstrate that Transformers are more robust in both theoretical settings and experimental evaluations, particularly in handling corrupted and out-of-distribution data.

Primary Area: learning theory

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Reciprocal Reviewing: I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 14123