Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Wavy Transformer

Satoshi Noguchi, Yoshinobu Kawahara

TL;DR

Wavy Transformer reframes transformer attention as diffusion on a complete graph and introduces a second-order wave dynamic in the attention pathway to mitigate token over-smoothing. By coupling a velocity path with velocity-aware normalization and FFN, the model preserves state-velocity relations and supports oscillatory, energy-conserving updates. Empirical results across NLP, CV, and sparse-graph tasks show that combining diffusion and wave dynamics improves performance with minimal parameter overhead and no extra hyperparameter tuning, while reducing representation collapse deeper in networks. This work provides a physics-inspired design principle for transformers, offering a practical strategy to scale deep architectures without sacrificing expressivity.

Abstract

Transformers have achieved remarkable success across natural language processing (NLP) and computer vision (CV). However, deep transformer models often suffer from an over-smoothing issue, in which token representations converge to similar values as they pass through successive transformer blocks. In this paper, we establish an equivalence between the hidden-state dynamics induced by stacked attention layers and graph neural diffusion on a complete graph. From this perspective, over-smoothing can be interpreted as a consequence of the dissipative nature of the underlying diffusion dynamics. Motivated by this physical interpretation, we propose Wavy Transformer, which consists of a novel attention layer based on second-order wavy dynamics. We also introduce a feed-forward network and a normalization layer designed to preserve the physical state-velocity relationship under the chain rule, thereby extending the transformer architecture. We further validate our proposed techniques on various transformer models for NLP and CV tasks. The results consistently demonstrate that Wavy Transformer improves performance with minimal additional parameters and no extra hyperparameter tuning.

Wavy Transformer

TL;DR

Wavy Transformer reframes transformer attention as diffusion on a complete graph and introduces a second-order wave dynamic in the attention pathway to mitigate token over-smoothing. By coupling a velocity path with velocity-aware normalization and FFN, the model preserves state-velocity relations and supports oscillatory, energy-conserving updates. Empirical results across NLP, CV, and sparse-graph tasks show that combining diffusion and wave dynamics improves performance with minimal parameter overhead and no extra hyperparameter tuning, while reducing representation collapse deeper in networks. This work provides a physics-inspired design principle for transformers, offering a practical strategy to scale deep architectures without sacrificing expressivity.

Abstract

Transformers have achieved remarkable success across natural language processing (NLP) and computer vision (CV). However, deep transformer models often suffer from an over-smoothing issue, in which token representations converge to similar values as they pass through successive transformer blocks. In this paper, we establish an equivalence between the hidden-state dynamics induced by stacked attention layers and graph neural diffusion on a complete graph. From this perspective, over-smoothing can be interpreted as a consequence of the dissipative nature of the underlying diffusion dynamics. Motivated by this physical interpretation, we propose Wavy Transformer, which consists of a novel attention layer based on second-order wavy dynamics. We also introduce a feed-forward network and a normalization layer designed to preserve the physical state-velocity relationship under the chain rule, thereby extending the transformer architecture. We further validate our proposed techniques on various transformer models for NLP and CV tasks. The results consistently demonstrate that Wavy Transformer improves performance with minimal additional parameters and no extra hyperparameter tuning.

Paper Structure

This paper contains 60 sections, 2 theorems, 43 equations, 5 figures, 19 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Attention and Over-smoothing
  4. Fundamental Dynamics on Manifolds
  5. Kernel Perspective of Diffusion and Wave Equations
  6. Diffusion Equation
  7. Wave Equation
  8. Wavy Transformer
  9. Attention as Graph Neural Diffusion
  10. Wavy Dynamics Based Attention
  11. Physically Consistent Layer Normalization and Feed-forward Network
  12. Layer Normalization for Velocity
  13. Feed-forward Network for Velocity
  14. Related Works
  15. Mitigating Over‑Smoothing: Existing Approaches and This Study
  16. ...and 45 more sections

Key Result

Theorem 1

Consider a solution to the diffusion equation. Then, the following holds: which implies that the norm of the solution is monotonically non-increasing in time.

Figures (5)

  • Figure 1: Schematic of Wavy Transformer block, combining wavy attention layers and velocity‑ specific layer‑normalization and feed‑forward layers. Each layer is designed to preserve state-velocity relationship. Post-LN is assumed here; the Pre-LN case is discussed in Appendix \ref{['ape:LN_difference']}xiong2020layer.
  • Figure 2: Comparison of cosine similarity across layers for diffuse and wave residual connections with 1-sigma interval shading.
  • Figure 3: Cosine similarity across layers for a mix (+ Full Wave) with 1-sigma interval shading.
  • Figure 4: Comparisons of cosine similarities with 1-sigma interval shading for several DeiT-based models.
  • Figure 5: (a) Post-LN Wavy Transformer layer; (b) Pre-LN Wavy Transformer layer.

Theorems & Definitions (2)

  • Theorem 1
  • Theorem 2