8
$\begingroup$

I have been working mechanically with transformers, hoping that with time clarity about what the query, key, and value matrices represent will develop; but I am still lost. Would greatly benefit from a simplified explanation.

$\endgroup$
1
  • $\begingroup$ They are matrices because a lot of queries, keys and values are processed in parallel, but I think it is easier to understand what each single entry in the matrix represents. $\endgroup$
    – user253751
    Jul 24, 2023 at 10:13

2 Answers 2

3
$\begingroup$

Something I found helpful was "Transformers for Software Engineers" - unrolling the matrix multiplications into a funky functional program which maps over vectors.

We can follow this approach backwards and see the K, Q, V naturally emerge.

I'm going to use pseudo-TypeScript and illustrate the residual stream as a JSON (I'll use the term "residual" to describe each individual element of the array, i.e. the original token plus any information we accumulate about that token in the context of the input):

let residualStream: Residual[] = [
  { "token": oneHot("Hello"), position: oneHot(100) },
  { "token": oneHot("World"), position: oneHot(101) },
  { "token": oneHot("."), position: oneHot(102) },
]

Our goal is to add useful features to the residuals until we have so many that predicting the next word becomes trivial. We might end up with something like this by the end (requires the residual stream to be wide enough):

residualStream = [
  { "token": oneHot("Hello"), position: oneHot(100), partOfSpeech: "noun", prevToken: oneHot("."), lastMatchingPosition: oneHot(30), whatCameNextLastTime: oneHot("Everyone"), ... },
  { "token": oneHot("World"), position: oneHot(101), partOfSpeech: "noun", prevToken: oneHot("Hello"), lastMatchingPosition: oneHot(50), whatCameNextLastTime: oneHot("!"), ... },
  { "token": oneHot("."), position: oneHot(102), partOfSpeech: "noun", prevToken: oneHot("World"), lastMatchingPosition: oneHot(10), whatCameNextLastTime: oneHot("Bye"), ... },
]

Feedforward layer is like vmapping a function of a single element:

type FeedForward = Residual => Residual
let feedforward: FeedForward
const applyFeedForward = residualStream => residualStream.map(feedforward)
residualStream = applyFeedForward(residualStream)

As you know, this is sufficient to calculate things like:

const getLengthOfToken: FeedForward = ({ token, ...rest }) =>
  ({ token, length: token.length(), ...rest })

but not something like

const getPreviousToken: FeedForward = ???

since the previous token is not derivable from the residual. So now we need to bust out the attention mechanism:

const getPreviousToken: AttentionHead

and we'll see that it'll break down into K, Q, V like you asked. The first step is likely to be picking out the relevant info:

type WV = Residual => Value
const getTokenValue: WV = residual => residual.token

type WK = Residual => Key
const getTokenPosition: WK = residual => residual.position

And the last step is likely to be storing the result back into the residual stream:

type WO = (res: Residual, z: Z) => Residual
// assuming we've figured out `z`, which'll be the next part of this answer
const recordPreviousToken: WO = (res, z) => { ...res, prevToken: z }

Note that these parts don't need to be nearly as wide as the residual stream, since they are just operating on specific features:

K = [oneHot(100), oneHot(101), oneHot(102)]
V = [oneHot("Hello"), oneHot("World"), oneHot(".")]
Z = [oneHot("!"), oneHot("Hello"), oneHot("World")]

What's more is, we probably want to do as many of these small operations like getPreviousToken in parallel as possible. So this is where we get multi-head attention.

Now we have found W_V, W_K, W_O, we only need to find W_Q now. Let's try to finish our function. We'll have to draw upon the defining feature of attention, which is: taking the residuals of the other tokens into account.

