sparse_attention¶

paddle.nn.functional. sparse_attention ( query, key, value, sparse_csr_offset, sparse_csr_columns, name=None ) [源代码] ¶

对 Transformer 模块中的 Attention 矩阵进行了稀疏化，从而减少内存消耗和计算量。

其稀疏数据排布通过 CSR 格式表示，CSR 格式包含两个参数，offset 和 colunms。计算公式为：

\[result=softmax(\frac{ Q * K^T }{\sqrt{d}}) * V\]

其中，Q，K，V 表示注意力模块的三个输入参数。这三个参数的维度是一样的。d 代表这三个参数的最后一个维度的大小。

警告

目前该 API 只在 CUDA11.3 及以上版本中使用。

参数¶

query (Tensor) - 输入的 Tensor，代表注意力模块中的 query，这是一个 4 维 Tensor，形状为：[batch_size, num_heads, seq_len, head_dim]，数据类型为 float32 或 float64。

key (Tensor) - 输入的 Tensor，代表注意力模块中的 key，这是一个 4 维 Tensor，形状为：[batch_size, num_heads, seq_len, head_dim]，数据类型为 float32 或 float64。

value (Tensor) - 输入的 Tensor，代表注意力模块中的 value，这是一个 4 维 Tensor，形状为：[batch_size, num_heads, seq_len, head_dim]，数据类型为 float32 或 float64。

sparse_csr_offset (Tensor) - 输入的 Tensor，注意力模块中的稀疏特性，稀疏特性使用 CSR 格式表示，offset 代表矩阵中每一行非零元的数量。这是一个 3 维 Tensor，形状为：[batch_size, num_heads, seq_len + 1]，数据类型为 int32。

sparse_csr_columns (Tensor) - 输入的 Tensor，注意力模块中的稀疏特性，稀疏特性使用 CSR 格式表示，colunms 代表矩阵中每一行非零元的列索引值。这是一个 3 维 Tensor，形状为：[batch_size, num_heads, sparse_nnz]，数据类型为 int32。

返回¶

Tensor，代表注意力模块的结果。这是一个 4 维 Tensor，形状为：[batch_size, num_heads, seq_len, head_dim]，数据类型为 float32 或 float64。

代码示例¶

          # required: skiptest
import paddle

paddle.disable_static()

# `query`, `key` and `value` all have shape [1, 1, 4, 2]
query = paddle.to_tensor([[[[0, 1, ], [2, 3],
                            [0, 1], [2, 3]]]], dtype="float32")
key = paddle.to_tensor([[[[0, 1], [2, 3],
                        [0, 1], [2, 3]]]], dtype="float32")
value = paddle.to_tensor([[[[0, 1], [2, 3],
                            [0, 1], [2, 3]]]], dtype="float32")


offset = paddle.to_tensor([[[0, 2, 4, 6, 8]]], dtype="int32")
columns = paddle.to_tensor([[[0, 1, 0, 1, 2, 3, 2, 3]]], dtype="int32")

print(offset.shape)  # (1, 1, 5)
print(columns.shape)  # (1, 1, 8)

key_padding_mask = paddle.to_tensor([[1, 1, 1, 0]], dtype="float32")
attention_mask = paddle.to_tensor([[1, 0, 1, 1],
                                [1, 1, 1, 1],
                                [1, 1, 1, 1],
                                [1, 1, 1, 1]], dtype="float32")
output_mask = paddle.nn.functional.sparse_attention(query, key,
                                                    value, offset, columns,
                                                    key_padding_mask=key_padding_mask,
                                                    attn_mask=attention_mask)
print(output_mask)
# Tensor(shape=[1, 1, 4, 2], dtype=float32, place=Place(gpu:0), stop_gradient=False,
#        [[[[0.        , 1.        ],
#           [1.99830270, 2.99830270],
#           [0.        , 1.        ],
#           [0.        , 1.        ]]]])

output = paddle.nn.functional.sparse_attention(query, key,
                                            value, offset, columns)
print(output)
# Tensor(shape=[1, 1, 4, 2], dtype=float32, place=Place(gpu:0), stop_gradient=False,
#        [[[[1.60885942, 2.60885954],
#           [1.99830270, 2.99830270],
#           [1.60885942, 2.60885954],
#           [1.99830270, 2.99830270]]]])

         