So basically a matmul is an einsum that drops the last coordinate of
the first operand and the second to last coordinate of the second
operand. Matrix multiplications really are sequences of dot products!
Linear algebra is slowly and painfully coming back to me.
Attached is a transcription of huggingface's perceiver attention that
works with the same example data. The 'keys/queries/values' axis ends
up being the sequence axis. They permute the matrices to exclude the
heads dimension so the dot products can be done with a normal matmul
rather than einsum and its string parsing.
import torch, math
def attention(queries, keys, values):
queries = queries.permute(0, 2, 1, 3)
keys = keys.permute(0, 2, 1, 3)
values = values.permute(0, 2, 1, 3)
# Take the dot product between the queries and keys to get the raw attention scores.
attention_scores = torch.matmul(queries, keys.transpose(-1, -2))
batch_size, num_heads, seq_len, q_head_dim = queries.shape
_, _, _, v_head_dim = values.shape
hiddens = num_heads * v_head_dim
attention_scores = attention_scores / math.sqrt(q_head_dim)
# Normalize the attention scores to probabilities
attention_probs = torch.nn.Softmax(dim=-1)(attention_scores)
context_layer = torch.matmul(attention_probs, values)
context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
new_context_layer_shape = context_layer.size()[:-2] + (hiddens,)
context_layer = context_layer.view(*new_context_layer_shape)
return context_layer
if __name__ == '__main__':
import jax
queries, keys, values = torch.from_numpy(jax.random.normal(jax.random.PRNGKey(0), (3, 1, 64, 8, 16)).to_py())
out = attention(queries, keys, values)
print(out)
