tf.keras.layers.MultiHeadAttention

Слой многоголового внимания.

Наследуется от: Layer, Module

tf.keras.layers.MultiHeadAttention(
    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, **kwargs
)

Это реализация многоголового внимания, основанная на статье "Внимание — всё, что вам нужно". Если query, key, value совпадают, то это самовнимание. Каждый временной шаг в query обращается к соответствующей последовательности в key, и возвращает вектор фиксированной ширины.

Этот слой сначала проецирует query, key и value. Это (по сути) список тензоров длиной num_attention_heads, где соответствующие формы — [размер_пакета, <размер_запроса>, размер_ключа], [размер_пакета, <размер_ключа>, размер_ключа], [размер_пакета, <размер_ключа>, размер_значения].

Затем тензоры запроса и ключа умножаются точечно и масштабируются. Они подвергаются softmax, чтобы получить вероятности внимания. Затем тензоры значений интерполируются с помощью этих вероятностей, а затем конкатенируются обратно в один тензор.

Наконец, тензор результата с последним измерением как размер_значения может принять линейное проектирование и вернуть результат.

Примеры:

Выполняет 1D перекрёстное внимание над двумя последовательными входными данными с маской внимания. Возвращает дополнительные веса внимания по головам.

layer = MultiHeadAttention(num_heads=2, key_dim=2)
target = tf.keras.Input(shape=[8, 16])
source = tf.keras.Input(shape=[4, 16])
output_tensor, weights = layer(target, source,
                               return_attention_scores=True)
print(output_tensor.shape)
(None, 8, 16)
print(weights.shape)
(None, 2, 8, 4)

Выполняет 2D самовнимание над тензором входных данных 5D по осям 2 и 3.

layer = MultiHeadAttention(num_heads=2, key_dim=2, attention_axes=(2, 3))
input_tensor = tf.keras.Input(shape=[5, 3, 4, 16])
output_tensor = layer(input_tensor, input_tensor)
print(output_tensor.shape)
(None, 5, 3, 4, 16)

Аргументы вызова:

• query: Запрос Tensor формы [B, T, dim].
• value: Значение Tensor формы [B, S, dim].
• key: Необязательный ключ Tensor формы [B, S, dim]. Если не указан, будет использован value для обоих key и value, что является наиболее распространённым случаем.
• attention_mask: Булева маска формы [B, T, S], которая предотвращает внимание к определённым позициям.
• return_attention_scores: Булево значение, указывающее, должен ли вывод быть выводом внимания (True), или (вывод_внимания, баллы_внимания) (False). По умолчанию False.
• training: Булево значение Python, указывающее, должен ли слой работать в режиме обучения (добавляя дропаут) или в режиме вывода (без дропаута). По умолчанию использует режим обучения родительского слоя/модели или False (режим вывода), если родительского слоя нет.