Функция scaled_dot_product_attention

Функция scaled_dot_product_attention из модуля torch.nn.functional вычисляет внимание на основе скалярного произведения запроса и ключа с последующим масштабированием. Она принимает тензоры запроса q, ключа k и значения v, а также необязательную маску attn_mask и коэффициент dropout. Функция оптимизирована для работы с графическими процессорами и поддерживает эффективные реализации для различных архитектур.

Синтаксис

torch.nn.functional.scaled_dot_product_attention(
	q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None
)

Пример

Давайте вычислим внимание для простых тензоров запроса, ключа и значения:

import torch
import torch.nn.functional as F

torch.manual_seed(0)

q = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
k = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
v = torch.tensor([[5.0, 6.0], [7.0, 8.0]])

res = F.scaled_dot_product_attention(q, k, v)
print(res)

Результат выполнения кода:

tensor([
	[5.5906, 6.5906],
	[6.3072, 7.3072],
])

Пример

Рассмотрим использование маски для скрытия определённых позиций:

import torch
import torch.nn.functional as F

torch.manual_seed(0)

q = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
k = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
v = torch.tensor([[5.0, 6.0], [7.0, 8.0]])
mask = torch.tensor([[0.0, -float('inf')], [0.0, 0.0]])

res = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
print(res)

Результат выполнения кода:

tensor([
	[5.1290, 6.1290],
	[6.3072, 7.3072],
])

Пример

Применим каузальную маску для авторегрессивного моделирования:

import torch
import torch.nn.functional as F

torch.manual_seed(0)

q = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
k = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
v = torch.tensor([[5.0, 6.0], [7.0, 8.0]])

res = F.scaled_dot_product_attention(q, k, v, is_causal=True)
print(res)

Результат выполнения кода:

tensor([
	[5.5906, 6.5906],
	[6.3072, 7.3072],
])

Пример

Изменим коэффициент масштабирования для управления значениями внимания:

import torch
import torch.nn.functional as F

torch.manual_seed(0)

q = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
k = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
v = torch.tensor([[5.0, 6.0], [7.0, 8.0]])

res = F.scaled_dot_product_attention(q, k, v, scale=0.5)
print(res)

Результат выполнения кода:

tensor([
	[5.5167, 6.5167],
	[6.6529, 7.6529],
])

Пример

Добавим dropout для регуляризации модели:

import torch
import torch.nn.functional as F

torch.manual_seed(0)

q = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
k = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
v = torch.tensor([[5.0, 6.0], [7.0, 8.0]])

res = F.scaled_dot_product_attention(q, k, v, dropout_p=0.5)
print(res)

Результат выполнения кода:

tensor([
	[ 0.0000,  0.0000],
	[12.6143, 14.6143],
])

Смотрите также

функцию softmax,
которая применяется для нормализации весов внимания
функцию cosine_similarity,
которая вычисляет косинусное сходство между векторами
функцию dropout,
которая применяет регуляризацию путем обнуления элементов
функцию multi_head_attention_forward,
которая реализует многоголовое внимание