File size: 2,865 Bytes
45507bb | 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 torch
import torch.nn.functional as F
from torch import nn
class Attention(nn.Module):
"""Multi-head (self or cross) attention on top of SDPA; `key_mask` is bool (B, S), True = valid."""
def __init__(self, d: int, heads: int, dropout: float = 0.0):
super().__init__()
assert d % heads == 0, "d_model must be divisible by n_heads"
self.h, self.dh = heads, d // heads
self.q = nn.Linear(d, d)
self.kv = nn.Linear(d, 2 * d)
self.o = nn.Linear(d, d)
self.dropout = dropout
def forward(self, x, ctx=None, key_mask=None):
ctx = x if ctx is None else ctx
B, T, D = x.shape
S = ctx.shape[1]
q = self.q(x).view(B, T, self.h, self.dh).transpose(1, 2)
kv = self.kv(ctx).view(B, S, 2, self.h, self.dh).permute(2, 0, 3, 1, 4)
mask = key_mask[:, None, None, :] if key_mask is not None else None
out = F.scaled_dot_product_attention(q, kv[0], kv[1], attn_mask=mask,
dropout_p=self.dropout if self.training else 0.0)
return self.o(out.transpose(1, 2).reshape(B, T, D))
class MLP(nn.Module):
def __init__(self, d: int, ratio: float = 4.0, dropout: float = 0.0):
super().__init__()
hidden = int(d * ratio)
self.net = nn.Sequential(nn.Linear(d, hidden), nn.GELU(), nn.Linear(hidden, d), nn.Dropout(dropout))
def forward(self, x):
return self.net(x)
class EncoderBlock(nn.Module):
"""Pre-LN transformer block."""
def __init__(self, d, heads, mlp_ratio=4.0, dropout=0.0):
super().__init__()
self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
self.attn = Attention(d, heads, dropout)
self.mlp = MLP(d, mlp_ratio, dropout)
def forward(self, x, key_mask=None):
h = self.ln1(x)
x = x + self.attn(h, key_mask=key_mask)
return x + self.mlp(self.ln2(x))
class DecisionBlock(nn.Module):
"""Choices read the context (cross-attn), then compare with each other (self-attn).
No positional information over choices -> the block is permutation-equivariant:
reordering the choices reorders the scores, nothing else.
"""
def __init__(self, d, heads, mlp_ratio=4.0, dropout=0.0):
super().__init__()
self.ln_q, self.ln_ctx = nn.LayerNorm(d), nn.LayerNorm(d)
self.cross = Attention(d, heads, dropout)
self.ln_s = nn.LayerNorm(d)
self.self_attn = Attention(d, heads, dropout)
self.ln_m = nn.LayerNorm(d)
self.mlp = MLP(d, mlp_ratio, dropout)
def forward(self, c, ctx, ctx_mask, choice_mask):
c = c + self.cross(self.ln_q(c), self.ln_ctx(ctx), key_mask=ctx_mask)
c = c + self.self_attn(self.ln_s(c), key_mask=choice_mask)
return c + self.mlp(self.ln_m(c))
|