Title: Noise Stability of Transformer Models
Primary area: foundation or frontier models, including LLMs
Keywords: transformers, simplicity bias, noise stability, regularization methods, spectral concentration

Abstract: Understanding simplicity biases in deep learning offers a promising path toward developing reliable AI. A common metric for this, inspired by Boolean function analysis, is average sensitivity, which captures a model's robustness to single-token perturbations. We argue that average sensitivity has two key limitations: it lacks a natural generalization to real-valued domains and fails to explain the "junta-like" input dependence we empirically observe in modern LLMs. To address these limitations, we propose *noise stability* as a more comprehensive simplicity metric. Noise stability expresses a model's robustness to correlated noise applied to *all* input coordinates simultaneously. We provide a theoretical analysis of noise stability for single-layer attention and ReLU MLP layers and tackle the multi-layer propagation problem with a covariance interval propagation approach. Building on this theory, we develop a practical *noise stability regularization* method. Experiments on algorithmic tasks show that our regularizer consistently catalyzes grokking and accelerates training by approximately $35$\%. Our results establish noise stability as a powerful tool for understanding and improving modern Transformers.

TL;DR: We introduce noise stability in Transformer models as an alternative proxy for explaining simplicity bias and propose a corresponding regularization method that we observe accelerates grokking.