Title: Wavy Transformer
Venue: NeurIPS 2025
Accepted: yes
Keywords: Transformers

Over-smoothing

Wave dynamics on graphs

Partial differential equations
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a modification to the transformer architecture based on the wave equation. The authors interpret the existing transformer layers as a discretization of the diffusive heat equation. Under this interpretation, they argue that the over-smoothing issue of the transformer -- where representations of different tokens become too similar at deep layers -- is because the tokens follow the dynamics of heat diffusion on a complete graph. They then modify the transformer layers to mimic the discretization of the wave equation, which does not follow diffusive dynamics. They conduct experiments to show that this modification (sometimes in combination with the original transformer layers) gives better performance on NLP and CV tasks.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. The paper is well-written with a clear logical flow that is easy to follow. The relevant background material is clearly presented. \n2. The motivation for the proposed modification is clearly explained. \n3. The proposed modification seeks to address a relevant issue observed by many other works.\n\nWeaknesses: \n1. The experimental results supporting the proposed modification are not particularly strong. The range of tasks selected is small, and where there are improvements, they seem marginal. Since the modification is quite fundamental, I expect to see a wider range of tasks and larger performance improvements.\n\n2. While the cosine similarity for the wavy transformer exhibits oscillations across layers (which is different from the monotone increase of the original transformer), the representations of the wavy transformer almost always have higher cosine similarity. This suggests that the proposed modification does not reduce over-smoothing. The authors state that the wavy transformer “extracts fundamentally different features”, but this is not a sufficient explanation of the effects of the proposed modification and why we should expect it to be better. This also seems to contradict the original motivation to alleviate over-smoothing.\n\nQUESTIONS:\n1. As mentioned above, given that the wavy transformer results in higher cosine similarity between representations, how do we explain the effects of the modification? If it does not reduce over-smoothing, then what problem is it solving?\n2. Is over-smoothing a phenomenon observed in models that use causal attention as well? If so, how do we understand it from the diffusive dynamics perspective?\n3. In essence, the paper posits that representations evolve layer-wise through an oscillatory process. Why would that be the case?\n4. You’re proposing the wave equation as an example for a non-diffusive process. Why the wave equation? There are other PDEs that would also not diffuse.\n5. From the last few paragraphs in Subsection 3.2, it’s not crystal clear what is your proposed modification. It seems like there are several options considered.\n6. It would be useful to include an ablation study showing the impact of Equation 15.\n7. “the wavy residual connection alone just achieves nearly the same-level performance as the diffusive residual connection” — Do the authors have a hypothesis for why this occurs?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 5, "CLARITY": 4, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-04"}

comments: SUMMARY:
The paper discusses the equivalence between hidden state dynamics and graph neural diffusion on a complete graph, and interprets over-smoothing as a consequence of the dissipative nature of diffusion dynamics.

STRENGTHS AND WEAKNESSES:
The research topic and its application to text is interesting.
However, the biggest problem is that the paper presents methods and theoretical results/interpretations that have been extensively studied in Graph Neural Networks research (and does not cite them). In particular:
   - The idea that over-smoothing is a consequence of the dissipative nature of diffusion dynamics has been previously established and discussed   [1,3]
   - The introduction of wavy dynamics to mitigate over-smoothing is not new and has been implemented in GNNs  [2, 3]. Wave equation based transformer blocks have also been studied in other contexts [4]


Moreover, the paper extensively discusses GNN and interprets attention layers as "graph neural diffusion on a complete graph". However, the experimental validation does not present any direct experimental results on traditional graph datasets.

Overall, the paper presents an interesting idea, which however needs to be better positioned and contrasted with existing works. This would also imply adding baselines in the experimental sections. 

[1] Gravina et al., 2025. On Oversquashing in Graph Neural Networks Through the Lens of Dynamical Systems

[2] Rusch, et al. 2024 Graph-Coupled Oscillator Network

