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),
layers.Dropout(rate)
])
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?
