Influence Malleability in Linearized Attention: Dual Implications of Non-Convergent NTK Dynamics

1. Problem Statement

The theoretical understanding of attention mechanisms, particularly within the framework of the Neural Tangent Kernel (NTK), remains incomplete. Conventional NTK theory posits that sufficiently wide neural networks operate in a "lazy training" regime (kernel regime), where the kernel remains constant during training, and the network converges to its infinite-width limit.

However, it is unclear whether attention mechanisms follow this trajectory. While attention is empirically successful, its learning dynamics lack rigorous characterization. The central question addressed by this work is: Do linearized attention mechanisms converge to their infinite-width NTK limits, and if not, what are the theoretical and practical implications of this non-convergence?

2. Methodology

The authors employ a combination of theoretical analysis and empirical validation using a parameter-free linearized attention mechanism.

• Linearized Attention Architecture:

  • They define a simplified attention mechanism: $f_{att}(X) = X X^T X$, where $X$ is the input matrix. This corresponds to scaled dot-product attention with identity projections and a linearized softmax approximation ($e^x \approx 1+x$).
  • This architecture is coupled with a two-layer ReLU MLP to form the MLP-Attn model.
  • Crucially, the attention transformation is transductive (computed over the entire training set) and parameter-free, allowing for exact kernel characterization.

• Theoretical Framework:

  • NTK Analysis: The authors derive the exact NTK for this sequential architecture and analyze its convergence properties as width $m \to \infty$.
  • Spectral Analysis: They investigate how the attention transformation affects the condition number of the Gram matrix ($G = XX^T$).
  • Influence Malleability: A new metric is introduced to quantify the sensitivity of a model's reliance on training data. It measures the "flip rate" of influence scores when training examples are perturbed (e.g., via adversarial attacks).

• Experimental Setup:

  • Datasets: MNIST, CIFAR-10, and Fashion-MNIST.
  • Baselines: Compared against standard 2-layer ReLU networks (2L-ReLU).
  • Metrics: NTK distance ($\|f_m - f_{NTK}\|$), Influence Flip Rate (under FGSM, PGD, and MIM attacks), and Spearman rank correlation of influence scores.

3. Key Contributions

A. Non-Convergence of Linearized Attention

The paper proves that linearized attention does not converge to its infinite-width NTK limit, even at large widths (up to $m=4096$).

• Empirical Evidence: Unlike 2L-ReLU networks, which show a monotonic decrease in NTK distance as width increases, MLP-Attn shows non-monotonic or increasing NTK distance.
• Theoretical Proof (Spectral Amplification): The authors prove that the attention transformation cubes the Gram matrix ($\tilde{G} = G^3$).
  • This amplifies the condition number: $\kappa(\tilde{G}) = \kappa(G)^3$.
  • To achieve convergence to the NTK limit, the network width $m$ must satisfy $m = \Omega(\kappa(G)^6)$.
  • For natural image datasets where $\kappa(G) \approx 10^3$, the required width is $m \approx 10^{18}$, which is practically impossible. Thus, attention operates firmly in the feature learning regime (where representations evolve) rather than the lazy training regime.

B. Introduction of Influence Malleability

The authors define Influence Malleability as the capacity of a model to dynamically alter its reliance on specific training examples when those examples are perturbed.

• Metric: Measured via the Flip Rate—the percentage of top-influential training examples that switch from "helpful" to "harmful" (or vice versa) after an adversarial perturbation.
• Finding: Linearized attention exhibits 6–9× higher influence malleability compared to standard ReLU networks.

C. Dual Implications of Malleability

The paper establishes a fundamental trade-off driven by the non-convergent, data-dependent nature of attention:

1. Benefit (Bias Reduction): The data-dependent kernel aligns with the task structure (spectral alignment), potentially reducing approximation error (bias) when the target function matches the data distribution's dominant modes.
2. Cost (Vulnerability): The same sensitivity makes the model highly susceptible to adversarial manipulation of training data. Small perturbations to influential examples can drastically alter the model's learning trajectory and predictions.

4. Key Results

• NTK Distance Trajectory:
  • 2L-ReLU: Distance decreases monotonically with width (e.g., MNIST: 45.1 $\to$ 39.2), confirming convergence to the kernel regime.
  • MLP-Attn: Distance increases or fluctuates (e.g., MNIST: 10.3 $\to$ 43.4; CIFAR-10: 3.7 $\to$ 12.6), confirming non-convergence.
• Influence Flip Rates:
  • Under PGD attacks ($\epsilon=0.3$), MLP-Attn on MNIST showed a 28.9% flip rate, whereas 2L-ReLU showed only 3.3%.
  • This gap persists across datasets and attack types, validating the theoretical prediction of heightened sensitivity.
• Adversarial Training:
  • Adversarial training forces standard ReLU networks to become more malleable (flip rate jumps from 3.3% to 43.4%), nearly matching attention's innate malleability. This suggests that attention's sensitivity is architecturally intrinsic, not just a result of optimization dynamics.
• Loss Landscape:
  • MLP-Attn finds broader, flatter minima in the loss landscape compared to the sharp minima of ReLU networks, consistent with feature learning dynamics.

5. Significance and Implications

• Theoretical Shift: The paper challenges the applicability of infinite-width NTK theory to attention mechanisms. It demonstrates that attention is fundamentally a feature learning mechanism, not a kernel method, due to the spectral amplification of the Gram matrix.
• Robustness vs. Expressivity: It provides a theoretical explanation for the "double-edged sword" of attention: its flexibility allows it to adapt to complex data structures (low bias) but makes it inherently fragile to data poisoning and adversarial examples (high variance/sensitivity).
• Practical Guidance:
  • Practitioners should be cautious when applying kernel-based theoretical tools to attention models.
  • The high malleability suggests that attention-based systems require rigorous data cleaning and robust training strategies, as their performance is tightly coupled to the quality of the training distribution.
  • The findings suggest that the "power" of attention (adaptability) and its "vulnerability" (instability) share a common origin: the departure from the kernel regime.

In summary, this work rigorously characterizes the learning dynamics of linearized attention, proving its non-convergence to the NTK limit and identifying influence malleability as the key metric explaining both its superior adaptability and its vulnerability to adversarial data manipulation.