SympFormer: Accelerated attention blocks via Inertial Dynamics on Density Manifolds

2026-03-17

AI Summary

This work addresses the slow convergence and limited computational efficiency of conventional self-attention mechanisms by introducing inertial dynamics into attention modeling for the first time. The authors formulate the self-attention block as a Hamiltonian momentum system on the Wasserstein-2 space of probability measures and derive an accelerated attention module through spatiotemporal discretization, which incorporates both position and velocity variables. The proposed method significantly improves convergence rates while maintaining the same number of oracle calls. Theoretical analysis shows that elliptical level-set distributions remain invariant under this framework and establishes a formal connection to Stein variational gradient flow. Empirical validation, combining particle-based approximation, linearized self-attention, and bilinear kernel techniques, demonstrates the effectiveness of the proposed module.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines

Abstract

Transformers owe much of their empirical success in natural language processing to the self-attention blocks. Recent perspectives interpret attention blocks as interacting particle systems, whose mean-field limits correspond to gradient flows of interaction energy functionals on probability density spaces equipped with Wasserstein-$2$-type metrics. We extend this viewpoint by introducing accelerated attention blocks derived from inertial Nesterov-type dynamics on density spaces. In our proposed architecture, tokens carry both spatial (feature) and velocity variables. The time discretization and the approximation of accelerated density dynamics yield Hamiltonian momentum attention blocks, which constitute the proposed accelerated attention architectures. In particular, for linear self-attention, we show that the attention blocks approximate a Stein variational gradient flow, using a bilinear kernel, of a potential energy. In this setting, we prove that elliptically contoured probability distributions are preserved by the accelerated attention blocks. We present implementable particle-based algorithms and demonstrate that the proposed accelerated attention blocks converge faster than the classical attention blocks while preserving the number of oracle calls.

Problem

• accelerated attention
• self-attention
• inertial dynamics
• density manifolds
• convergence speed

Innovation

• accelerated attention
• inertial dynamics
• density manifold
• Hamiltonian momentum
• Stein variational gradient flow