This is a bit vague so I'll begin by indicating the motivation. I am looking for ways to [do something interesting or useful] with the self-attention in transformer models. Ultimately the self-attention computation boils down to a bilinear form $x^{T}By$ where $x,y$ are vectors of embeddings of "tokens" and $B=W_{Q}^{T}W_{K}$ is a matrix composed of the attention matrices.

Now, I would like to formulate some sort of regularization term for the training process of a transformer that encourages $f_x(y)=x^{T}By$ to be approximately sparse for every vector $x$. In the sense that there are a few choices of $y$ that give very high values of $f_x(y)$ while all other choices of $y$ give very low values of $f_x(y)$.

And this is where I am stuck and fail to formulate my desideratum formally. I tried to think probabilistically about this and read a bit about Hanson-Wright inequalities but I think they do not capture the aspect I am looking for - namely, the simulataneous constraining of $f_x()$ to be sharply concentrated for all $x$.

Any suggestions will be appreciated. (Inclduing explanations why this is a dead-end).