#!/usr/bin/env python3 """Train an LSTM to predict the next item in controlled number sequences.""" from __future__ import annotations import argparse import math import torch from torch import nn from torch.nn import functional as F torch.set_num_threads(1) VOCABULARY_SIZE = 10 NUMBER_COUNT = 8 def make_sequence_data( count: int, sequence_length: int, seed: int, ) -> tuple[torch.Tensor, torch.Tensor]: """Create sequences whose first token selects counting up or down. Control ID 0 means count up and ID 1 means count down. Number values 0..7 are stored as token IDs 2..9. Inputs and shifted targets both have shape ``(count, sequence_length)``. """ if count < 1 or sequence_length < 3: raise ValueError("count must be positive and sequence_length at least three") generator = torch.Generator(device="cpu").manual_seed(seed) controls = torch.randint(0, 2, (count,), generator=generator) starts = torch.randint(0, NUMBER_COUNT, (count,), generator=generator) complete = torch.empty((count, sequence_length + 1), dtype=torch.long) complete[:, 0] = controls complete[:, 1] = starts + 2 current = starts for position in range(2, sequence_length + 1): step = torch.where(controls == 0, 1, -1) current = (current + step) % NUMBER_COUNT complete[:, position] = current + 2 return complete[:, :-1], complete[:, 1:] def manual_lstm_cell( x: torch.Tensor, previous_hidden: torch.Tensor, previous_cell: torch.Tensor, input_weight: torch.Tensor, hidden_weight: torch.Tensor, bias: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]: """Calculate one LSTM state update from the four gate equations.""" gates = x @ input_weight.T + previous_hidden @ hidden_weight.T + bias input_raw, forget_raw, candidate_raw, output_raw = gates.chunk(4, dim=-1) input_gate = torch.sigmoid(input_raw) forget_gate = torch.sigmoid(forget_raw) candidate = torch.tanh(candidate_raw) output_gate = torch.sigmoid(output_raw) cell = forget_gate * previous_cell + input_gate * candidate hidden = output_gate * torch.tanh(cell) return hidden, cell, { "input": input_gate, "forget": forget_gate, "candidate": candidate, "output": output_gate, } class NextStepLSTM(nn.Module): """Embedding, LSTM, and linear output layer for next-token prediction.""" def __init__(self, embedding_size: int = 8, hidden_size: int = 24) -> None: super().__init__() self.embedding = nn.Embedding(VOCABULARY_SIZE, embedding_size) self.lstm = nn.LSTM(embedding_size, hidden_size, batch_first=True) self.output = nn.Linear(hidden_size, VOCABULARY_SIZE) def forward(self, token_ids: torch.Tensor) -> torch.Tensor: vectors = self.embedding(token_ids) states, _ = self.lstm(vectors) return self.output(states) def memoryless_baseline_loss(sequence_length: int) -> float: """Loss of a rule that knows counting but cannot remember the control token.""" first_position_loss = math.log(NUMBER_COUNT) later_position_loss = math.log(2.0) return (first_position_loss + (sequence_length - 1) * later_position_loss) / sequence_length def evaluate( model: NextStepLSTM, inputs: torch.Tensor, targets: torch.Tensor, ) -> tuple[float, float]: """Return cross-entropy and token accuracy without changing the model.""" model.eval() with torch.no_grad(): logits = model(inputs) loss = F.cross_entropy(logits.reshape(-1, VOCABULARY_SIZE), targets.reshape(-1)) accuracy = (logits.argmax(dim=-1) == targets).float().mean() return loss.item(), accuracy.item() def train_model(steps: int = 300, sequence_length: int = 10) -> NextStepLSTM: """Train on deterministic mini-batches from a generated CPU dataset.""" if steps < 1: raise ValueError("steps must be positive") torch.manual_seed(2026) train_inputs, train_targets = make_sequence_data(512, sequence_length, seed=30) batch_generator = torch.Generator(device="cpu").manual_seed(31) model = NextStepLSTM() optimizer = torch.optim.Adam(model.parameters(), lr=0.015) model.train() for _ in range(steps): indices = torch.randint(0, len(train_inputs), (64,), generator=batch_generator) inputs = train_inputs[indices] targets = train_targets[indices] optimizer.zero_grad() logits = model(inputs) loss = F.cross_entropy(logits.reshape(-1, VOCABULARY_SIZE), targets.reshape(-1)) loss.backward() nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() return model def run(steps: int = 300, smoke_test: bool = False) -> dict[str, float]: """Train, compare with a memoryless baseline, and evaluate held-out data.""" sequence_length = 10 effective_steps = min(steps, 3) if smoke_test else steps validation_inputs, validation_targets = make_sequence_data( 256, sequence_length, seed=40 ) model = train_model(effective_steps, sequence_length) validation_loss, validation_accuracy = evaluate( model, validation_inputs, validation_targets ) baseline_loss = memoryless_baseline_loss(sequence_length) print("input and target shapes:", tuple(validation_inputs.shape), tuple(validation_targets.shape)) print("logit shape:", tuple(model(validation_inputs[:4]).shape)) print(f"memoryless baseline loss: {baseline_loss:.4f}") print(f"held-out loss: {validation_loss:.4f}") print(f"held-out token accuracy: {validation_accuracy:.3f}") if smoke_test: print("smoke test: learning thresholds skipped") else: assert validation_loss < baseline_loss * 0.70 assert validation_accuracy > 0.82 return { "baseline_loss": baseline_loss, "validation_loss": validation_loss, "validation_accuracy": validation_accuracy, } def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--steps", type=int, default=300) parser.add_argument( "--smoke-test", action="store_true", help="run at most three optimizer steps and skip learning thresholds", ) return parser.parse_args() if __name__ == "__main__": arguments = parse_args() run(steps=arguments.steps, smoke_test=arguments.smoke_test)