"""A minimal Transformer encoder implemented with basic PyTorch modules. The script includes shape and masking checks plus a tiny synthetic task. It is intended for learning and reproduction, not as a high-performance kernel. """ from __future__ import annotations import argparse import math import torch from torch import Tensor, nn def set_seed(seed: int) -> None: torch.manual_seed(seed) class SinusoidalPositionalEncoding(nn.Module): def __init__(self, d_model: int, max_length: int = 512) -> None: super().__init__() position = torch.arange(max_length, dtype=torch.float32).unsqueeze(1) div_term = torch.exp( torch.arange(0, d_model, 2, dtype=torch.float32) * (-math.log(10000.0) / d_model) ) encoding = torch.zeros(max_length, d_model) encoding[:, 0::2] = torch.sin(position * div_term) encoding[:, 1::2] = torch.cos(position * div_term) self.register_buffer("encoding", encoding.unsqueeze(0), persistent=False) def forward(self, x: Tensor) -> Tensor: return x + self.encoding[:, : x.size(1)].to(dtype=x.dtype) class MultiHeadSelfAttention(nn.Module): def __init__(self, d_model: int, n_heads: int, dropout: float = 0.0) -> None: super().__init__() if d_model % n_heads != 0: raise ValueError("d_model must be divisible by n_heads") self.n_heads = n_heads self.d_head = d_model // n_heads self.qkv = nn.Linear(d_model, 3 * d_model) self.out_proj = nn.Linear(d_model, d_model) self.dropout = nn.Dropout(dropout) def forward( self, x: Tensor, key_padding_mask: Tensor | None = None, need_weights: bool = False, ) -> tuple[Tensor, Tensor | None]: batch, length, d_model = x.shape qkv = self.qkv(x).view( batch, length, 3, self.n_heads, self.d_head ) q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0) scores = q @ k.transpose(-2, -1) scores = scores / math.sqrt(self.d_head) if key_padding_mask is not None: if key_padding_mask.shape != (batch, length): raise ValueError("key_padding_mask must have shape [batch, length]") scores = scores.masked_fill( key_padding_mask[:, None, None, :], torch.finfo(scores.dtype).min, ) weights = torch.softmax(scores, dim=-1) context = self.dropout(weights) @ v context = context.transpose(1, 2).contiguous().view(batch, length, d_model) output = self.out_proj(context) return output, weights if need_weights else None class EncoderLayer(nn.Module): def __init__( self, d_model: int, n_heads: int, d_ff: int, dropout: float = 0.1, norm_first: bool = False, ) -> None: super().__init__() self.norm_first = norm_first self.self_attn = MultiHeadSelfAttention(d_model, n_heads, dropout) self.ffn = nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Dropout(dropout), nn.Linear(d_ff, d_model), ) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout) def forward( self, x: Tensor, key_padding_mask: Tensor | None = None, need_weights: bool = False, ) -> tuple[Tensor, Tensor | None]: if self.norm_first: attn_out, weights = self.self_attn( self.norm1(x), key_padding_mask, need_weights ) x = x + self.dropout1(attn_out) x = x + self.dropout2(self.ffn(self.norm2(x))) return x, weights attn_out, weights = self.self_attn(x, key_padding_mask, need_weights) x = self.norm1(x + self.dropout1(attn_out)) x = self.norm2(x + self.dropout2(self.ffn(x))) return x, weights class TransformerEncoder(nn.Module): def __init__( self, vocab_size: int, d_model: int = 64, n_heads: int = 4, d_ff: int = 128, n_layers: int = 2, dropout: float = 0.1, norm_first: bool = False, ) -> None: super().__init__() self.d_model = d_model self.embedding = nn.Embedding(vocab_size, d_model, padding_idx=0) self.position = SinusoidalPositionalEncoding(d_model) self.layers = nn.ModuleList( [ EncoderLayer(d_model, n_heads, d_ff, dropout, norm_first) for _ in range(n_layers) ] ) self.final_norm = nn.LayerNorm(d_model) if norm_first else nn.Identity() self.reset_parameters() def reset_parameters(self) -> None: for parameter in self.parameters(): if parameter.dim() > 1: nn.init.xavier_uniform_(parameter) with torch.no_grad(): self.embedding.weight[0].zero_() def forward( self, token_ids: Tensor, key_padding_mask: Tensor | None = None, need_weights: bool = False, ) -> tuple[Tensor, list[Tensor]]: x = self.embedding(token_ids) * math.sqrt(self.d_model) x = self.position(x) all_weights: list[Tensor] = [] for layer in self.layers: x, weights = layer(x, key_padding_mask, need_weights) if weights is not None: all_weights.append(weights) return self.final_norm(x), all_weights class FirstLastMatchClassifier(nn.Module): def __init__(self, vocab_size: int) -> None: super().__init__() self.encoder = TransformerEncoder(vocab_size=vocab_size) self.classifier = nn.Linear(64, 2) def forward( self, token_ids: Tensor, key_padding_mask: Tensor | None = None ) -> Tensor: encoded, _ = self.encoder(token_ids, key_padding_mask) return self.classifier(encoded[:, 0]) def make_batch( batch_size: int, sequence_length: int, vocab_size: int ) -> tuple[Tensor, Tensor]: token_ids = torch.randint(1, vocab_size, (batch_size, sequence_length)) labels = torch.randint(0, 2, (batch_size,)) token_ids[labels == 1, -1] = token_ids[labels == 1, 0] negative_rows = labels == 0 negative_last = torch.randint(1, vocab_size, (int(negative_rows.sum()),)) negative_first = token_ids[negative_rows, 0] negative_last = torch.where( negative_last == negative_first, (negative_last % (vocab_size - 1)) + 1, negative_last, ) token_ids[negative_rows, -1] = negative_last return token_ids, labels def run_invariant_checks() -> None: encoder = TransformerEncoder(vocab_size=32, dropout=0.0) encoder.eval() token_ids = torch.tensor([[1, 2, 3, 4, 0], [5, 6, 7, 0, 0]]) padding_mask = token_ids.eq(0) encoded, weights = encoder(token_ids, padding_mask, need_weights=True) assert encoded.shape == (2, 5, 64) assert len(weights) == 2 assert weights[0].shape == (2, 4, 5, 5) torch.testing.assert_close( weights[0].sum(dim=-1), torch.ones_like(weights[0].sum(dim=-1)) ) masked_probability = weights[0].masked_select( padding_mask[:, None, None, :].expand_as(weights[0]) ) assert torch.count_nonzero(masked_probability).item() == 0 print("invariant checks: passed") def train_toy_task(steps: int, seed: int) -> None: set_seed(seed) vocab_size = 64 model = FirstLastMatchClassifier(vocab_size) optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4) loss_fn = nn.CrossEntropyLoss() model.train() for step in range(1, steps + 1): token_ids, labels = make_batch(64, 12, vocab_size) loss = loss_fn(model(token_ids), labels) optimizer.zero_grad(set_to_none=True) loss.backward() optimizer.step() if step == 1 or step % 20 == 0 or step == steps: print(f"step={step:03d} loss={loss.item():.4f}") def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--steps", type=int, default=80) parser.add_argument("--seed", type=int, default=7) parser.add_argument( "--check-only", action="store_true", help="run shape and masking checks without training", ) return parser.parse_args() def main() -> None: args = parse_args() print(f"torch={torch.__version__}") run_invariant_checks() if not args.check_only: train_toy_task(args.steps, args.seed) if __name__ == "__main__": main()