type AttentionHead = (residual: Residual, otherResiduals[]) => Value
const getPreviousTokenValue = (residual, otherResiduals[]) => (
  sum(otherResiduals.map(otherResidual =>
    getTokenValue(otherResidual) // W_V
    * scaledSoftmax( // handwaving this part; in reality this would apply over Q * K^T as a whole
      getTokenPosition(otherResidual) // W_K
      .dot(getPreviousTokenPosition(residual)) // we found W_Q!
  ))
)

We found W_Q! It's:

type WQ = Residual => Query
const getPreviousTokenPosition = residual => shiftLeft(residual.position)

To recap:

  • W_V gets a piece of info ("value") from each residual
  • W_K positions ("keys") the residual in N-dimensional space in a way appropriate for the attention head's task
  • W_Q looks ("queries") in N-dimensional space for other residuals relevant to the attention head's task
  • W_O sums up the info/results ("output") and stores it back into the residual stream

As a final exercise, I challenge you to accomplish to challenge the attention head design; that is, can you accomplish the same thing with less than 3 "reads"? (3 is counting K, Q, V)

I tried this and I concluded that the general way to express an arbitrary comparison of two vectors (in our case, residuals), is a bilinear form. And the way to speed that up so that it's not N_INPUT^2 * D_RESIDUAL is to perform a low-rank approximation, which is exactly what self-attention does by way of the dimension-reducing maps $W_K$ and $W_Q$ approximating a rank factorization $A = W_Q W_K^T$. But I believe the various "fast attention" algorithms each have a different take on this.

BTW, the example I used is from "In-Context Learning and Induction Heads". Highly recommended reading for understanding how attention heads work. See if you can come up with similar pseudo-code as I did, but for the rest of the Induction Head circuit.

$\endgroup$
1
  • $\begingroup$ Thank you! This is awesome! $\endgroup$
    – Chinmay
    Jul 24, 2023 at 19:15
5
$\begingroup$

A very very distant connection can be seen between the self-attention layer and the word2vec model. I think that this might be helpful to in order to gain more intution.

Starting from the word2vec model, what it does is it provides a separate embedding for every word in your vocabulary.

Given an input vector $x_i$, which is usually a one-hot encoding in the case of word2vec, you multiply it with the embedding matrix $V$ and get the word embedding vector: $v_i = x_i V$.

Now, if you have a sequence $x = \{x_0, x_1, ..., x_{T-1}\}$, then each embedding vector is calculated without taking into consideration the rest of the sequence, i.e. without taking into account the context.

What you want to do is compute an embedding $z_i$ for every $x_i$, such that $z_i$ captures the context and depends on the rest of the sequence. And if we have the same $x_i$ in a sequence with different elements, then $z_i$ will be different.

The self-attention layer of the transformer proposes the following solution:

  1. We will calculate the individual embedding of each of the elements of the sequence and we will call that the value embeddings $\{v_0, v_1, ..., v_{T-1}\}$, where $v_i = x_i V$.

  2. The final embedding $z_i$ for each of the elements $x_i$ will be a weighted sum between all of the value embeddings, i.e.: $$ z_i = \sum_j \alpha_{i,j} v_j $$

The weight $\alpha_{i,j}$ dictates how much of the value embedding of $x_j$ should be present in the mixture when calculating $z_i$. What we want is to have a high value of $\alpha_{i,j}$ if $x_i$ is tightly connected to $x_j$, and a low value otherwise. Thus, we want to estimate something like a proximity score between $x_i$ and $x_j$. Usually this is simply done by taking the scalar product between the two $score(x_i, x_j) = x_i x_j^T$. But then $score(x_i, x_j) = score(x_j, x_i)$, which is not what we want. We don't actually care how close are $x_i$ and $x_j$, rather we want to know how relevant is $x_j$ to $x_i$ - one-way relation. Later we would like to know how relevant is $x_i$ to $x_j$ and we expect the score to be different.

And so the attention score is actually computed as $score(x_i, x_j) = x_i W x_j^{T}$, where $W$ is an additional embedding matrix, which is actually decomposed into two matrices: $W=QK^T$. Now for the score we have $score(x_i, x_j) = x_i Q K^T x_j^{T}$, and we can say that $q_i = x_i Q$ is our query embedding of $x_i$ and $k_j = x_j K$ is our key embedding of $x_j$.

Now, except for making this nice query-key-value analogy, there is another reason for decomposing $W$ into $W=QK^T$, but I think that this answer is already long enough. Feel free to checkout a blog post that I wrote about the transformer model if you want to read more: https://pi-tau.github.io/posts/transformer/

$\endgroup$

You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .