Changing MLP layers in Transformers to Probabilistic encoders

idea

In Transformer Feed-Forward Layers Are Key-Value Memories, the interaction between the input $x$ and each column $k_i$ of the first parameter matrix is viewed as a conditional distribution of $k_i$ given $x$:

Link to original

What if we employ the idea of VAE and make $x$ parameterizing $p(k_i|x)$ as Gaussian, something like:

$$\begin{array}{rl}{\mu }_x& =M(x)\\ & \\ p(k_i|x)& \propto e^{\frac{-1}{2}(k_i-{\mu }_x{)}^2}\end{array}$$

which is a Gaussian with mean give by $x$ and a fixed std.

The intuition is, given an input $x$ that produces ${\mu }_x$, if $k_i$ is distributed near ${\mu }_x$ then it is more likely to be related to $x$. This will force the model to learn $k_i$ such that it’s distributed near ${\mu }_x$.

But perhaps the MLP already has this property ? As $x\cdot k_i$ is bigger when $x$ and $k_i$ are more co-linear, the “almost co-linear” region surrounding $x$ can be seen as a distribution parameterized (partly or entirely) by $x$.