Giorgio Gilestro
aca7b394a3
configs/neural/{N0,N1,N2,N5}.yaml -> {bridge,collapse,grounding,architectures}.yaml,
results dirs likewise. Updated experiment/output.dir fields, comments/docstrings, and
docs; regenerated the four result manifests (now carrying the real git commit). No
functional path resolution referenced the codes (the Makefile globs configs/neural/*.yaml
and tests use inline configs), so nothing breaks. 92 tests green.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
109 lines
4.7 KiB
Python
109 lines
4.7 KiB
Python
"""Autoregressive MLP generative model (Stage C, for the `architectures` generality axis).
|
|
|
|
A causal feed-forward next-token model: token ``i`` is predicted from the concatenated
|
|
(causally-masked) embeddings of all earlier tokens. Deliberately a *different* inductive
|
|
bias from the GRU — if collapse appears here too, it is a property of the transmission
|
|
operator, not of any one architecture.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
from .config import ModelCfg, SyntheticCfg
|
|
from .oracle import Oracle
|
|
from .torch_models import _BaseTorchGenerator
|
|
from .train import device_generator, seed_everything
|
|
|
|
|
|
def _make_mlp_net(V: int, L: int, embed: int, hidden: int):
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
class MLPNet(nn.Module):
|
|
"""Predict every position from a causally-masked flatten of prior embeddings."""
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.V, self.L, self.E = V, L, embed
|
|
self.bos = V
|
|
self.embed = nn.Embedding(V + 1, embed)
|
|
self.net = nn.Sequential(
|
|
nn.Linear(L * embed, hidden), nn.ReLU(),
|
|
nn.Linear(hidden, hidden), nn.ReLU(),
|
|
nn.Linear(hidden, V),
|
|
)
|
|
# lower-triangular INCLUSIVE mask over input positions: position i sees inputs
|
|
# 0..i (the input is already shifted by one, so this is strictly causal on x).
|
|
mask = torch.tril(torch.ones(L, L))
|
|
self.register_buffer("mask", mask)
|
|
|
|
def _context(self, inp): # inp: (B, L) input tokens
|
|
B = inp.shape[0]
|
|
emb = self.embed(inp) # (B, L, E)
|
|
m = self.mask.to(emb.dtype) # (L, L)
|
|
# ctx[b, i] = concat_j ( emb[b, j] * mask[i, j] ) -> (B, L, L*E)
|
|
ctx = emb.unsqueeze(1) * m.unsqueeze(0).unsqueeze(-1) # (B, L, L, E)
|
|
return ctx.reshape(B, self.L, self.L * self.E)
|
|
|
|
def forward(self, x): # x: (B, L) targets
|
|
B = x.shape[0]
|
|
bos = torch.full((B, 1), self.bos, dtype=torch.long, device=x.device)
|
|
inp = torch.cat([bos, x[:, :-1]], dim=1)
|
|
ctx = self._context(inp)
|
|
return self.net(ctx) # (B, L, V)
|
|
|
|
def step_logits(self, prefix): # prefix: (B, pos) tokens so far
|
|
"""Logits for the next token given the tokens generated so far."""
|
|
B, pos = prefix.shape
|
|
bos = torch.full((B, 1), self.bos, dtype=torch.long, device=prefix.device)
|
|
inp = torch.cat([bos, prefix], dim=1)[:, : self.L] # (B, <=L)
|
|
if inp.shape[1] < self.L:
|
|
pad = torch.zeros((B, self.L - inp.shape[1]), dtype=torch.long,
|
|
device=prefix.device)
|
|
inp = torch.cat([inp, pad], dim=1)
|
|
ctx = self._context(inp) # (B, L, L*E)
|
|
return self.net(ctx[:, pos, :]) # logits at position `pos`
|
|
|
|
return MLPNet()
|
|
|
|
|
|
class MLPGenerator(_BaseTorchGenerator):
|
|
"""Autoregressive feed-forward generative model over token sequences."""
|
|
|
|
def fit(self, X: np.ndarray, rng: np.random.Generator) -> None:
|
|
import torch
|
|
|
|
g = seed_everything(int(rng.integers(2 ** 31)))
|
|
net = _make_mlp_net(self.V, self.L, self.mcfg.embed, self.mcfg.hidden).to(self.device)
|
|
net.train()
|
|
opt = torch.optim.Adam(net.parameters(), lr=self.mcfg.lr)
|
|
loss_fn = torch.nn.CrossEntropyLoss()
|
|
data = torch.as_tensor(np.asarray(X), dtype=torch.long, device=self.device)
|
|
n, bs = data.shape[0], self.mcfg.batch_size
|
|
for _ in range(self.mcfg.epochs):
|
|
perm = torch.randperm(n, generator=g).to(self.device)
|
|
for i in range(0, n, bs):
|
|
batch = data[perm[i:i + bs]]
|
|
logits = net(batch)
|
|
loss = loss_fn(logits.reshape(-1, self.V), batch.reshape(-1))
|
|
opt.zero_grad()
|
|
loss.backward()
|
|
opt.step()
|
|
net.eval()
|
|
self.net = net
|
|
|
|
def sample(self, n: int, rng: np.random.Generator) -> np.ndarray:
|
|
import torch
|
|
|
|
if self.net is None:
|
|
raise RuntimeError("MLPGenerator.sample called before fit")
|
|
g = device_generator(int(rng.integers(2 ** 31)), self.device)
|
|
prefix = torch.empty((n, 0), dtype=torch.long, device=self.device)
|
|
with torch.no_grad():
|
|
for pos in range(self.L):
|
|
logits = self.net.step_logits(prefix)
|
|
probs = torch.softmax(logits, dim=-1)
|
|
tok = torch.multinomial(probs, 1, generator=g)
|
|
prefix = torch.cat([prefix, tok], dim=1)
|
|
return prefix.cpu().numpy()
|