10.5 Многоголовое внимание

На практике при одном и том же наборе запросов, ключей и значений мы хотим, чтобы модель могла, опираясь на один и тот же механизм внимания, выучивать разные виды поведения, а затем объединять их как знания, чтобы захватывать зависимости разного масштаба внутри последовательности (например, краткосрочные и долгосрочные зависимости). Поэтому полезно позволить механизму внимания комбинировать разные подпространства представлений (representation subspaces) запросов, ключей и значений.

Для этого вместо одной агрегации внимания мы можем преобразовать запросы, ключи и значения с помощью $h$ независимо выученных линейных проекций (linear projections). Затем эти $h$ преобразованных запросов, ключей и значений параллельно подаются в агрегацию внимания. Наконец, выходы $h$ агрегаций внимания конкатенируются и преобразуются ещё одной обучаемой линейной проекцией, чтобы получить итоговый выход. Такая конструкция называется многоголовым вниманием (multihead attention) (Vaswani et al., 2017). Каждая из $h$ агрегаций внимания называется головой (head). На рис. 10.5.1 показано многоголовое внимание, в котором обучаемые линейные преобразования реализованы полносвязными слоями.

Многоголовое внимание: несколько голов конкатенируются, а затем проходят линейное преобразование

Рис. 10.5.1. Многоголовое внимание: несколько голов конкатенируются, а затем проходят линейное преобразование

10.5.1 Модель

Прежде чем реализовывать многоголовое внимание, опишем модель математически. Даны запрос $q \in R^{d_{q}}$ , ключ $k \in R^{d_{k}}$ и значение $v \in R^{d_{v}}$ . Каждая голова внимания $h_{i}$ ( $i = 1, \dots, h$ ) вычисляется так:

$h_{i} = f (W_{i}^{(q)} q, W_{i}^{(k)} k, W_{i}^{(v)} v) \in R^{p_{v}},$

где обучаемые параметры включают $W_{i}^{(q)} \in R^{p_{q} \times d_{q}}$ , $W_{i}^{(k)} \in R^{p_{k} \times d_{k}}$ и $W_{i}^{(v)} \in R^{p_{v} \times d_{v}}$ , а $f$ — функция агрегации внимания. Функцией $f$ может быть аддитивное внимание или масштабированное скалярное произведение внимания из раздела 10.3. Выход многоголового внимания проходит ещё одно линейное преобразование, которое соответствует конкатенации $h$ голов, поэтому обучаемый параметр — $W_{o} \in R^{p_{o} \times h p_{v}}$ :

$W_{o} h_{1} ⋮ h_{h} \in R^{p_{o}} .$

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

MXNet:

from d2l import mxnet as d2l
import math
from mxnet import autograd, np, npx
from mxnet.gluon import nn
npx.set_np()

PyTorch:

from d2l import torch as d2l
import math
import torch
from torch import nn

TensorFlow:

from d2l import tensorflow as d2l
import tensorflow as tf

PaddlePaddle:

from d2l import paddle as d2l
import warnings
warnings.filterwarnings("ignore")
import math
import paddle
from paddle import nn

10.5.2 Реализация

На практике при реализации в качестве каждой головы внимания обычно выбирают масштабированное скалярное произведение внимания. Чтобы избежать резкого роста вычислительных затрат и затрат на параметры, положим $p_{q} = p_{k} = p_{v} = p_{o} / h$ . Заметим, что если задать число выходов линейных преобразований запросов, ключей и значений равным $p_{q} h = p_{k} h = p_{v} h = p_{o}$ , то $h$ голов можно вычислять параллельно. В реализации ниже $p_{o}$ задаётся параметром num_hiddens.

MXNet:

#@save
class MultiHeadAttention(nn.Block):
    """Многоголовое внимание"""
    def __init__(self, num_hiddens, num_heads, dropout, use_bias=False,
                 **kwargs):
        super(MultiHeadAttention, self).__init__(**kwargs)
        self.num_heads = num_heads
        self.attention = d2l.DotProductAttention(dropout)
        self.W_q = nn.Dense(num_hiddens, use_bias=use_bias, flatten=False)
        self.W_k = nn.Dense(num_hiddens, use_bias=use_bias, flatten=False)
        self.W_v = nn.Dense(num_hiddens, use_bias=use_bias, flatten=False)
        self.W_o = nn.Dense(num_hiddens, use_bias=use_bias, flatten=False)

    def forward(self, queries, keys, values, valid_lens):
        # Формы queries, keys, values:
        # (batch_size, число запросов или пар «ключ — значение», num_hiddens)
        # Форма valid_lens:
        # (batch_size,) или (batch_size, число запросов)
        # После преобразования формы queries, keys, values:
        # (batch_size*num_heads, число запросов или пар «ключ — значение»,
        # num_hiddens/num_heads)
        queries = transpose_qkv(self.W_q(queries), self.num_heads)
        keys = transpose_qkv(self.W_k(keys), self.num_heads)
        values = transpose_qkv(self.W_v(values), self.num_heads)

        if valid_lens is not None:
            # По оси 0 первый элемент (скаляр или вектор) повторяем num_heads раз,
            # затем так же повторяем второй элемент и так далее.
            valid_lens = valid_lens.repeat(self.num_heads, axis=0)

        # Форма output: (batch_size*num_heads, число запросов,
        # num_hiddens/num_heads)
        output = self.attention(queries, keys, values, valid_lens)
        
        # Форма output_concat: (batch_size, число запросов, num_hiddens)
        output_concat = transpose_output(output, self.num_heads)
        return self.W_o(output_concat)

