Маскировка MultiHeadAttention с тензорным потоком

Question

Я пытался создать пользовательскую маску для целевых комбинаций запросов и ключей для моего слоя MultiHeadAttention, но не могу понять, как использовать эту маску слоя.

Вот пример с фиктивным набором данных (размер пакета 1):

      key     = tf.ones([1, 32 , 128])
mask    = tf.concat([
    tf.concat([tf.zeros([16 , 16]) , tf.zeros([16 , 16]) ] , 0) ,
    tf.concat([tf.zeros([16 , 16]) , tf.ones([16 , 16])  ] , 0) ] , 1)
mask    = mask[tf.newaxis, tf.newaxis, : , : ]


# key shape  -> ( 1 , 32 , 128 )
# mask shape -> ( 1 , 1,  32 , 32 )

когда я печатаю mask[0][0].numpy()Я получил :

Теперь используем следующий слой (1 голова, внимание к себе):

      mha_layer =  tf.keras.layers.MultiHeadAttention( num_heads=1, key_dim=128 )
attention_output, attention_scores = mha_layer(  key , key , attention_mask=mask  ,  return_attention_scores=True)

Я получаю следующие оценки внимания ( attention_scores[0][0].numpy()):

Здесь темно-фиолетовый цвет соответствует 0,0, желтый — 0,06, зелено-синий — 0,03.

Я ожидал бы, что зелено-синяя часть будет 0,0 с из-за маскировки.

Я неправильно использую маскировку? или невозможно замаскировать целые запросы/ключи?

Я надеюсь, что мой вопрос имеет смысл и что он не слишком очевиден. Заранее спасибо, если поможете :)