jax.nn.dot_product_attention

jax.nn.dot_product_attention#

jax.nn.dot_product_attention(query, key, value, bias=None, mask=None, *, scale=None, is_causal=False, query_seq_lengths=None, key_value_seq_lengths=None, implementation=None)[source]#

缩放点积注意力函数。

计算查询、键和值张量上的注意力函数

\[\mathrm{Attention}(Q, K, V)=\mathrm{softmax}(\frac{QK^T}{\sqrt{d_k}})V\]

如果我们将 logits 定义为 \(QK^T\) 的输出，并将 probs 定义为 \(softmax\) 的输出。

在整个函数中，我们使用以下大写字母来表示数组的形状

B = batch size
S = length of the key/value (source)
T = length of the query (target)
N = number of attention heads
H = dimensions of each attention head
K = number of key/value heads
G = number of groups, which equals to N // K
参数：

  • query (ArrayLike) – 查询数组；形状 (BTNH|TNH)

  • key (ArrayLike) – 密钥数组：形状为 (BSKH|SKH)。当 K 等于 N 时，执行多头注意力 (MHA https://arxiv.org/abs/1706.03762)。否则，如果 N 是 K 的倍数，则执行分组查询注意力 (GQA https://arxiv.org/abs/2305.13245)；如果 K == 1（GQA 的一种特殊情况），则执行多查询注意力 (MQA https://arxiv.org/abs/1911.02150)。

  • value (ArrayLike) – 值数组，应与 key 数组具有相同的形状。

  • bias (ArrayLike | None | None) – 可选，要添加到 logits 的偏差数组；形状必须为 4D 且可广播到 (BNTS|NTS)。

  • mask (ArrayLike | None | None) – 可选，用于过滤 logits 的掩码数组。它是一个布尔掩码，其中 True 表示元素应参与注意力。对于加性掩码，用户应将其传递给 bias。形状必须为 4D 且可广播到 (BNTS|NTS)。

  • scale (float | None | None) – logits 的缩放因子。如果为 None，则缩放因子将设置为查询头部维度（即 H）的平方根的倒数。

  • is_causal (bool) – 如果为真，则应用因果注意力。请注意，某些实现（如 xla）将生成一个掩码张量并将其应用于 logits 以掩盖注意力矩阵的非因果部分，但其他实现（如 cudnn）将避免计算非因果区域，从而提高速度。

  • query_seq_lengths (ArrayLike | None | None) – 查询的序列长度的 int32 数组；形状为 (B)

  • key_value_seq_lengths (ArrayLike | None | None) – 密钥和值的序列长度的 int32 数组；形状为 (B)

  • implementation (Literal['xla', 'cudnn'] | None | None) – 用于控制使用哪个实现后端的字符串。支持的字符串为 xla、cudnn（cuDNN 闪存注意力）。默认为 None，这将自动选择最佳可用后端。请注意，cudnn 只支持一部分形状/数据类型，如果不受支持，则会抛出异常。

返回值:

与 query 形状相同的注意力输出数组。

返回类型:

数组

内容