PyTorch:

#@save
class MultiHeadAttention(nn.Module):
    """Многоголовое внимание"""
    def __init__(self, key_size, query_size, value_size, num_hiddens,
                 num_heads, dropout, bias=False, **kwargs):
        super(MultiHeadAttention, self).__init__(**kwargs)
        self.num_heads = num_heads
        self.attention = d2l.DotProductAttention(dropout)
        self.W_q = nn.Linear(query_size, num_hiddens, bias=bias)
        self.W_k = nn.Linear(key_size, num_hiddens, bias=bias)
        self.W_v = nn.Linear(value_size, num_hiddens, bias=bias)
        self.W_o = nn.Linear(num_hiddens, num_hiddens, bias=bias)

    def forward(self, queries, keys, values, valid_lens):
        # Формы queries, keys, values:
        # (batch_size, число запросов или пар «ключ — значение», num_hiddens)
        # Форма valid_lens:
        # (batch_size,) или (batch_size, число запросов)
        # После преобразования формы queries, keys, values:
        # (batch_size*num_heads, число запросов или пар «ключ — значение»,
        # num_hiddens/num_heads)
        queries = transpose_qkv(self.W_q(queries), self.num_heads)
        keys = transpose_qkv(self.W_k(keys), self.num_heads)
        values = transpose_qkv(self.W_v(values), self.num_heads)

        if valid_lens is not None:
            # По оси 0 первый элемент (скаляр или вектор) повторяем num_heads раз,
            # затем так же повторяем второй элемент и так далее.
            valid_lens = torch.repeat_interleave(
                valid_lens, repeats=self.num_heads, dim=0)

        # Форма output: (batch_size*num_heads, число запросов,
        # num_hiddens/num_heads)
        output = self.attention(queries, keys, values, valid_lens)

        # Форма output_concat: (batch_size, число запросов, num_hiddens)
        output_concat = transpose_output(output, self.num_heads)
        return self.W_o(output_concat)

TensorFlow:

#@save
class MultiHeadAttention(tf.keras.layers.Layer):
    """Многоголовое внимание"""
    def __init__(self, key_size, query_size, value_size, num_hiddens,
                 num_heads, dropout, bias=False, **kwargs):
        super().__init__(**kwargs)
        self.num_heads = num_heads
        self.attention = d2l.DotProductAttention(dropout)
        self.W_q = tf.keras.layers.Dense(num_hiddens, use_bias=bias)
        self.W_k = tf.keras.layers.Dense(num_hiddens, use_bias=bias)
        self.W_v = tf.keras.layers.Dense(num_hiddens, use_bias=bias)
        self.W_o = tf.keras.layers.Dense(num_hiddens, use_bias=bias)
    
    def call(self, queries, keys, values, valid_lens, **kwargs):
        # Формы queries, keys, values:
        # (batch_size, число запросов или пар «ключ — значение», num_hiddens)
        # Форма valid_lens:
        # (batch_size,) или (batch_size, число запросов)
        # После преобразования формы queries, keys, values:
        # (batch_size*num_heads, число запросов или пар «ключ — значение»,
        # num_hiddens/num_heads)
        queries = transpose_qkv(self.W_q(queries), self.num_heads)
        keys = transpose_qkv(self.W_k(keys), self.num_heads)
        values = transpose_qkv(self.W_v(values), self.num_heads)
        
        if valid_lens is not None:
            # По оси 0 первый элемент (скаляр или вектор) повторяем num_heads раз,
            # затем так же повторяем второй элемент и так далее.
            valid_lens = tf.repeat(valid_lens, repeats=self.num_heads, axis=0)
            
        # Форма output: (batch_size*num_heads, число запросов,
        # num_hiddens/num_heads)
        output = self.attention(queries, keys, values, valid_lens, **kwargs)
        
        # Форма output_concat: (batch_size, число запросов, num_hiddens)
        output_concat = transpose_output(output, self.num_heads)
        return self.W_o(output_concat)

