About attention: the Query, Key and Value vectors (before the linear transformations) are just the entire sequence, that is being inputted, or just each token? Chat-GPT nor Youtube didn't give me a clear answer. But, I thought. If we feed in each sequence straight into the Attention mechanism, then the linear layers, which are supposed to transform these inputs, won't be able to accept that, because sequences can be different length, and linear layers' input shape is fixed. And if we process each token independently, we have to store the other tokens somewhere, and then have loops to iterate through each token etc.

So, I decided to find some code, where people create that Multi-Head attention, and here's what I found:

import tensorflow as tf
from tensorflow.keras import layers

class MultiHeadAttention(layers.Layer):
    def __init__(self, model_dim, n_heads, rate=0.1, initializer='glorot_uniform'):
        super(MultiHeadAttention, self).__init__()
        self.n_heads = n_heads
        self.model_dim = model_dim

        assert model_dim % self.n_heads == 0

        self.head_dim = model_dim // self.n_heads

        self.wq = layers.Dense(model_dim, kernel_initializer=initializer)
        self.wk = layers.Dense(model_dim, kernel_initializer=initializer)
        self.wv = layers.Dense(model_dim, kernel_initializer=initializer)
        self.dropout1 = layers.Dropout(rate)
        self.dropout2 = layers.Dropout(rate)
        self.wo = layers.Dense(model_dim, kernel_initializer=initializer)

    def split_heads(self, x, batch_size):
        x = tf.reshape(x, (batch_size, -1, self.n_heads, self.head_dim))
        return tf.transpose(x, perm=[0, 2, 1, 3])

    def call(self, q, k, v, mask=None):
        batch_size = tf.shape(q)[0]

        q = self.wq(q)  
        k = self.wk(k)  
        v = self.wv(v)  

        q = self.split_heads(q, batch_size) 
        k = self.split_heads(k, batch_size)  
        v = self.split_heads(v, batch_size) 

        dh = tf.cast(self.head_dim, tf.float32)
        qk = tf.matmul(q, k, transpose_b=True)
        scaled_qk =  qk / tf.math.sqrt(dh)
        if mask is not None:
            scaled_qk += (mask * -1e9) 

        attn = self.dropout1(tf.nn.softmax(scaled_qk, axis=-1))
        attn = tf.matmul(attn, v) 

        attn = tf.transpose(attn, perm=[0, 2, 1, 3]) 
        original_size_attention = tf.reshape(attn, (batch_size, -1, self.model_dim)) 

        output = self.dropout2(self.wo(original_size_attention))
        return output
class TransformerBlock(layers.Layer):
    def __init__(self, emb_dim, n_heads, mlp_dim, 
                 rate=0.1, initializer='glorot_uniform', eps=1e-6, activation='gelu'):
        super(TransformerBlock, self).__init__()
        self.attn = MultiHeadAttention(emb_dim, n_heads, initializer=initializer)
        self.mlp = tf.keras.Sequential([
            layers.Dense(mlp_dim, activation=activation, kernel_initializer=initializer), 
            layers.Dense(emb_dim, kernel_initializer=initializer),
        self.ln1 = layers.LayerNormalization(epsilon=eps)
        self.ln2 = layers.LayerNormalization(epsilon=eps)

    def call(self, inputs, mask=None):
        x = self.ln1(inputs)
        x = inputs + self.attn(x, x, x, mask) 
        x = x + self.mlp(self.ln2(x))
        return x

So, here I see that they basically feed in the entire sequence into the Multi-Head attention:

x = self.ln1(inputs) x = inputs + self.attn(x, x, x, mask)

But again, sequences can be different length, but linear layers accept inputs of fixed length. Is it me not understanding something? I also read about Padding and Masking, that is used during training. Is it also used during inference?


3 Answers 3



If the sequence length of $Q, K, V$ is $L$, the embedding size is $E$, and the number of heads is $H$, then weight matrices are of the order $E \times (E // H)$ to transform $E$-sized embeddings to $(E // H)$ sized embeddings. We will have $H$ such embeddings; we get the same $E$-sized vectors upon concatenation. The attention weights will be of size $L \times L$ to enforce every token is attending to every other token in self-attention mode.


During the training, we would process the entire sequence at once and train the Transformer with the teacher-forcing algorithm usually. If the input sequences vary in length, we would truncate or pad them to a same length (of 512 tokens for instance).

During the inference, the generation of the output sequence is done via ancestral sampling. We start with a special token (SOS - Start-of-sequence) and sample the next token, and then feed this new token back in again and continue.


Yes, the query, key and value are the entire sequence. If the sequence is smaller than the maximum sequence length (which is the size of the linear layers), then you use padding tokens to complete the sequence and masking in the attention. If it's longer you cannot process the entire sequence with a standard self attention layer.


You must log in to answer this question.

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