Self-Attention for Online Decision-Making and Control
Gautam Goel, University of California, Berkeley (UC Berkeley
|
Self-attention is the key algorithmic module which powers the Transformer neural architecture. However, the softmax nonlinearity appearing in self-attention makes theoretical analysis quite challenging. We study the training dynamics of gradient descent in a softmax self-attention layer trained to perform linear regression and propose a simple first-order optimization algorithm which converges to the globally optimal self-attention parameters at a geometric rate. Our analysis proceeds in two steps. First, we show that in the infinite-data limit the regression problem solved by the self-attention layer is equivalent to a nonconvex matrix factorization problem. Second, we exploit this connection to design a novel ``structure-aware" variant of gradient descent which efficiently optimizes the original finite-data regression objective. Our optimization algorithm features several innovations over standard gradient descent, including a preconditioner and regularizer which help avoid spurious stationary points, and a data-dependent spectral initialization of parameters which lie near the manifold of global minima with high probability. As an application of our results, we show that our algorithm can be used to obtain sublinear regret for the problem of online stochastic regression. |
 |