PaddlePaddle:

#@save
class MultiHeadAttention(nn.Layer):
    """Многоголовое внимание"""
    def __init__(self, key_size, query_size, value_size, num_hiddens,
                 num_heads, dropout, bias=False, **kwargs):
        super(MultiHeadAttention, self).__init__(**kwargs)
        self.num_heads = num_heads
        self.attention = d2l.DotProductAttention(dropout)
        self.W_q = nn.Linear(query_size, num_hiddens, bias_attr=bias)
        self.W_k = nn.Linear(key_size, num_hiddens, bias_attr=bias)
        self.W_v = nn.Linear(value_size, num_hiddens, bias_attr=bias)
        self.W_o = nn.Linear(num_hiddens, num_hiddens, bias_attr=bias)

    def forward(self, queries, keys, values, valid_lens):
        # Формы queries, keys, values:
        # (batch_size, число запросов или пар «ключ — значение», num_hiddens)
        # Форма valid_lens:
        # (batch_size,) или (batch_size, число запросов)
        # После преобразования формы queries, keys, values:
        # (batch_size*num_heads, число запросов или пар «ключ — значение»,
        # num_hiddens/num_heads)
        queries = transpose_qkv(self.W_q(queries), self.num_heads)
        keys = transpose_qkv(self.W_k(keys), self.num_heads)
        values = transpose_qkv(self.W_v(values), self.num_heads)
        if valid_lens is not None:
            # По оси 0 первый элемент (скаляр или вектор) повторяем num_heads раз,
            # затем так же повторяем второй элемент и так далее.
            valid_lens = paddle.repeat_interleave(
                valid_lens, repeats=self.num_heads, axis=0)

        # Форма output: (batch_size*num_heads, число запросов,
        # num_hiddens/num_heads)
        output = self.attention(queries, keys, values, valid_lens)

        # Форма output_concat: (batch_size, число запросов, num_hiddens)
        output_concat = transpose_output(output, self.num_heads)
        return self.W_o(output_concat)

Чтобы все головы вычислялись параллельно, класс MultiHeadAttention выше использует две функции транспонирования, определённые ниже. Функция transpose_output обращает операцию функции transpose_qkv.

MXNet:

#@save
def transpose_qkv(X, num_heads):
    """Меняет форму для параллельного вычисления нескольких голов внимания"""
    # Форма входа X: (batch_size, число запросов или пар «ключ — значение», num_hiddens)
    # Форма выхода X: (batch_size, число запросов или пар «ключ — значение», num_heads,
    # num_hiddens/num_heads)
    X = X.reshape(X.shape[0], X.shape[1], num_heads, -1)

    # Форма выхода X: (batch_size, num_heads, число запросов или пар «ключ — значение»,
    # num_hiddens/num_heads)
    X = X.transpose(0, 2, 1, 3)

    # Форма итогового выхода: (batch_size*num_heads, число запросов или пар «ключ — значение»,
    # num_hiddens/num_heads)
    return X.reshape(-1, X.shape[2], X.shape[3])


#@save
def transpose_output(X, num_heads):
    """Обращает операцию функции transpose_qkv"""
    X = X.reshape(-1, num_heads, X.shape[1], X.shape[2])
    X = X.transpose(0, 2, 1, 3)
    return X.reshape(X.shape[0], X.shape[1], -1)

PyTorch:

#@save
def transpose_qkv(X, num_heads):
    """Меняет форму для параллельного вычисления нескольких голов внимания"""
    # Форма входа X: (batch_size, число запросов или пар «ключ — значение», num_hiddens)
    # Форма выхода X: (batch_size, число запросов или пар «ключ — значение», num_heads,
    # num_hiddens/num_heads)
    X = X.reshape(X.shape[0], X.shape[1], num_heads, -1)

    # Форма выхода X: (batch_size, num_heads, число запросов или пар «ключ — значение»,
    # num_hiddens/num_heads)
    X = X.permute(0, 2, 1, 3)

    # Форма итогового выхода: (batch_size*num_heads, число запросов или пар «ключ — значение»,
    # num_hiddens/num_heads)
    return X.reshape(-1, X.shape[2], X.shape[3])


#@save
def transpose_output(X, num_heads):
    """Обращает операцию функции transpose_qkv"""
    X = X.reshape(-1, num_heads, X.shape[1], X.shape[2])
    X = X.permute(0, 2, 1, 3)
    return X.reshape(X.shape[0], X.shape[1], -1)

