SympFormer: Accelerated attention blocks via Inertial Dynamics on Density Manifolds

TL;DR

This paper introduces accelerated attention blocks inspired by inertial dynamics on density manifolds, leading to faster convergence in transformer models while maintaining computational efficiency.

Contribution

It extends the particle system interpretation of attention by incorporating inertial Nesterov-type dynamics, resulting in Hamiltonian momentum attention blocks with improved convergence.

Findings

Accelerated attention blocks converge faster than classical ones.

The proposed method preserves elliptically contoured distributions.

Particle-based algorithms demonstrate improved efficiency.

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.