SympFormer: Accelerated attention blocks via Inertial Dynamics on Density Manifolds
This work addresses the computational efficiency problem in transformer architectures for natural language processing, representing an incremental improvement with a novel method for a known bottleneck.
The paper tackled the slow convergence of classical attention blocks in transformers by introducing accelerated attention blocks based on inertial dynamics on density manifolds, demonstrating that the proposed blocks converge faster while preserving the number of oracle calls.
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.