import json, math, os, sys import numpy as np import torch import torch.nn as nn sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import get_dataset, train_model, evaluate, sweep_baseline, make_report SEED = 2716 L = 8 H = 64 EPOCHS = 12 BATCH = 128 # Union of all learning rates tried by either side. LR_GRID = [1e-3, 2e-3, 3e-3, 5e-3] GAMMAS = [0.3, 0.75, 1.25] class StepGRU(nn.Module): """An 8-step GRU stack: one GRUCell per observed dynamics step.""" def __init__(self, mode='iid', gamma=0.75, hidden=H, seed=0): super().__init__() self.inp = nn.Linear(3, hidden) self.cells = nn.ModuleList([nn.GRUCell(hidden, hidden) for _ in range(L)]) self.head = nn.Linear(hidden, 1) self.mode, self.gamma = mode, float(gamma) self._initialize(seed) @staticmethod def _corr(gamma): ix = np.arange(L) return (1.0 + np.abs(ix[:, None] - ix[None, :])) ** (-gamma) def _initialize(self, seed): # Explicit deterministic initialization; only cross-step covariance differs. gen = torch.Generator().manual_seed(int(seed)) with torch.no_grad(): # Standard shared input projection; initialized identically by mode. self.inp.weight.copy_(torch.randn(self.inp.weight.shape, generator=gen) / math.sqrt(3)) self.inp.bias.zero_() C = self._corr(self.gamma) chol = np.linalg.cholesky(C + 1e-8 * np.eye(L)) for name in ('weight_ih', 'weight_hh', 'bias_ih', 'bias_hh'): ps = [getattr(c, name) for c in self.cells] # GRU input matrices differ at step 0 (3 inputs) versus later # steps (hidden inputs); apply the shared-depth process within # each homogeneous shape group. groups = {} for k, p in enumerate(ps): groups.setdefault(tuple(p.shape), []).append((k, p)) for shape, items in groups.items(): if self.mode == 'tied': z = torch.randn(shape, generator=gen) for _, p in items: p.copy_(z) elif self.mode == 'iid' or len(items) < 2: for _, p in items: p.copy_(torch.randn(shape, generator=gen)) else: z = torch.randn((len(items),) + shape, generator=gen) sub = torch.as_tensor(chol[np.ix_([k for k, _ in items], [k for k, _ in items])], dtype=z.dtype) z = sub @ z.reshape(len(items), -1) z = z.reshape((len(items),) + shape) for j, (_, p) in enumerate(items): p.copy_(z[j]) # Match the usual PyTorch GRU scale approximately and make variants comparable. for p in self.head.parameters(): p.copy_(torch.randn(p.shape, generator=gen) * 0.1) self.head.bias.zero_() for c in self.cells: c.weight_ih.mul_(1.0 / math.sqrt(3)) c.weight_hh.mul_(1.0 / math.sqrt(H)) c.bias_ih.zero_(); c.bias_hh.zero_() def forward(self, x): seq = x.view(x.shape[0], L, 3) h = torch.zeros(x.shape[0], H, device=x.device, dtype=x.dtype) for k, cell in enumerate(self.cells): z = self.inp(seq[:, k]) h = cell(z, h) return self.head(h) def make_train(mode, gamma, lr, signature=None): def run(seed): torch.manual_seed(seed); np.random.seed(seed) d = get_dataset('dynamics', seed, n_train=400, n_test=400) net = StepGRU(mode=mode, gamma=gamma, seed=seed) trained, metric, hist = train_model(net, d, epochs=EPOCHS, lr=lr, batch=BATCH, log=print) if signature is not None and trained is not None: with torch.no_grad(): # Trained-model behavior: adjacent hidden-state response covariance # measured over the actual test trajectories. x = d['xte'] seq = x.view(x.shape[0], L, 3) dev = next(trained.parameters()).device seq = seq.to(dev) h = torch.zeros(x.shape[0], H, device=dev) hs = [] for k, cell in enumerate(trained.cells): h = cell(trained.inp(seq[:, k]), h); hs.append(h.detach().cpu().numpy()) hs = np.stack(hs, axis=1) vals = [] for lag in range(1, L): a, b = hs[:, :-lag].reshape(-1, H), hs[:, lag:].reshape(-1, H) vals.append(float(np.mean(a*b) / (np.std(a)*np.std(b)+1e-8))) signature.append({'seed': int(seed), 'metric': float(metric), 'hidden_corr_lags': vals}) return float(metric) return run def main(): # Baseline sweep includes every LR used by the idea, satisfying search-space parity. base_grid = [{'lr': lr, 'mode': 'iid'} for lr in LR_GRID] base = sweep_baseline(lambda cfg: make_train('iid', 0.75, cfg['lr']), base_grid) # Same three-setting idea sweep, evaluated on all eight paired seeds. idea_configs = [{'lr': lr, 'gamma': g} for lr, g in zip([base['best_cfg']['lr'], 2e-3, 5e-3], GAMMAS)] tried = [] best = None for cfg in idea_configs: sig = [] r = evaluate(make_train('power', cfg['gamma'], cfg['lr'], sig)) tried.append({'cfg': cfg, 'result': r, 'signature_samples': sig}) if best is None or r['mean'] < best['result']['mean']: best = {'cfg': cfg, 'result': r, 'signature_samples': sig} idea = best['result'] base['idea_grid'] = tried # Re-test matched trained models for the mechanism signature at the selected config. # Predicted c_lag ~ (1+lag)^(-gamma); compare slope of measured hidden correlations. obs = np.asarray([v for row in best['signature_samples'] for v in row['hidden_corr_lags']]) lags = np.tile(np.arange(1, L), len(best['signature_samples'])) positive = obs > 1e-5 observed_slope = float(np.polyfit(np.log1p(lags[positive]), np.log(obs[positive]), 1)[0]) if positive.sum() > 3 else float('nan') predicted_slope = -float(best['cfg']['gamma']) # Honest tolerance: this nonlinear trained-state signature is only confirmed if # the sign and exponent are reasonably close, not merely because initialization was set so. confirmed = bool(np.isfinite(observed_slope) and observed_slope < 0 and abs(observed_slope - predicted_slope) < 0.45) signature = { 'quantity': 'trained hidden-state cross-step correlation on dynamics test trajectories', 'gamma': best['cfg']['gamma'], 'predicted_loglog_slope': predicted_slope, 'observed_loglog_slope': observed_slope, 'mean_abs_corr_by_lag': np.mean([r['hidden_corr_lags'] for r in best['signature_samples']], axis=0).tolist(), 'confirmed': confirmed } report = make_report('dynamics', 'rnn_small_step_gru', base, idea, {'mechanism_signature': signature, 'baseline_architecture': '8 independent GRUCell steps', 'idea_architecture': 'same 8 GRUCell steps, coordinatewise power-law covariance'}) report['idea_sweep'] = tried report['protocol_note'] = 'Dynamics chosen because the idea concerns recurrent depth and stability; 8 paired seeds, baseline LR sweep, idea 3-config sweep.' with open('bench_report.json', 'w') as f: json.dump(report, f, indent=2) print(json.dumps(report, indent=2)) if __name__ == '__main__': main()