Self attention is not available as a Keras layer at the moment. The layers that you can find in the
tensorflow.keras docs are two:
AdditiveAttention() layers, implementing Bahdanau attention,
Attention() layers, implementing Luong attention.
For self-attention, you need to write your own custom layer.
I suggest you to take a look at this TensorFlow tutorial on how to implement Transformers from scratch. The Transformer is the model that popularized the concept of self-attention, and by studying it you can figure out a more general implementation. In particular, check the section Multi-Head Attention, where they develop a custom
MultiHeadAttention() layer. That is where all the attention-related action happens. In particular, study how the
Q tensors are used in it in order to compute the attention formula.
It won't be easy but it's certainly a super interesting exercise. If you make some cool model out of it please share the link to a GitHub repository. Good luck!
There is a trick you can use: since self-attention is of multiplicative kind, you can use an
Attention() layer and feed the same tensor twice (for Q, V, and indirectly K too).
You can't build a model in the Sequential way, you need the functional one. So you'd get something like:
attention = Attention(use_scale=True)(X, X)
where X is the tensor on which you want to get self-attention.
use_scale=True arg: that is a scaling of the self-attention tensor, analogous to the one that happens in the original Transformer paper. Its purpose is to prevent vanishing gradient (that happens in extreme regions of the softmax). The only difference is that in this level, the scaling parameter is learned instead of being a fixed scalar. It defaults to