File size: 2,447 Bytes
5d3bddc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | from __future__ import annotations
import math
import torch
from torch import nn
from torch.nn import functional as F
VOCAB_SIZE = 64
class ContentAddressedMemory(nn.Module):
"""A learned key-value tape with differentiable content addressing."""
def __init__(self, width: int = 24) -> None:
super().__init__()
self.key_embedding = nn.Embedding(VOCAB_SIZE, width)
self.value_embedding = nn.Embedding(VOCAB_SIZE, width)
self.output = nn.Linear(width, VOCAB_SIZE)
self.log_beta = nn.Parameter(torch.tensor(math.log(10.0)))
def forward(
self,
keys: torch.Tensor,
values: torch.Tensor,
query: torch.Tensor,
*,
return_attention: bool = False,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
memory_keys = F.normalize(self.key_embedding(keys), dim=-1)
query_key = F.normalize(self.key_embedding(query), dim=-1)
beta = self.log_beta.exp().clamp(1.0, 30.0)
scores = torch.einsum("bsd,bd->bs", memory_keys, query_key) * beta
attention = scores.softmax(dim=-1)
read = torch.einsum(
"bs,bsd->bd",
attention,
self.value_embedding(values),
)
logits = self.output(read)
if return_attention:
return logits, attention
return logits
class FixedStateGRU(nn.Module):
"""A larger recurrent control that compresses the tape into one state."""
def __init__(self, embedding_dim: int = 8, hidden_dim: int = 24) -> None:
super().__init__()
self.embedding = nn.Embedding(VOCAB_SIZE * 3, embedding_dim)
self.gru = nn.GRU(embedding_dim, hidden_dim, batch_first=True)
self.output = nn.Linear(hidden_dim, VOCAB_SIZE)
def forward(
self,
keys: torch.Tensor,
values: torch.Tensor,
query: torch.Tensor,
) -> torch.Tensor:
batch, slots = keys.shape
tape = torch.empty(
batch,
slots * 2 + 1,
dtype=torch.long,
device=keys.device,
)
tape[:, 0 : slots * 2 : 2] = keys
tape[:, 1 : slots * 2 : 2] = values + VOCAB_SIZE
tape[:, -1] = query + VOCAB_SIZE * 2
_, state = self.gru(self.embedding(tape))
return self.output(state[-1])
def parameter_count(model: nn.Module) -> int:
return sum(parameter.numel() for parameter in model.parameters())
|