Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

SympFormer: Accelerated attention blocks via Inertial Dynamics on Density Manifolds

Viktor Stein, Wuchen Li, Gabriele Steidl

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.

SympFormer: Accelerated attention blocks via Inertial Dynamics on Density Manifolds

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--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.
Paper Structure (23 sections, 6 theorems, 81 equations, 2 figures, 4 tables, 1 algorithm)

This paper contains 23 sections, 6 theorems, 81 equations, 2 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Nesterov's Acceleration Method
  3. Time Discretizations of Linearly Damped Hamiltonian Systems
  4. Explicit Euler method
  5. Conformally symplectic Euler method
  6. Exponential Euler method
  7. Adams-Bashforth (AB-2) method
  8. Transformers
  9. Accelerated Transformers with Linear Attention
  10. Accelerated Flow
  11. Preservation of Elliptically Contoured Distributions
  12. Accelerated Linear Attention Dynamics
  13. Accelerated Transformers with Softmax attention
  14. Accelerated Flow
  15. Accelerated Softmax Attention Dynamics
  16. ...and 8 more sections

Key Result

Proposition 4.1

The accelerated linear self-attention transformer flow in sys_1 satisfies Here, $(\alpha_t)_{t > 0}$ are non-negative damping parameters.

Figures (2)

  • Figure 1: The accelerated attention block followed by an accelerated MLP block as described in \ref{['algo:SympFormer']}.
  • Figure 2: Validation loss (circles) and training loss (lines) on the tinystories data set after 10000 optimization steps, illustrating the results from \ref{['tab:tinystories_softmax_1']}.

Theorems & Definitions (10)

  • Proposition 4.1: Accelerated linear attention PDE
  • Example 4.2: Elliptically contoured distributions
  • Proposition 4.3: Preservation of elliptically contoured distributions
  • Corollary 4.4: Preservation of centered elliptically contoured distributions
  • Proposition 4.5
  • Proposition 5.1: Accelerated softmax attention PDE
  • Proposition 5.2
  • Remark 5.3
  • Example 5.4
  • Example 5.5