def DotProductAttention(query, key, value, mask, dropout, mode, rng):
"""Core dot product self-attention.
Args:
query: array of representations
key: array of representations
value: array of representations
mask: attention-mask, gates attention
dropout: float: dropout rate
mode: 'eval' or 'train': whether to use dropout
rng: JAX PRNGKey: subkey for disposable use
Returns:
Self attention for q, k, v arrays.
"""
depth = np.shape(query)[-1]
dots = np.matmul(query, np.swapaxes(key, -1, -2)) / np.sqrt(depth)
if mask is not None:
# TODO(kitaev): workaround for https://github.com/google/jax/issues/850
# We must ensure that both mask and the -1e9 constant have a data dependency
# on the input. Broadcasted copies of these use a lot of memory, so they
# should be computed at runtime (rather than being global constants).
if math.backend_name() == 'jax':
mask = jax.lax.tie_in(dots, mask)
# JAX's `full_like` already ties in -1e9 to dots.
dots = np.where(mask, dots, np.full_like(dots, -1e9))
# Softmax.
dots = np.exp(dots - math.logsumexp(dots, axis=-1, keepdims=True))
if dropout >= 1.0:
raise ValueError('Dropout rates must be lower than 1.')
if dropout is not None and dropout > 0.0 and mode == 'train':
keep = math.random.bernoulli(rng, 1.0 - dropout, dots.shape)
dots = np.where(keep, dots / (1.0 - dropout), np.zeros_like(dots))
out = np.matmul(dots, value)
return out
def forward_unbatched(self, x, *, weights, state, update_state):
del update_state
if self.share_qk:
w_q, w_v, w_o = weights
else:
w_q, w_k, w_v, w_o = weights
q = np.matmul(x, w_q)
k = None
if not self.share_qk:
k = np.matmul(x, w_k)
v = np.matmul(x, w_v)
mask_fn = functools.partial(mask_self_attention,
causal=self.causal,
exclude_self=self.share_qk)
q_info = kv_info = jax.lax.tie_in(x, np.arange(q.shape[-2]))
o, _ = attend(
q,
k,
v,
q_chunk_len=self.chunk_len,
kv_chunk_len=self.chunk_len,
n_chunks_before=self.n_chunks_before,
n_chunks_after=self.n_chunks_after,
mask_fn=mask_fn,
q_info=q_info,
kv_info=kv_info,
dropout=self.attention_dropout,
rng=None, # TODO(kitaev): support RNG
)
out = np.matmul(o, w_o)
return out, state
def forward_unbatched(self, q_antecedent, kv_antecedent, mask=None, *,
weights, state, rng, update_state):
del update_state
attend_rng, output_rng = jax.random.split(rng)
w_q, w_k, w_v, w_o = weights
q = np.matmul(q_antecedent, w_q)
k = np.matmul(kv_antecedent, w_k)
v = np.matmul(kv_antecedent, w_v)
if not self.masked:
assert mask is None
q_info = kv_info = mask_fn = None
else:
# mask is a boolean array (True means "is valid token")
assert mask is not None
q_info = None
kv_info = (~mask).astype(np.int32) # pylint: disable=invalid-unary-operand-type
def mask_fn(dots, q_info, kv_info):
del q_info
mask = jax.lax.convert_element_type(kv_info, np.float32)
dots = dots - 1e9 * mask
return dots
o, _ = attend(
q, k, v,
mask_fn=mask_fn, q_info=q_info, kv_info=kv_info,
dropout=self.attention_dropout, rng=attend_rng,
)
out = np.matmul(o, w_o)
out = apply_broadcasted_dropout(out, self.output_dropout, output_rng)
return out, state
def forward(self, x, weights):
if self._use_bias:
if not isinstance(weights, (tuple, list)):
raise ValueError(f'Weights should be a (w, b) tuple or list; '
f'instead got: {weights}')
w, b = weights
return jnp.matmul(x, w) + b # Affine map.
else:
w = weights
return jnp.matmul(x, w) # Linear map.
def DotProductAttention(queries, keys, values, mask, dropout, mode, rng):
"""Computes new activations via masked attention-weighted sum of values.
This function is the core of the attention mechanism. It:
- computes per-head attention weights from per-head `(queries, keys)`,
- applies `mask` to screen out positions that come from padding tokens,
- optionally applies dropout to attention weights, and
- uses attention weights to combine per-head `values` vectors.
Args:
queries: Per-head activations representing attention queries.
keys: Per-head activations representing attention keys.
values: Per-head activations to be combined by computed attention weights.
mask: Mask that distinguishes positions with real content vs. padding.
dropout: Probababilistic rate for dropout applied to attention activations
(based on query-key pairs) before dotting them with values.
mode: Either 'train' or eval'. Dropout applies only in 'train' mode.
rng: Single-use random number generator (JAX PRNG key).
Returns:
Per-head activations resulting from masked per-head attention-weighted
sum of per-head values.
"""
d_feature = queries.shape[-1]
dots = jnp.matmul(queries, jnp.swapaxes(keys, -1,
-2)) / jnp.sqrt(d_feature)
if mask is not None:
# TODO(kitaev): workaround for https://github.com/google/jax/issues/850
# We must ensure that both mask and the -1e9 constant have a data dependency
# on the input. Broadcasted copies of these use a lot of memory, so they
# should be computed at runtime (rather than being global constants).
if math.backend_name() == 'jax':
mask = jax.lax.tie_in(dots, mask)
# JAX's `full_like` already ties in -1e9 to dots.
dots = jnp.where(mask, dots, jnp.full_like(dots, -1e9))
# Softmax.
dots = jnp.exp(dots - math.logsumexp(dots, axis=-1, keepdims=True))
if dropout >= 1.0:
raise ValueError('Dropout rates must be lower than 1.')
if dropout is not None and dropout > 0.0 and mode == 'train':
keep = math.random.bernoulli(rng, 1.0 - dropout, dots.shape)
dots = jnp.where(keep, dots / (1.0 - dropout), jnp.zeros_like(dots))
out = jnp.matmul(dots, values)
return out
def log_gaussian_diag_pdf(x, mu, diag_sigma): # pylint: disable=invalid-name
"""Compute log N(x | mu, eye(diag_sigma))."""
a = mu.shape[-1] * jnp.log(2 * jnp.pi)
b = jnp.sum(jnp.log(diag_sigma), axis=-1)
y = x - mu / diag_sigma
y = jnp.expand_dims(y, axis=-1)
xm = jnp.expand_dims(x - mu, axis=-2)
c = jnp.matmul(xm, y)
c = jnp.squeeze(jnp.squeeze(c, axis=-1), axis=-1)
return -0.5 * (a + b + c)
def log_gaussian_pdf(x, mu, sigma): # pylint: disable=invalid-name
"""Compute log N(x | mu, sigma)."""
a = mu.shape[-1] * jnp.log(2 * jnp.pi)
_, b = jnp.linalg.slogdet(sigma)
y = jnp.linalg.solve(sigma, x - mu)
y = jnp.expand_dims(y, axis=-1)
xm = jnp.expand_dims(x - mu, axis=-2)
c = jnp.matmul(xm, y)
c = jnp.squeeze(jnp.squeeze(c, axis=-1), axis=-1)
return -0.5 * (a + b + c)
def forward_unbatched(self, x, mask=None, *, weights, state, update_state):
del update_state
if self.share_qk:
w_q, w_v, w_o = weights
else:
w_q, w_k, w_v, w_o = weights
q = np.matmul(x, w_q)
k = None
if not self.share_qk:
k = np.matmul(x, w_k)
v = np.matmul(x, w_v)
mask_fn = functools.partial(mask_self_attention,
causal=self.causal,
exclude_self=self.share_qk,
masked=self.masked)
q_info = kv_info = jax.lax.tie_in(x, np.arange(q.shape[-2]))
assert (mask is not None) == self.masked
if self.masked:
# mask is a boolean array (True means "is valid token")
ones_like_mask = jax.lax.tie_in(x,
np.ones_like(mask, dtype=np.int32))
kv_info = kv_info * np.where(mask, ones_like_mask, -ones_like_mask)
o, _ = attend(
q,
k,
v,
q_chunk_len=self.chunk_len,
kv_chunk_len=self.chunk_len,
n_chunks_before=self.n_chunks_before,
n_chunks_after=self.n_chunks_after,
mask_fn=mask_fn,
q_info=q_info,
kv_info=kv_info,
dropout=self.attention_dropout,
rng=None, # TODO(kitaev): support RNG
)
out = np.matmul(o, w_o)
return out, state
def forward_unbatched(self, x, *, weights, state, update_state):
w_q, w_v, w_o = weights
q = np.matmul(x, w_q)
v = np.matmul(x, w_v)
if update_state:
_, old_rng = state
rng = jax.random.fold_in(old_rng, 0)
hash_rng = jax.random.fold_in(rng, 1)
buckets = self.hash_vectors(q, hash_rng)
state = (buckets, rng)
else:
buckets, rng = state
rng = jax.random.fold_in(rng, 2)
seqlen = x.shape[0]
assert int(buckets.shape[0]) == self.n_hashes * seqlen
ticker = jax.lax.tie_in(x, np.arange(self.n_hashes * seqlen))
buckets_and_t = seqlen * buckets + (ticker % seqlen)
buckets_and_t = jax.lax.stop_gradient(buckets_and_t)
# Hash-based sort ("s" at the start of variable names means "sorted")
sbuckets_and_t, sticker = jax.lax.sort_key_val(buckets_and_t,
ticker,
dimension=-1)
_, undo_sort = jax.lax.sort_key_val(sticker, ticker, dimension=-1)
sbuckets_and_t = jax.lax.stop_gradient(sbuckets_and_t)
sticker = jax.lax.stop_gradient(sticker)
undo_sort = jax.lax.stop_gradient(undo_sort)
st = (sticker % seqlen)
sq = np.take(q, st, axis=0)
sv = np.take(v, st, axis=0)
mask_fn = functools.partial(mask_self_attention,
causal=self.causal,
exclude_self=True)
q_info = st
so, slogits = attend(
sq,
k=None,
v=sv,
q_chunk_len=self.chunk_len,
n_chunks_before=self.n_chunks_before,
n_chunks_after=self.n_chunks_after,
mask_fn=mask_fn,
q_info=q_info,
dropout=self.attention_dropout,
rng=rng,
)
def unsort_for_output_impl(so, slogits):
o = np.take(so, undo_sort, axis=0)
# Sorting is considerably faster than gather, but first we need to get the
# XLA compiler to abandon the idea of fusing this sort with the input sort
# (which introduces a computation cycle and leads to a crash).
# TODO(kitaev): remove "sticker_" variable if XLA is fixed.
sticker_ = sticker + jax.lax.convert_element_type(
slogits[0] > 0, sticker.dtype)
_, logits = jax.lax.sort_key_val(sticker_, slogits, dimension=-1)
return o, logits
def unsort_for_output_vjp(so, slogits):
"""Custom gradient for unsort_for_output."""
so = jax.lax.stop_gradient(so)
slogits = jax.lax.stop_gradient(slogits)
o, logits = unsort_for_output_impl(so, slogits)
def vjpfun(o_logits_grads):
so_grad = np.take(o_logits_grads[0], sticker, axis=0)
# TODO(kitaev): this exists to match the forward pass, but I'm not sure
# if it's actually required.
buckets_and_t_ = buckets_and_t + jax.lax.convert_element_type(
o_logits_grads[1][0] > 0, buckets_and_t.dtype)
_, slogits_grad = jax.lax.sort_key_val(buckets_and_t_,
o_logits_grads[1],
dimension=-1)
return (so_grad, slogits_grad)
return (o, logits), vjpfun
unsort_for_output = jax.custom_transforms(unsort_for_output_impl)
jax.defvjp_all(unsort_for_output, unsort_for_output_vjp)
o, logits = unsort_for_output_impl(so, slogits)
if self.n_hashes > 1:
o = np.reshape(o, (self.n_hashes, seqlen, o.shape[-1]))
logits = np.reshape(logits, (self.n_hashes, seqlen, 1))
probs = np.exp(logits - logsumexp(logits, axis=0, keepdims=True))
o = np.sum(o * probs, axis=0)
assert o.shape == (seqlen, w_v.shape[-1])
out = np.matmul(o, w_o)
return out, state
def attend(
q,
k=None,
v=None,
q_chunk_len=None,
kv_chunk_len=None,
n_chunks_before=0,
n_chunks_after=0,
mask_fn=None,
q_info=None,
kv_info=None,
dropout=0.0,
rng=None,
):
"""Dot-product attention, with optional chunking and/or masking.
Args:
q: Query vectors, shape [q_len, d_qk]
k: Key vectors, shape [kv_len, d_qk]; or None
v: Value vectors, shape [kv_len, d_v]
q_chunk_len: Set to non-zero to enable chunking for query vectors
kv_chunk_len: Set to non-zero to enable chunking for key/value vectors
n_chunks_before: Number of adjacent previous chunks to attend to
n_chunks_after: Number of adjacent subsequent chunks to attend to
mask_fn: TODO(kitaev) doc
q_info: Query-associated metadata for masking
kv_info: Key-associated metadata for masking
dropout: Dropout rate
rng: RNG for dropout
Returns:
A tuple (output, dots_logsumexp). The output has shape [q_len, d_v], and
dots_logsumexp has shape [q_len]. The logsumexp of the attention
probabilities is useful for combining multiple rounds of attention (as in
LSH attention).
"""
assert v is not None
share_qk = (k is None)
if q_info is None:
q_info = np.arange(q.shape[-2])
if kv_info is None and not share_qk:
kv_info = np.arange(v.shape[-2])
# Split q/k/v into chunks along the time axis, if desired.
if q_chunk_len is not None:
q = np.reshape(q, (-1, q_chunk_len, q.shape[-1]))
q_info = np.reshape(q_info, (-1, q_chunk_len))
if share_qk:
assert kv_chunk_len is None or kv_chunk_len == q_chunk_len
k = q
kv_chunk_len = q_chunk_len
kv_info = q_info
elif kv_chunk_len is not None:
k = np.reshape(k, (-1, kv_chunk_len, k.shape[-1]))
kv_info = np.reshape(kv_info, (-1, kv_chunk_len))
if kv_chunk_len is not None:
v = np.reshape(v, (-1, kv_chunk_len, v.shape[-1]))
if share_qk:
k = length_normalized(k)
k = k / np.sqrt(k.shape[-1])
# Optionally include adjacent chunks.
if q_chunk_len is not None or kv_chunk_len is not None:
assert q_chunk_len is not None and kv_chunk_len is not None
else:
assert n_chunks_before == 0 and n_chunks_after == 0
k = look_adjacent(k, n_chunks_before, n_chunks_after)
v = look_adjacent(v, n_chunks_before, n_chunks_after)
kv_info = look_adjacent(kv_info, n_chunks_before, n_chunks_after)
# Dot-product attention.
dots = np.matmul(q, np.swapaxes(k, -1, -2))
# Masking
if mask_fn is not None:
dots = mask_fn(dots, q_info[..., :, None], kv_info[..., None, :])
# Softmax.
dots_logsumexp = logsumexp(dots, axis=-1, keepdims=True)
dots = np.exp(dots - dots_logsumexp)
if dropout > 0.0:
assert rng is not None
# Dropout is broadcast across the bin dimension
dropout_shape = (dots.shape[-2], dots.shape[-1])
# TODO(kitaev): verify that tie-in is safe to remove (in light of jax fix)
keep_prob = jax.lax.tie_in(dots, 1.0 - dropout)
keep = jax.random.bernoulli(rng, keep_prob, dropout_shape)
multiplier = keep.astype(dots.dtype) / jax.lax.tie_in(keep, keep_prob)
dots = dots * multiplier
# The softmax normalizer (dots_logsumexp) is used by multi-round LSH attn.
out = np.matmul(dots, v)
out = np.reshape(out, (-1, out.shape[-1]))
dots_logsumexp = np.reshape(dots_logsumexp, (-1, ))
return out, dots_logsumexp
def forward_unbatched(self, x, *, weights, state, update_state):
w_q, w_v, w_o = weights
q = np.matmul(x, w_q)
v = np.matmul(x, w_v)
if update_state:
_, old_rng = state
rng = jax.random.fold_in(old_rng, 0)
hash_rng = jax.random.fold_in(rng, 1)
buckets = self.hash_vectors(q, hash_rng)
state = (buckets, rng)
else:
buckets, rng = state
rng = jax.random.fold_in(rng, 2)
seqlen = x.shape[0]
assert int(buckets.shape[0]) == self.n_hashes * seqlen
ticker = jax.lax.tie_in(x, np.arange(self.n_hashes * seqlen))
buckets_and_t = seqlen * buckets + (ticker % seqlen)
buckets_and_t = jax.lax.stop_gradient(buckets_and_t)
# Hash-based sort ("s" at the start of variable names means "sorted")
sbuckets_and_t, sticker = jax.lax.sort_key_val(buckets_and_t,
ticker,
dimension=-1)
_, undo_sort = jax.lax.sort_key_val(sticker, ticker, dimension=-1)
sbuckets_and_t = jax.lax.stop_gradient(sbuckets_and_t)
sticker = jax.lax.stop_gradient(sticker)
undo_sort = jax.lax.stop_gradient(undo_sort)
st = (sticker % seqlen)
sq = np.take(q, st, axis=0)
sv = np.take(v, st, axis=0)
mask_fn = functools.partial(mask_self_attention,
causal=self.causal,
exclude_self=True)
q_info = st
so, slogits = attend(
sq,
k=None,
v=sv,
q_chunk_len=self.chunk_len,
n_chunks_before=self.n_chunks_before,
n_chunks_after=self.n_chunks_after,
mask_fn=mask_fn,
q_info=q_info,
dropout=self.attention_dropout,
rng=rng,
)
# np.take(so, undo_sort, axis=0); np.take(slogits, undo_sort, axis=0) would
# also work, but these helpers include performance optimizations for TPU.
o = permute_via_gather(so, undo_sort, sticker, axis=0)
logits = permute_via_sort(slogits, sticker, buckets_and_t, axis=-1)
if self.n_hashes > 1:
o = np.reshape(o, (self.n_hashes, seqlen, o.shape[-1]))
logits = np.reshape(logits, (self.n_hashes, seqlen, 1))
probs = np.exp(logits - logsumexp(logits, axis=0, keepdims=True))
o = np.sum(o * probs, axis=0)
assert o.shape == (seqlen, w_v.shape[-1])
out = np.matmul(o, w_o)
return out, state