This series of figures examines properties of the model weights without considering any specific input to the model.

Value Matrix

Once the attention mechanism has selected a target token, the hidden_state from that token is multiplied by a Value Matrix and the result is added to the current token's hidden_state. There is a different value matrix for each head and it is low rank, just like the key and query matrices.

We can decompose the value matrix for each head into left and right singular vectors using SVD. We can then consider that the value matrix copies (and scales) the information in a subspace (defined by the left singular vectors) of the target token's hidden_state and copies it to the subspace defined by the right singular vectors.

Self, Ortho, and Rotation Mapping

When considering the subspaces spanned by the input and output-side singular vectors, it is interesting to decompose the matrix into the sum of three matrices:

• self-mapping: A matrix where the left and right singular vectors are the same.
• ortho-mapping: A matrix where all the left singular vectors are orthogonal to all the right singular vectors.
• rotation-mapping: The leftover part that corresponds to a rotation of the input subspace that is then copied into the output subspace.

This decomposition may not be unique. I use this code to calculate how much of the matrix consists of each of these components:

U, _, Vh = scipy.sparse.linalg.svds(ValueMatrix, 64)
U = np.fliplr(U)
Vh = np.flipud(Vh)
cross_prod = Vh @ U

self_mapping = np.abs(np.diagonal(cross_prod))
ortho_mapping = np.linalg.norm(Vh - cross_prod @ U.T, axis=1)
rot_mapping = np.linalg.norm((cross_prod * (1 - np.eye(cross_prod.shape[0]))) @ U.T, axis=1)

Figure 2 shows how the value matrices decompose into these three components. The decomposition of a random matrix similar to these value matrices is indicated in the figure. It is clear that the training drives the value matrices to contain much more self-mapping than normal.

It is expected that the matrices mostly consist of the ortho-mapping component since they have a 64 dimensional subspace in a 2048 dimensional space. Both self and ortho-mapping would require picking the same 64 dimensional basis vectors twice--for the input and output basis.

QK Matrix

The attention mechanism typically has the form $x^T Q^T K x$ with a separate $Q$ and $K$ for each head. We can combine the $Q$ and $K$ matrices into a single matrix that captures what the attention mechanism is doing. This isn't typically done because these matrices are low rank, for example they have dimensions 2048 x 64 in this model and the combined $Q^T K$ matrix would be 2048 x 2048 and much more computationally expensive to work with. We can examine the $Q^TK$ matrix to see if it has any interesting properties. From here on, I will refer to this matrix as $M$.

The following figures show the same analysis as performed for the value matrix earlier, but the meaning is different. For the value matrix, the left and right singular vectors specify the subspace being copied from and the subspace being copied to. For the $Q^T K$ matrix, they instead specify the subspaces used in the current token and the past tokens (query and key).