TensorFlow:

#@save
def transpose_qkv(X, num_heads):
    """Меняет форму для параллельного вычисления нескольких голов внимания"""
    # Форма входа X: (batch_size, число запросов или пар «ключ — значение», num_hiddens)
    # Форма выхода X: (batch_size, число запросов или пар «ключ — значение», num_heads,
    # num_hiddens/num_heads)
    X = tf.reshape(X, shape=(X.shape[0], X.shape[1], num_heads, -1))

    # Форма выхода X: (batch_size, num_heads, число запросов или пар «ключ — значение»,
    # num_hiddens/num_heads)
    X = tf.transpose(X, perm=(0, 2, 1, 3))

    # Форма итогового выхода: (batch_size*num_heads, число запросов или пар «ключ — значение»,
    # num_hiddens/num_heads)
    return tf.reshape(X, shape=(-1, X.shape[2], X.shape[3]))


#@save
def transpose_output(X, num_heads):
    """Обращает операцию функции transpose_qkv"""
    X = tf.reshape(X, shape=(-1, num_heads, X.shape[1], X.shape[2]))
    X = tf.transpose(X, perm=(0, 2, 1, 3))
    return tf.reshape(X, shape=(X.shape[0], X.shape[1], -1))

PaddlePaddle:

#@save
def transpose_qkv(X, num_heads):
    """Меняет форму для параллельного вычисления нескольких голов внимания"""
    # Форма входа X: (batch_size, число запросов или пар «ключ — значение», num_hiddens)
    # Форма выхода X: (batch_size, число запросов или пар «ключ — значение», num_heads,
    # num_hiddens/num_heads)
    X = X.reshape((X.shape[0], X.shape[1], num_heads, -1))

    # Форма выхода X: (batch_size, num_heads, число запросов или пар «ключ — значение»,
    # num_hiddens/num_heads)
    X = X.transpose((0, 2, 1, 3))

    # Форма итогового выхода: (batch_size*num_heads, число запросов или пар «ключ — значение»,
    # num_hiddens/num_heads)
    return X.reshape((-1, X.shape[2], X.shape[3]))


#@save
def transpose_output(X, num_heads):
    """Обращает операцию функции transpose_qkv"""
    X = X.reshape((-1, num_heads, X.shape[1], X.shape[2]))
    X = X.transpose((0, 2, 1, 3))
    return X.reshape((X.shape[0], X.shape[1], -1))

Проверим класс MultiHeadAttention на небольшом примере, в котором ключи и значения совпадают. Форма выхода многоголового внимания — (batch_size, num_queries, num_hiddens).

MXNet:

num_hiddens, num_heads = 100, 5
attention = MultiHeadAttention(num_hiddens, num_heads, 0.5)
attention.initialize()

PyTorch:

num_hiddens, num_heads = 100, 5
attention = MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens,
                               num_hiddens, num_heads, 0.5)
attention.eval()

TensorFlow:

num_hiddens, num_heads = 100, 5
attention = MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens,
                               num_hiddens, num_heads, 0.5)

PaddlePaddle:

num_hiddens, num_heads = 100, 5
attention = MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens,
                               num_hiddens, num_heads, 0.5)
attention.eval()

MXNet, PyTorch, PaddlePaddle:

batch_size, num_queries = 2, 4
num_kvpairs, valid_lens =  6, d2l.tensor([3, 2])
X = d2l.ones((batch_size, num_queries, num_hiddens))
Y = d2l.ones((batch_size, num_kvpairs, num_hiddens))
attention(X, Y, Y, valid_lens).shape

TensorFlow:

batch_size, num_queries = 2, 4
num_kvpairs, valid_lens = 6, d2l.tensor([3, 2])
X = tf.ones((batch_size, num_queries, num_hiddens))
Y = tf.ones((batch_size, num_kvpairs, num_hiddens))
attention(X, Y, Y, valid_lens, training=False).shape

Резюме

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

Упражнения

Визуализируйте веса внимания для каждой головы в этом эксперименте.
Предположим, что есть обученная модель на основе многоголового внимания, и мы хотим удалить наименее важные головы внимания, чтобы ускорить предсказание. Как спланировать эксперимент, чтобы измерить важность голов внимания?

MXNet: Обсуждение

PyTorch: Обсуждение

PaddlePaddle: Обсуждение

Литература

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., & Polosukhin, I. (2017). Attention is all you need. Advances in Neural Information Processing Systems, 5998–6008.

Keyboard shortcuts

Погружение в глубокое обучение

10.5 Многоголовое внимание

10.5.1 Модель

10.5.2 Реализация

Резюме

Упражнения

Литература