CachedMultiHeadAttention 層

[來源]

CachedMultiHeadAttention 類別

keras_nlp.layers.CachedMultiHeadAttention(
    num_heads,
    key_dim,
    value_dim=None,
    dropout=0.0,
    use_bias=True,
    output_shape=None,
    attention_axes=None,
    kernel_initializer="glorot_uniform",
    bias_initializer="zeros",
    kernel_regularizer=None,
    bias_regularizer=None,
    activity_regularizer=None,
    kernel_constraint=None,
    bias_constraint=None,
    seed=None,
    **kwargs
)

具有快取支援的 MultiHeadAttention 層。

此層適用於自回歸解碼。它可以用於快取解碼器自注意力和交叉注意力。正向傳播可以在以下三種模式之一中進行

• 無快取，與一般的多頭注意力相同。
• 靜態快取（cache_update_index 為 None）。在這種情況下，將使用快取的鍵/值投影，並且輸入值將被忽略。
• 更新的快取（cache_update_index 不是 None）。在這種情況下，將使用輸入計算新的鍵/值投影，並將其拼接至指定索引處的快取中。

請注意，快取僅在推論期間有用，不應在訓練期間使用。

我們在下方使用符號 B、T、S，其中 B 是批次維度，T 是目標序列長度，S 是來源序列長度。請注意，在生成解碼期間，T 通常為 1（您正在生成長度為 1 的目標序列以預測下一個標記）。

呼叫參數

• query：形狀為 (B, T, dim) 的查詢 Tensor。
• value：形狀為 (B, S*, dim) 的值 Tensor。如果 cache 為 None，則 S* 必須等於 S 並與 attention_mask 的形狀匹配。如果 cache 不為 None，則 S* 可以是任何小於 S 的長度，並且計算出的值將在 cache_update_index 處拼接至 cache 中。
• key：形狀為 (B, S*, dim) 的可選鍵 Tensor。如果 cache 為 None，則 S* 必須等於 S 並與 attention_mask 的形狀匹配。如果 cache 不為 None，則 S* 可以是任何小於 S 的長度，並且計算出的值將在 cache_update_index 處拼接至 cache 中。
• attention_mask：形狀為 (B, T, S) 的布林遮罩。attention_mask 會阻止注意力集中在某些位置。布林遮罩指定哪些查詢元素可以關注哪些鍵值元素，1 表示關注，0 表示不關注。廣播可以發生在缺少的批次維度和頭維度。
• cache：密集浮點數張量。鍵值快取，形狀為 [B, 2, S, num_heads, key_dims]，其中 S 必須與 attention_mask 形狀一致。此參數旨在生成過程中使用，以避免重新計算中間狀態。
• cache_update_index：整數或整數張量，用於更新 cache 的索引（通常是在運行生成時正在處理的當前標記的索引）。如果在設置 cache 時 cache_update_index=None，則不會更新快取。
• training：一個布林值，表示該層應處於訓練模式還是推理模式。

回傳值

一個 (attention_output, cache) 元組。attention_output 是計算結果，形狀為 (B, T, dim)，其中 T 表示目標序列形狀，如果 output_shape 為 None，則 dim 是查詢輸入的最後一個維度。否則，多頭輸出將投影到 output_shape 指定的形狀。cache 是更新後的快取。