| """Causal Transformer for the selective copy task.""" |
|
|
| import torch |
| import torch.nn as nn |
|
|
|
|
| class CausalSelfAttention(nn.Module): |
| def __init__(self, d_model, n_heads, dropout=0.0): |
| super().__init__() |
| assert d_model % n_heads == 0 |
| self.n_heads = n_heads |
| self.d_head = d_model // n_heads |
| self.qkv = nn.Linear(d_model, 3 * d_model, bias=False) |
| self.proj = nn.Linear(d_model, d_model, bias=False) |
| self.attn_drop_p = dropout |
|
|
| def forward(self, x): |
| B, T, C = x.shape |
| q, k, v = self.qkv(x).split(C, dim=-1) |
| def split_heads(t): |
| return t.view(B, T, self.n_heads, self.d_head).transpose(1, 2) |
| out = torch.nn.functional.scaled_dot_product_attention( |
| split_heads(q), split_heads(k), split_heads(v), |
| is_causal=True, |
| dropout_p=self.attn_drop_p if self.training else 0.0, |
| ) |
| return self.proj(out.transpose(1, 2).contiguous().view(B, T, C)) |
|
|
|
|
| class TransformerBlock(nn.Module): |
| def __init__(self, d_model, n_heads, dropout=0.0): |
| super().__init__() |
| self.ln1 = nn.LayerNorm(d_model) |
| self.attn = CausalSelfAttention(d_model, n_heads, dropout) |
| self.ln2 = nn.LayerNorm(d_model) |
| self.mlp = nn.Sequential( |
| nn.Linear(d_model, 4 * d_model), nn.GELU(), |
| nn.Linear(4 * d_model, d_model), nn.Dropout(dropout), |
| ) |
|
|
| def forward(self, x): |
| x = x + self.attn(self.ln1(x)) |
| x = x + self.mlp(self.ln2(x)) |
| return x |
|
|
|
|
| class TransformerModel(nn.Module): |
| def __init__(self, d_input, d_model, d_output, n_layers=2, n_heads=4, max_len=4096, dropout=0.0, **kwargs): |
| super().__init__() |
| self.input_proj = nn.Linear(d_input, d_model) |
| self.pos_emb = nn.Embedding(max_len, d_model) |
| self.blocks = nn.ModuleList([TransformerBlock(d_model, n_heads, dropout) for _ in range(n_layers)]) |
| self.ln_f = nn.LayerNorm(d_model) |
| self.head = nn.Linear(d_model, d_output) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| """(B, T, d_input) → (B, T, d_output)""" |
| B, T, _ = x.shape |
| if T > self.pos_emb.num_embeddings: |
| raise ValueError(f"Sequence length {T} exceeds max_len {self.pos_emb.num_embeddings}") |
| h = self.input_proj(x) + self.pos_emb(torch.arange(T, device=x.device)) |
| for block in self.blocks: |
| h = block(h) |
| return self.head(self.ln_f(h)) |
|
|
| @staticmethod |
| def extra_kwargs(model_cfg) -> dict: |
| return {"n_heads": model_cfg.n_heads} |
|
|
