Answers.Body: <p>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.</p>
  <p>Starting from the word2vec model, what it does is it provides a separate embedding for every word in your vocabulary.</p>
  <p>Given an input vector <span class="math-container">$x_i$</span>, which is usually a one-hot encoding in the case of word2vec, you multiply it with the embedding matrix <span class="math-container">$V$</span> and get the word embedding vector: <span class="math-container">$v_i = x_i V$</span>.</p>
  <p>Now, if you have a sequence <span class="math-container">$x = \{x_0, x_1, ..., x_{T-1}\}$</span>, then each embedding vector is calculated without taking into consideration the rest of the sequence, i.e. without taking into account the context.</p>
  <p>What you want to do is compute an embedding <span class="math-container">$z_i$</span> for every <span class="math-container">$x_i$</span>, such that <span class="math-container">$z_i$</span> captures the context and depends on the rest of the sequence. And if we have the same <span class="math-container">$x_i$</span> in a sequence with different elements, then <span class="math-container">$z_i$</span> will be different.</p>
  <p>The self-attention layer of the transformer proposes the following solution:</p>
  <ol>
  <li><p>We will calculate the individual embedding of each of the elements of the sequence and we will call that the <em>value embeddings</em> <span class="math-container">$\{v_0, v_1, ..., v_{T-1}\}$</span>, where <span class="math-container">$v_i = x_i V$</span>.</p>
  </li>
  <li><p>The final embedding <span class="math-container">$z_i$</span> for each of the elements <span class="math-container">$x_i$</span> will be a weighted sum between all of the value embeddings, i.e.:
  <span class="math-container">$$ z_i = \sum_j \alpha_{i,j} v_j $$</span></p>
  </li>
  </ol>
  <p>The weight <span class="math-container">$\alpha_{i,j}$</span> dictates how much of the value embedding of <span class="math-container">$x_j$</span> should be present in the mixture when calculating <span class="math-container">$z_i$</span>. What we want is to have a high value of <span class="math-container">$\alpha_{i,j}$</span> if <span class="math-container">$x_i$</span> is tightly connected to <span class="math-container">$x_j$</span>, and a low value otherwise. Thus, we want to estimate something like a proximity score between <span class="math-container">$x_i$</span> and <span class="math-container">$x_j$</span>. Usually this is simply done by taking the scalar product between the two <span class="math-container">$score(x_i, x_j) = x_i x_j^T$</span>. But then <span class="math-container">$score(x_i, x_j) = score(x_j, x_i)$</span>, which is not what we want. We don't actually care how close are <span class="math-container">$x_i$</span> and <span class="math-container">$x_j$</span>, rather we want to know how relevant is <span class="math-container">$x_j$</span> to <span class="math-container">$x_i$</span> - one-way relation. Later we would like to know how relevant is <span class="math-container">$x_i$</span> to <span class="math-container">$x_j$</span> and we expect the score to be different.</p>
  <p>And so the attention score is actually computed as <span class="math-container">$score(x_i, x_j) = x_i W x_j^{T}$</span>, where <span class="math-container">$W$</span> is an additional embedding matrix, which is actually decomposed into two matrices: <span class="math-container">$W=QK^T$</span>. Now for the score we have <span class="math-container">$score(x_i, x_j) = x_i Q K^T x_j^{T}$</span>, and we can say that <span class="math-container">$q_i = x_i Q$</span> is our query embedding of <span class="math-container">$x_i$</span> and <span class="math-container">$k_j = x_j K$</span> is our key embedding of <span class="math-container">$x_j$</span>.</p>
  <p>Now, except for making this nice query-key-value analogy, there is another reason for decomposing <span class="math-container">$W$</span> into <span class="math-container">$W=QK^T$</span>, 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: <a href="https://pi-tau.github.io/posts/transformer/" rel="nofollow noreferrer">https://pi-tau.github.io/posts/transformer/</a></p>
  
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2023-07-24T12:40:37.430Z
Answers.LastActivityDate: 2023-08-07T08:49:39.777Z
Answers.LastEditDate: 2023-08-07T08:49:39.777Z
Answers.UserDisplayName: pi-tau
Answers.Body: <p>Something I found helpful was &quot;<a href="https://blog.nelhage.com/post/transformers-for-software-engineers/" rel="nofollow noreferrer">Transformers for Software Engineers</a>&quot; - unrolling the matrix multiplications into a funky functional program which maps over vectors.</p>
  <p>We can follow this approach backwards and see the K, Q, V naturally emerge.</p>
  <p>I'm going to use pseudo-TypeScript and illustrate the residual stream as a JSON (I'll use the term &quot;residual&quot; 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):</p>
  <pre><code>let residualStream: Residual[] = [
    { &quot;token&quot;: oneHot(&quot;Hello&quot;), position: oneHot(100) },
    { &quot;token&quot;: oneHot(&quot;World&quot;), position: oneHot(101) },
    { &quot;token&quot;: oneHot(&quot;.&quot;), position: oneHot(102) },
  ]
  </code></pre>
  <p>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):</p>
  <pre><code>residualStream = [
    { &quot;token&quot;: oneHot(&quot;Hello&quot;), position: oneHot(100), partOfSpeech: &quot;noun&quot;, prevToken: oneHot(&quot;.&quot;), lastMatchingPosition: oneHot(30), whatCameNextLastTime: oneHot(&quot;Everyone&quot;), ... },
    { &quot;token&quot;: oneHot(&quot;World&quot;), position: oneHot(101), partOfSpeech: &quot;noun&quot;, prevToken: oneHot(&quot;Hello&quot;), lastMatchingPosition: oneHot(50), whatCameNextLastTime: oneHot(&quot;!&quot;), ... },
    { &quot;token&quot;: oneHot(&quot;.&quot;), position: oneHot(102), partOfSpeech: &quot;noun&quot;, prevToken: oneHot(&quot;World&quot;), lastMatchingPosition: oneHot(10), whatCameNextLastTime: oneHot(&quot;Bye&quot;), ... },
  ]
  </code></pre>
  <p>Feedforward layer is like <a href="https://jax.readthedocs.io/en/latest/_autosummary/jax.vmap.html" rel="nofollow noreferrer"><code>vmap</code></a>ping a function of a single element:</p>
  <pre><code>type FeedForward = Residual =&gt; Residual
  let feedforward: FeedForward
  const applyFeedForward = residualStream =&gt; residualStream.map(feedforward)
  residualStream = applyFeedForward(residualStream)
  </code></pre>
  <p>As you know, this is sufficient to calculate things like:</p>
  <pre><code>const getLengthOfToken: FeedForward = ({ token, ...rest }) =&gt;
    ({ token, length: token.length(), ...rest })
  </code></pre>
  <p>but not something like</p>
  <pre><code>const getPreviousToken: FeedForward = ???
  </code></pre>
  <p>since the previous token is not derivable from the residual. So now we need to bust out the attention mechanism:</p>
  <pre><code>const getPreviousToken: AttentionHead
  </code></pre>
  <p>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:</p>
  <pre><code>type WV = Residual =&gt; Value
  const getTokenValue: WV = residual =&gt; residual.token
  
  type WK = Residual =&gt; Key
  const getTokenPosition: WK = residual =&gt; residual.position
  </code></pre>
  <p>And the last step is likely to be storing the result back into the residual stream:</p>
  <pre><code>type WO = (res: Residual, z: Z) =&gt; Residual
  // assuming we've figured out `z`, which'll be the next part of this answer
  const recordPreviousToken: WO = (res, z) =&gt; { ...res, prevToken: z }
  </code></pre>
  <p>Note that these parts don't need to be nearly as wide as the residual stream, since they are just operating on specific features:</p>
  <pre><code>K = [oneHot(100), oneHot(101), oneHot(102)]
  V = [oneHot(&quot;Hello&quot;), oneHot(&quot;World&quot;), oneHot(&quot;.&quot;)]
  Z = [oneHot(&quot;!&quot;), oneHot(&quot;Hello&quot;), oneHot(&quot;World&quot;)]
  </code></pre>
  <p>What's more is, we probably want to do as many of these small operations like <code>getPreviousToken</code> in parallel as possible. So this is where we get multi-head attention.</p>
  <p>Now we have found <code>W_V</code>, <code>W_K</code>, <code>W_O</code>, we only need to find <code>W_Q</code> 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.</p>
  <pre><code>type AttentionHead = (residual: Residual, otherResiduals[]) =&gt; Value
  const getPreviousTokenValue = (residual, otherResiduals[]) =&gt; (
    sum(otherResiduals.map(otherResidual =&gt;
      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!
    ))
  )
  </code></pre>
  <p>We found <code>W_Q</code>! It's:</p>
  <pre><code>type WQ = Residual =&gt; Query
  const getPreviousTokenPosition = residual =&gt; shiftLeft(residual.position)
  </code></pre>
  <p>To recap:</p>
  <ul>
  <li><code>W_V</code> gets a piece of info (&quot;value&quot;) from each residual</li>
  <li><code>W_K</code> positions (&quot;keys&quot;) the residual in N-dimensional space in a way appropriate for the attention head's task</li>
  <li><code>W_Q</code> looks (&quot;queries&quot;) in N-dimensional space for other residuals relevant to the attention head's task</li>
  <li><code>W_O</code> sums up the info/results (&quot;output&quot;) and stores it back into the residual stream</li>
  </ul>
  <p>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 &quot;reads&quot;? (3 is counting K, Q, V)</p>
  <p>I tried this and I concluded that the general way to express an arbitrary comparison of two vectors (in our case, residuals), is a <a href="https://en.wikipedia.org/wiki/Bilinear_form" rel="nofollow noreferrer">bilinear form</a>. And the way to speed that up so that it's not <code>N_INPUT^2 * D_RESIDUAL</code> is to perform a <a href="https://en.wikipedia.org/wiki/Low-rank_approximation" rel="nofollow noreferrer">low-rank approximation</a>, which is exactly what self-attention does by way of the dimension-reducing maps <span class="math-container">$W_K$</span> and <span class="math-container">$W_Q$</span> approximating a <a href="https://en.wikipedia.org/wiki/Rank_factorization" rel="nofollow noreferrer">rank factorization</a> <span class="math-container">$A = W_Q W_K^T$</span>. But I believe the various &quot;fast attention&quot; algorithms each have a different take on this.</p>
  <p>BTW, the example I used is from <a href="https://transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html" rel="nofollow noreferrer">&quot;In-Context Learning and Induction Heads&quot;</a>. 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.</p>
  
Answers.ContentLicense: CC BY-SA 4.0
Answers.CreationDate: 2023-07-24T03:22:24.660Z
Answers.LastActivityDate: 2023-07-25T14:13:44.137Z
Answers.LastEditDate: 2023-07-25T14:13:44.137Z
Answers.UserDisplayName: llllvvuu
Body: <p>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.</p>
  
Tags: <neural-networks><transformer><attention>
Title: Can someone help me understand the intuition behind the query, key and value matrices in the transformer architecture?
UserDisplayName: Chinmay
ThreadText: # Can someone help me understand the intuition behind the query, key and value matrices in the transformer architecture?
  
  (Asked by: username_0 on 2023-07-24)
  
  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.
  
  **Tags:** neural-networks, transformer, attention
  
  ---
  
  ## 2 answers
  
  ### Answer by username_2 on 2023-07-24. Score: 3 | Accepted answer
  
  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.
  
  
  ---
  
  ### Answer by username_1 on 2023-07-24. Score: 5
  
  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/>
  
  
  ---
  
  