[3] Eliasof et al., 2021. PDE-GCN: Novel Architectures for Graph Neural Networks Motivated by Partial Differential Equations

[4] Deng et al., 2025. Denoising Hamiltonian Network for Physical Reasoning

QUESTIONS:
See weaknesses.

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 2
IMPACT: 2
ORIGINALITY: 1
DATE: 2025-07-02

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper establishes a link between the dynamics of stacked attention layers and diffusion processes on a complete graph, revealing over-smoothing as a consequence of dissipative diffusion dynamics. Motivated by this insight, the authors propose the Wavy Transformer, which introduces a novel attention mechanism based on second-order wave-like dynamics. Unlike diffusion, wave dynamics preserve energy and exhibit oscillatory behavior, helping to potentially mitigate over-smoothing. The authors find that Wavy Transformer achieves consistent performance improvements on NLP and CV tasks, without significant increases in parameters or hyperparameter tuning.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n- Proposes a novel approach that modifies transformer dynamics using physical principles, offering a principled approach aimed at reducing over-smoothing.\n- Provides clear explanations and construction of the proposed architecture components.\n\nWeaknesses:\n\n- A potential drawback is the added computational overhead introduced by the method. In the case of vision tasks, the authors report that the inference throughput drops by about 50% when using wave dynamics, which could limit practical usage of the method.\n- See the Questions section for more discussion\n\nQUESTIONS:\n- Since the computational cost seems to be mainly discussed for inference of vision models, how does the inference computational performance of Wavy Transformer compare for NLP models?\n- How much computational overhead does the method introduce during training? Additionally, how does incorporating wave dynamics affect training convergence? Are there notable differences in convergence behavior compared to standard transformers?\n- Part of the motivation for this method was to mitigate over-smoothing in transformers. However, based on Figures 2 and 5, wave dynamics don’t appear to reduce over-smoothing (at least based on the token-wise cosine similarity metric), although some differences emerge in its layer-wise behavior. Could the authors clarify this, and if this is indeed the case, could they elaborate on other explanations or benefits of Wavy Transformer beyond over-smoothing?\n- In Section 5.1.4, the authors note oscillatory behavior in cosine similarity for wavy residual connections, but this behavior seems less clear in the DeiT experiments (Figures 5 and 6). Are there reasons for this, and any insights into which behavior might be expected or beneficial for the model? Also, does such behavior persist or vary across prompts and different datasets?\n- Have the authors considered applying Wavy Transformer to decoder-only architectures for language tasks as well?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-03"}

comments: SUMMARY:
This paper addresses the over-smoothing problem in deep Transformer models by reinterpreting attention layers as a form of graph diffusion on a complete graph. Based on this view, the authors propose the Wavy Transformer, which replaces the diffusive dynamics with wave-based (second-order) dynamics that preserve energy and reduce smoothing. The architecture also introduces corresponding modifications to feed-forward and normalization layers to maintain consistency with the proposed dynamics. Experiments on both NLP and vision tasks demonstrate consistent improvements with minimal parameter overhead.

STRENGTHS AND WEAKNESSES:
## Strengths
1. Offers a solid theoretical interpretation of attention as graph diffusion, grounding the over-smoothing issue in physical dynamics.
2. Introduces a novel wave-based update rule that mitigates over-smoothing by preserving energy and promoting oscillatory behavior.
3. Demonstrates consistent gains across NLP and CV tasks.

## Weaknesses
1. Evaluation is limited to relatively small models (BERT hidden size 256, DeiT-Ti), with no experiments on large-scale settings.
2. Computational cost (runtime, memory) is not discussed, even though parameter count remains unchanged, the added velocity tracking may incur overhead.

QUESTIONS:
1. How does the method scale to larger models?
2. Could you comment on the runtime/memory cost introduced by the velocity tracking, compared to a standard Transformer block?

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 4
IMPACT: 3
ORIGINALITY: 4
DATE: 2025-07-03