def calculate_dot_product_similarities(
query: Tensor,
key: Tensor,
) -> Tensor:
"""
Calculate similarity scores between queries and keys using dot product.
Args:
query: embedding vector of query of shape (B, h, T, d_k)
key: embedding vector of key of shape (B, h, T, d_k)
Returns: Similarities (closeness) between q and k of shape (B, h, T, T) where
last (T, T) represents relations between all query elements in T sequence
against all key elements in T sequence. If T is people in an organization,
(T,T) represents all (cartesian product) social connections among them.
The relation considers d_k number of features.
"""
# --------------------------------------------------------------------------------
# Relationship between k and q as the first MatMul using dot product similarity:
# (B, h, T, d_k) @ (B, hH, d_k, T) ---> (B, h, T, T)
# --------------------------------------------------------------------------------
similarities = query @ key.transpose(-2, -1) # dot product
return similarities # shape:(B, h, T, T)
Code to demonstrate
def calculate_dot_product_similarities(
query: Tensor,
key: Tensor,
) -> Tensor:
"""
Calculate similarity scores between queries and keys using dot product.
Args:
query: embedding vector of query of shape (B, h, T, d_k)
key: embedding vector of key of shape (B, h, T, d_k)
Returns: Similarities (closeness) between q and k of shape (B, h, T, T) where
last (T, T) represents relations between all query elements in T sequence
against all key elements in T sequence. If T is people in an organization,
(T,T) represents all (cartesian product) social connections among them.
The relation considers d_k number of features.
"""
# --------------------------------------------------------------------------------
# Relationship between k and q as the first MatMul using dot product similarity:
# (B, h, T, d_k) @ (B, hH, d_k, T) ---> (B, h, T, T)
# --------------------------------------------------------------------------------
similarities = query @ key.transpose(-2, -1) # dot product
return similarities # shape:(B, h, T, T)
def scale(
similarities: Tensor,
d_k: int
) -> Tensor:
"""
Standardize the variance of the dot product similarities using the standard deviation
of the dot product of the normal distributions std=sqrt(d_k) so that the variance will
be 1.0 approx.
Citation:
> While for small values of dk the two mechanisms perform similarly, additive attention
> outperforms dot product attention without scaling for larger values of dk [3].
> We suspect that for large values of d_k, the dot products grow large in magnitude,
> pushing the softmax function into regions where it has extremely small gradients.
> To counteract this effect, we scale the dot products by sqrt(d_k).
The last (T, T) of the shape (B,h,T,T) is the matrix that represents the similarities
as the dot product between (q,k) from every q from sequence length T and k from the
sequence length T. The dimensions of q and k are both d_k, and q, k are expected to
follow normal distribution where the mean is 0 and variance is 1. The variance of the
two normal distributions q, k is expected to be d_k. Hence, standardize the (T,T)
with its standard deviation std=sqrt(d_k) so that the variance will be approximately 1.
Then, the later softmax will be smoothed out so that not to pick up higher value.
Args:
similarities: Similarities matrix shape (B, h, T, T)
d_k: dimension of the
Returns: scaled similarities matrix of shape (B, h, T, T)
"""
# --------------------------------------------------------------------------------
# Scaling factor to standardize (div by standard deviation) the product [email protected]
# of two zero centered normal distributions q, k. The variance of the product
# is head_size d_k. See https://stats.stackexchange.com/a/52699/105137.
# --------------------------------------------------------------------------------
std = torch.sqrt(torch.tensor(d_k, dtype=TYPE_FLOAT)) # standard deviation
# --------------------------------------------------------------------------------
# Scale similarities of each head by std so that the variance is approx 1.
# Scaling regularize the softmax output so as not to overfit to features, by which
# features in query and key can relate among themselves better.
# Otherwise, features with higher value will be peaked by softmax, (which is good
# for use as classification head but not for Bag of Words to incorporate features
# to make them related), hence only specific features in query and key will be
# connected.
# --------------------------------------------------------------------------------
scaled = similarities / std # scaled dot product
return scaled
def mask(
similarities: Tensor,
mask_matrix: Tensor
) -> Tensor:
"""
Args:
similarities: matrix to mask of shape (B,H,T,T)
mask_matrix: boolean matrix of which elements in (T,T) to mask fill.
Returns: masked similarity matrix
"""
# --------------------------------------------------------------------------------
# mask to make uni-direction (left to right only) for algorithm such as GPT.
# Skip masking for bi-directional e.g .BERT,
# --------------------------------------------------------------------------------
# exp(-inf) = 0 masks the similarities so that it will be uni-directional.
assert (
similarities.ndim == 4 and # (B,H,T,T)
similarities.shape[-2] == similarities.shape[-1] and
similarities.shape[-1] == mask_matrix.shape[-1]
)
masked = similarities.masked_fill(mask=mask_matrix, value=float('-inf'))
return masked
def calculate_attentionscalculate_attention_values(
similarities,
values
):
"""
For every q element, create a Bag of Words that encodes the relationships with
other elements (including itself) in T, using (q,k) relationship value as the
strength of the relationships.
Citation:
> On each of these projected versions of queries, keys and values we then perform
> the attention function in parallel, yielding d_v-dimensional output values.
```
bows = []
for row in similarities: # similarity matrix of shape (T,T)
bow = sum([ # bow:shape(d_v,)
# each column in row is (q,k) similarity score s
s*v for (s,v) in zip(row,values) # k:shape(), v:shape(d_v,)
= ])
bows.append(bow) # bows:shape(T,d_v)
```
Args:
similarities: q to k relationship strength matrix of shape (B, h, T, T)
values: elements of sequence with length T of shape (B, h, T, d_v)
Returns: Bag of Words for every q element of shape (B, h, T, d_v)
"""
return similarities @ values # (B,h,T,T) @ (B,h,T,d_v) -> (B,h,T,d_v)
class ScaledDotProductAttention(nn.Module):
"""
Class to implement Scaled Dot Product Attention (Figure 2 left in the paper).
"""
def __init__(self, do_mask: bool, max_time_steps: Optional[int]):
"""
Args:
max_time_steps: max sequence length or time steps T
"""
mask_matrix: Optional[Tensor]
super().__init__()
if do_mask:
mask_matrix = torch.tril(torch.ones(max_time_steps, max_time_steps)) == 0
else:
mask_matrix = None
self.register_buffer("mask_matrix", mask_matrix)
assert (
(not do_mask and self.mask_matrix is None) or
(do_mask and self.mask_matrix.ndim == 2 and self.mask_matrix.shape[-1] == max_time_steps)
)
def forward(
self,
q: Tensor,
k: Tensor,
v: Tensor,
):
"""Calculate the scaled dot product attention.
Args:
q: query of shape (B,h,T,d)
k: key of shape (B,h,T,d)
v: value of shape (B,h,T,d)
"""
# --------------------------------------------------------------------------------
# First MatMul in the Scaled Dot Product Attention to calculate the similarities
# matrix between (q,k) for every (q,k) combinations in Q, K.
# This is cartesian product matrix of shape (T, T) for every head and batch.
# The number of features in similarities matrix is B*H*T*T which will be
# (32 * 8 * 512 * 512) which is 64M. Each feature has 512 / H = 64 dimensions
# of float32, hence the size is 16G bytes of memory requirement.
# --------------------------------------------------------------------------------
similarities: Tensor = calculate_dot_product_similarities(
query=q,
key=k,
)
# --------------------------------------------------------------------------------
# Scale (standardize) the dot product similarity matrix with its standard deviation.
# --------------------------------------------------------------------------------
d_k = k.shape[-1] # head size
similarities = scale(similarities=similarities, d_k=d_k)
# --------------------------------------------------------------------------------
# Mask if required
# --------------------------------------------------------------------------------
if self.mask_matrix is not None:
similarities = mask(similarities=similarities, mask_matrix=self.mask_matrix)
# --------------------------------------------------------------------------------
# Normalize by softmax.
# --------------------------------------------------------------------------------
similarities = softmax(similarities, dim=-1)
# --------------------------------------------------------------------------------
# Second MatMul to generate attention value for each token in sequence of length T
# --------------------------------------------------------------------------------
attentions: Tensor = calculate_attentions(
similarities=similarities,
values=v
) # shape: (B,H,T,d)
return attentions
Andrej Karpathy explained by regarding a sentence as a graph as in CS25 I Stanford Seminar - Transformers United 2023: Introduction to Transformers w/ Andrej Karpathy.
