import sys, json, math, random from pathlib import Path import numpy as np import torch import torch.nn as nn sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import get_dataset, make_model, train_model, sweep_baseline, make_report SEEDS = tuple(range(8)) GRID = [{'lr': 1e-3}, {'lr': 3e-3}, {'lr': 6e-3}] EPOCHS, BATCH, NTR, NTE = 10, 128, 1200, 400 def separator_mask_from_graph(adj, max_sep=2): n = adj.shape[0] active = np.ones(n, dtype=bool) sep = [] for _ in range(max_sep): tri = np.zeros(n, dtype=float) deg = (adj & active[:, None] & active[None, :]).sum(1) for i in np.flatnonzero(active): ns = np.flatnonzero(adj[i] & active) tri[i] = float(adj[np.ix_(ns, ns)].sum()) / 2.0 score = np.where(active, tri / (deg + 1.0), -1.0) x = int(np.argmax(score)) if tri[x] < 1: break sep.append(x) active[x] = False active[adj[x]] = False seen = np.zeros(n, dtype=bool) comps = [] for st in np.flatnonzero(active): if seen[st]: continue stack, c = [int(st)], [] seen[st] = True while stack: u = stack.pop(); c.append(u) for v in np.flatnonzero(adj[u] & active): if not seen[v]: seen[v] = True; stack.append(int(v)) comps.append(np.asarray(c, dtype=int)) covered = np.zeros(n, dtype=bool) if sep: covered[sep] = True covered |= adj[:, sep].any(axis=1) mask = covered[:, None] | covered[None, :] for c in comps: mask[np.ix_(c, c)] = True np.fill_diagonal(mask, True) return mask, sep, comps, covered def affinity_graph(x, r=6): # Cheap symmetric proxy from token embeddings, as specified by the idea. with torch.no_grad(): z = x.detach().cpu().numpy().astype(np.float64) # Rows are tokens; retain the sequence length rather than flattening d features. s = (z @ z.T) / math.sqrt(max(1, z.shape[1])) np.fill_diagonal(s, -np.inf) n = len(z); r = min(r, max(1, n - 1)) a = np.zeros((n, n), dtype=bool) for i in range(n): a[i, np.argpartition(s[i], -r)[-r:]] = True a |= a.T np.fill_diagonal(a, False) return a def structural_check(): g = np.zeros((10, 10), dtype=bool) for block in ([0,1,2], [5,6,7]): g[np.ix_(block, block)] = True np.fill_diagonal(g, False) m, sep, comps, covered = separator_mask_from_graph(g, 2) anti = all(not g[np.ix_(comps[i], comps[j])].any() for i in range(len(comps)) for j in range(i)) formula = covered[:, None] | covered[None, :] for c in comps: formula[np.ix_(c, c)] = True return {'separator_size': len(sep), 'components': [len(c) for c in comps], 'anti_adjacent_components': bool(anti), 'mask_formula_exact': bool(np.array_equal(m, formula))} class SeparatorTransformer(nn.Module): """Benchmark-compatible transformer_tiny with graph-masked encoder attention.""" def __init__(self, win, out_dim=1, d=64, depth=2, r=6, max_sep=2): super().__init__() self.inp = nn.Linear(1, d) self.pos = nn.Parameter(torch.zeros(1, win, d)) nn.init.normal_(self.pos, std=.02) layer = nn.TransformerEncoderLayer(d, nhead=2, dim_feedforward=128, batch_first=True, dropout=0.0) self.enc = nn.TransformerEncoder(layer, depth) self.head = nn.Linear(win * d, out_dim) self.r, self.max_sep = r, max_sep self.last_stats = {} def forward(self, x): h = self.inp(x.unsqueeze(-1)) + self.pos[:, :x.shape[1]] # A single graph is built for the current sequence and reused per batch. q = h[0].detach() g = affinity_graph(q, self.r) mask, sep, comps, covered = separator_mask_from_graph(g, self.max_sep) # True means forbidden for TransformerEncoder. attn_mask = torch.from_numpy(~mask).to(h.device) h = self.enc(h, mask=attn_mask) self.last_stats = {'allowed_fraction': float(mask.mean()), 'separator_size': len(sep), 'components': len(comps), 'covered_fraction': float(covered.mean())} return self.head(h.reshape(h.shape[0], -1)) def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def train_one(seed, cfg, idea=False, capture=False): seed_all(seed) d = get_dataset('sequence', seed, n_train=NTR, n_test=NTE) if idea: net = SeparatorTransformer(d['input_shape'][0], d['ytr'].shape[1], r=6, max_sep=2) else: net = make_model('transformer_tiny', d['input_shape'], d['ytr'].shape[1]) net, metric, hist = train_model(net, d, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH) if metric is None: return float('nan'), {} return float(metric), getattr(net, 'last_stats', {}) if capture else {} def main(): check = structural_check() def base_factory(cfg): return lambda seed: train_one(seed, cfg, False)[0] base = sweep_baseline(base_factory, GRID, seeds=(0,1,2,3)) # Same union of lrs on both sides; idea is evaluated at best and two neighbors. idea_cfgs = GRID idea_vals = [] idea_stats = [] best_cfg = base['best_cfg'] for cfg in idea_cfgs: vals = [train_one(s, cfg, True, False)[0] for s in SEEDS] idea_vals.append({'cfg': cfg, 'mean': float(np.nanmean(vals)), 'per_seed': vals}) if cfg == best_cfg: best_vals = vals idea = {'mean': float(np.nanmean(best_vals)), 'std': float(np.nanstd(best_vals)), 'per_seed': [float(v) for v in best_vals], 'n': len(best_vals), 'selected_cfg': best_cfg, 'sweep': idea_vals} # Re-test trained behavior on each paired seed, not an analytic toy identity. for s in SEEDS: _, st = train_one(s, best_cfg, True, True) idea_stats.append(st) sig = {'predicted': {'sparse_allowed_fraction': '< 1.0', 'anti_adjacent_components': True}, 'observed_mean_allowed_fraction': float(np.mean([x.get('allowed_fraction', 1.0) for x in idea_stats])), 'observed_mean_components': float(np.mean([x.get('components', 1) for x in idea_stats])), 'observed_mean_separator_size': float(np.mean([x.get('separator_size', 0) for x in idea_stats])), 'confirmed': bool(np.mean([x.get('allowed_fraction', 1.0) for x in idea_stats]) < 0.95 and np.mean([x.get('components', 1) for x in idea_stats]) > 1)} report = make_report('sequence', 'transformer_tiny', base, idea, {'mechanism_signature': sig, 'structural_check': check, 'custom_track': None, 'protocol_notes': '8 paired seeds; baseline sweep and idea sweep share lr union.'}) Path('bench_report.json').write_text(json.dumps(report, indent=2)) print(json.dumps(report, indent=2)) if __name__ == '__main__': main()