I am wondering what are the trainable weights inside an attention-powered transformer. I figure the feed-forward layer contain trainable weights and the token embeddings, but what other parts contain trainable weights? i.e. when you store a trained model, what parts are you actually storing
I also wonder what weights inside the attention-mechanism must be temporary stored, or being held in memory, during a training iteration until the backpropagated gradient descent comes back?
Computation of one a single head of the scaled dot-product attention for values $V$, keys $K$ and queries $Q$ does not need any parameters. The output of the attention is simply (Eq. 1 of the Attention is all you need paper).
$$ \mathrm{Attention}(Q, K, V) = \mathrm{softmax}\left(\frac{QK}{\sqrt{d}} \right) V $$
However, in the multi-head setup, the keys, values, and queries are a result of a head-specific projection (unnumbered equation on the top of page 5)
$$ \mathrm{MultiHead}(Q, K, V) = \mathrm{Concat}(\mathrm{head}_1,\ldots,\mathrm{head}_h)W^O $$
where $$ \mathrm{head}_i = \mathrm{Attention}(QW_i^Q, KW_i^K, VW_i^V). $$
Here, all the projection matrices $W$ are trainable parameters. The Transformer Base architecture has dimensions 512, 8 heads and each head has a dimension 64. This means that the matrices $W_i^Q$, $W_i^K$, and $W_i^V$ have 512×64 parameters and there are eight of them (in practice, they are grouped into large matrices and the output of the large matrix multiplication is then reshaped accordingly), this is 262k parameters. Matrix $W^O$ projects a concatenation of 8 64-dimensional heads into 512 model, dimensions, i.e., further 262k parameters. In some implementations, there are also biases to the projections, but this is negligible in terms of parameter count.
The Transformer Base architecture thus has over 0.5M parameters in each self-attention sublayer. Compared to the feed-forward layers (2×512×2048≈2M), it is approx 4 times less.
During the forward and backward pass, there must a node for each of the computation (each dot product, each softmax) and all the intermediate results needs to be stored in order compute the gradient.
