Attention and Differential Attention functions.

_.py

def Attention(X, W_q, W_k, W_v):
  Q = X @ W_q
  K = X @ W_k
  V = X @ W_v
  # Q, K, V: [b, n, d]
  s = 1 / sqrt(d)
  A = Q @ K.transpose(−1,−2) ∗ s
  return softmax(A) @ V


def DiffAttn(X, W_q, W_k, W_v, λ):
  Q1, Q2 = split(X @ W_q)
  K1, K2 = split(X @ W_k)
  V = X @ W_v
  # Qi, Ki: [b, n, d]; V: [b, n, 2d]
  s = 1 / sqrt(d)
  A1 = Q1 @ K1.transpose(−1,−2) ∗ s
  A2 = Q2 @ K2.transpose(−1,−2) ∗ s
  return (softmax(A1)− λ softmax(A2)) @ V