How to Capture Higher-order Correlations? Generalizing Matrix Softmax Attention to Kronecker Computation

Josh Alman, Zhao Song

In the classical transformer attention scheme, we are given three $n\times d$ size matrices $Q,K,V$ (the query, key, and value tokens), and the goal is to compute a new $n \time…