| 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)) |
|
|
