Внимание по головам в TensorFlow

Класс MultiHeadAttention считает несколько параллельных голов внимания и возвращает один тензор той же длины, что и входная цепочка. При создании задают число голов и размерность пространства ключей. На вход подают запрос и значения; ключи можно передать отдельно или взять те же, что и значения.

Создадим слой с двумя головами и короткую цепочку из двух шагов в одной партии. Пропустим её через слой с одинаковыми запросом и значениями и выведем форму ответа:

import tensorflow as tf

layer = tf.keras.layers.MultiHeadAttention(
	num_heads=2, key_dim=4)
seq = tf.constant([[[1., 2., 3., 4., 5., 6., 7., 8.],
                    [0., 1., 0., 1., 0., 1., 0., 1.]]])
out = layer(seq, seq)
print(out.shape)
# выведет (1, 2, 8)

Первая ось - размер партии, вторая - число шагов, третья - длина вектора на каждом шаге. Если длина вектора совпадает с тем, что ожидает слой, форма не меняется, меняется только содержимое ячеек.

Три шага и одна голова дают тот же принцип: партия из одной строки, три позиции, два числа в каждом шаге:

import tensorflow as tf

layer = tf.keras.layers.MultiHeadAttention(
	num_heads=1, key_dim=2)
seq = tf.constant([[[1., 0.], [0., 1.], [1., 1.]]])
out = layer(seq, seq)
print(out.shape)
# выведет (1, 3, 2)

Соберите одну партию с двумя шагами и вектором длины шесть. Создайте слой с тремя головами и размерностью ключей два. Выведите форму тензора после прогона цепочки через себя.

Задайте цепочку из четырёх шагов с вектором длины четыре в одной партии. Узнайте форму ответа при двух головах и размерности ключей два.

Подготовьте две одинаковые цепочки из трёх шагов и восьми чисел в каждом. Примените слой с четырьмя головами и ключами размерности четыре и выведите форму результата.