Table of Contents
Fetching ...

Scaled Dot-Product Attention implements projection of inputs onto a common surface

Terence D Sanger

TL;DR

This paper addresses the mathematical understanding of scaled dot-product attention (SDPA). It shows that SDPA is equivalent to projecting each query onto a surface defined by nearby keys, with weights following a Gaussian-like distribution: $z_{ij} = (1/C_i) \exp(-|| q(i) - k(j) || / (2 \sigma^2))$, while the standard SDPA uses $y = softmax( (q k^T) / sqrt(d) ) v$. In the self-attention case, $q = k = v$, so the input is projected onto the space of nearby samples; in cross-attention, the projection is onto the context defined by the keys. Empirically, projection-SDPA matches standard SDPA in modeling power but is faster to train and may incur a small accuracy drop on translation tasks, motivating further optimization and extensions to time-varying nonlinear dependencies.

Abstract

Scaled dot-product attention (SDPA) is a fundamental component responsible for the success of large-language models and other nonlinear signal processing applications. The rationale for SDPA has been based upon "query, key, value" concepts borrowed from database theory, but these concepts are difficult to reconcile with standard methods in mathematical signal processing. We show that SDPA can be rewritten in a different but mathematically equivalent form as a projection of the input vectors onto a common surface determined by the inputs themselves. Therefore SDPA discovers nonlinear dependencies in the input that are time-dependent and context-dependent. The rewritten form of SDPA permits increased speed of both feedforward and learning algorithms, but more importantly suggests potential extensions. In the context of language, we re-interpret the role of SDPA as finding a time-dependent contextual meaning determined by the surface on which the set of input vectors lies. Input token embeddings are then modified by the local context surface. This interpretation differs substantially from the concept of "self-attention", and provides a strong justification for the use of SDPA for time-series data with time-varying local nonlinear dependencies.

Scaled Dot-Product Attention implements projection of inputs onto a common surface

TL;DR

This paper addresses the mathematical understanding of scaled dot-product attention (SDPA). It shows that SDPA is equivalent to projecting each query onto a surface defined by nearby keys, with weights following a Gaussian-like distribution: , while the standard SDPA uses . In the self-attention case, , so the input is projected onto the space of nearby samples; in cross-attention, the projection is onto the context defined by the keys. Empirically, projection-SDPA matches standard SDPA in modeling power but is faster to train and may incur a small accuracy drop on translation tasks, motivating further optimization and extensions to time-varying nonlinear dependencies.

Abstract

Scaled dot-product attention (SDPA) is a fundamental component responsible for the success of large-language models and other nonlinear signal processing applications. The rationale for SDPA has been based upon "query, key, value" concepts borrowed from database theory, but these concepts are difficult to reconcile with standard methods in mathematical signal processing. We show that SDPA can be rewritten in a different but mathematically equivalent form as a projection of the input vectors onto a common surface determined by the inputs themselves. Therefore SDPA discovers nonlinear dependencies in the input that are time-dependent and context-dependent. The rewritten form of SDPA permits increased speed of both feedforward and learning algorithms, but more importantly suggests potential extensions. In the context of language, we re-interpret the role of SDPA as finding a time-dependent contextual meaning determined by the surface on which the set of input vectors lies. Input token embeddings are then modified by the local context surface. This interpretation differs substantially from the concept of "self-attention", and provides a strong justification for the use of SDPA for time-series data with time-varying local nonlinear dependencies.
Paper Structure (5 sections, 10 equations, 3 figures)

This paper contains 5 sections, 10 equations, 3 figures.

Figures (3)

  • Figure 1: Illustration of projection SDPA. Blue squares are nearby/prior input vectors, solid red square is the current input vector. The output of the layer is shown by the open red square which is a linear combination of the nearby vectors. The combination weights (dotted lines) are proportional to the Gaussian-weighted distance (solid lines) from the current vector to the nearby/prior vectors. The overall effect of the layer is to project each input (solid red square) onto an output vector (open red square) given by a linear combination of other inputs (blue squares).
  • Figure 2: a. Accuracy (fraction of correct word matches) on test set and validation set during training for standard and projection SDPA models during 10 epochs of training. b. Cross-entropy loss on training set and validation set.
  • Figure 3: Sample input prompts (English) and network outputs (Spanish) after training for Transformer model using projection SDPA layers.