Матричная ассоциативная память (линейное внимание) и векторная ассоциативная память

В линейном внимании в матрицу памяти S пара ключ-значение добавляется как внешнее произведение их векторов.

S = S + kᵀv

В векторную память пара ключ-значение добавляется как поэлементное произведение.

S = S + k⊙v

Чтение из линейного внимания – это матричное умножение ключа и состояния.

o = kS

Чтение из векторной памяти – это поэлементное произведение ключа и состояния.

o = k⊙S

import torch
import torch.nn.functional as F

D = 1024

# Векторы ключа и значения
k, v = F.normalize(torch.randn(2, D))

# Запись в матрицу памяти линейного внимания
S = torch.outer(k, v)

# Чтение из памяти по ключу
o = k @ S

# Смотрим что прочиталось
r = torch.dot(o, v)
print(r)
# 1.0

# Запись в векторную память пары ключ-значение как поэлементное произведение
S = k * v
S = F.normalize(S.unsqueeze(0)).squeeze(0)

# Читаем из векторной памяти по ключу
o = S * k
o = F.normalize(o.unsqueeze(0)).squeeze(0)

# Смотрим что прочиталось
r = torch.dot(o, v)
print(r)
# 0.58