import json, sys, time 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, train_model, evaluate, sweep_baseline, make_report SEED = 2052 N_SUPPORT = 64 DH = 64 EPOCHS = 20 # Union of step sizes is shared by baseline and idea configs. LR_GRID = [1e-3, 3e-3, 1e-2] SIGMA_GRID = [0.4, 0.6, 0.8] records = {'baseline': {}, 'idea': {}} class FrozenAttentionPrompt(nn.Module): """One frozen single-head attention layer, with trainable prompt only.""" def __init__(self, d_in=32, dh=DH, n_prompt=N_SUPPORT, seed=0): super().__init__() r = np.random.RandomState(seed) q = r.normal(0, 1/np.sqrt(dh), (dh, dh)).astype('float32') k = r.normal(0, 1/np.sqrt(dh), (dh, dh)).astype('float32') v = r.normal(0, 1/np.sqrt(dh), (dh, 1)).astype('float32') self.register_buffer('Q', torch.from_numpy(q)) self.register_buffer('K', torch.from_numpy(k)) self.register_buffer('V', torch.from_numpy(v)) self.prompt = nn.Parameter(torch.zeros(n_prompt, dh)) nn.init.normal_(self.prompt, std=0.05) self.d_in, self.dh = d_in, dh def embed(self, x): z = torch.zeros(x.shape[0], self.dh, device=x.device, dtype=x.dtype) z[:, :self.d_in] = x z[:, self.d_in] = 1. return z def forward(self, x): e = self.embed(x) q = e @ self.Q keys = self.prompt @ self.K logits = q @ keys.T / np.sqrt(self.dh) w = torch.softmax(logits, dim=1) vals = self.prompt @ self.V return w @ vals def set_seed(seed): np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def dataset(seed): d = get_dataset('sequence', seed, n_train=400, n_test=400) # Bench train_model expects torch tensors; get_dataset already supplies them. return d def baseline_fn(cfg): def run(seed): set_seed(seed) ds = dataset(seed) model = FrozenAttentionPrompt(seed=seed) model, metric, _ = train_model(model, ds, epochs=EPOCHS, lr=cfg['lr'], batch=128, weight_decay=0.0, log=lambda *_: None) if model is None: return float('nan') with torch.no_grad(): device = next(model.parameters()).device x = ds['xte'].to(device); pred = model(x) # observed trained-model attention entropy and prompt norm e = model.embed(x[:32]); logits = (e @ model.Q) @ (model.prompt @ model.K).T / np.sqrt(DH) records['baseline'][int(seed)] = {'prompt_norm': float(model.prompt.norm()), 'attention_entropy': float((-(torch.softmax(logits,1)*torch.log_softmax(logits,1)).sum(1)).mean())} return float(metric) return run def analytic_prompt(xs, ys, sigma, seed): r = np.random.RandomState(seed) q = r.normal(0, 1/np.sqrt(DH), (DH, DH)) k = r.normal(0, 1/np.sqrt(DH), (DH, DH)) v = r.normal(0, 1/np.sqrt(DH), (DH, 1)) A = (q.T @ k)[:33] / np.sqrt(DH) B = np.vstack([A, v.T]) c = np.column_stack([xs / sigma**2, -np.sum(xs*xs, axis=1)/(2*sigma**2), ys]) P = np.linalg.lstsq(B, c.T, rcond=None)[0].T return P, q, k, v def idea_fn(cfg): def run(seed): ds = dataset(seed) xs = ds['xtr'][:N_SUPPORT].numpy().astype('float64') ys = ds['ytr'][:N_SUPPORT].numpy().astype('float64').reshape(-1,1) P,q,k,v = analytic_prompt(xs, ys, cfg['sigma'], seed) xt = ds['xte'].numpy().astype('float64') e = np.zeros((len(xt), DH)); e[:,:32] = xt; e[:,32] = 1 logits = e @ ((q.T @ k) @ P.T) / np.sqrt(DH) logits -= logits.max(1, keepdims=True) w = np.exp(logits); w /= w.sum(1, keepdims=True) pred = w @ (P @ v) metric = float(np.mean((pred[:,0] - ds['yte'].numpy())**2)) target = xt @ xs.T / cfg['sigma']**2 - (xs*xs).sum(1)[None,:]/(2*cfg['sigma']**2) # common query term is intentionally omitted, as it cancels in softmax. residual = float(np.sqrt(np.mean((e @ ((q.T@k)@P.T)/np.sqrt(DH)-target)**2))) records['idea'].setdefault(cfg['sigma'], {})[int(seed)] = { 'prompt_norm': float(np.linalg.norm(P,axis=1).mean()), 'logit_residual': residual, 'entropy': float(np.mean(-(w*np.log(np.maximum(w,1e-30))).sum(1))) } return metric return run def main(): # Baseline sweep covers all rates used by idea-side comparison. base_grid = [{'lr': x} for x in LR_GRID] base = sweep_baseline(baseline_fn, base_grid) # Idea sweep has three nearby bandwidth settings; each is evaluated on all 8 seeds. idea_runs = [] for sigma, lr in zip(SIGMA_GRID, LR_GRID): r = evaluate(idea_fn({'sigma': sigma, 'lr': lr})) idea_runs.append({'cfg': {'sigma': sigma, 'lr': lr}, 'result': r}) best = min(idea_runs, key=lambda z: z['result']['mean']) extra = {'track_choice': 'sequence: multi-token temporal-window correlations require sequence attention', 'best_sigma': best['cfg']['sigma'], 'idea_sweep': idea_runs, 'trained_baseline_behavior': records['baseline']} sigs = [records['idea'][s] for s in SIGMA_GRID if s in records['idea']] if sigs: means = [float(np.mean([x['prompt_norm'] for x in z.values()])) for z in sigs] ref = means[1] best_resid = float(np.mean([x['logit_residual'] for x in records['idea'][best['cfg']['sigma']].values()])) extra['mechanism_prediction'] = {'prediction': 'prompt norm scales as sigma^-2; affine logit residual is near zero', 'observed_prompt_norms': dict(zip(SIGMA_GRID, means)), 'predicted_norm_ratio_sigma_0.4_to_0.8': 4.0, 'observed_norm_ratio_sigma_0.4_to_0.8': means[0]/means[-1], 'observed_best_logit_residual': best_resid, 'confirmed': bool(abs(means[0]/means[-1]-4.0) < 0.5 and best_resid < 1e-4)} report = make_report('sequence', 'frozen_random_attention_prompt', base, best['result'], extra) report['baseline_sweep_shared_lr_grid'] = LR_GRID report['cut'] = 'Only sequence track tested; no CIFAR/MNIST transfer and no latency/FLOP study.' Path('bench_report.json').write_text(json.dumps(report, indent=2)) print(json.dumps(report, indent=2)) if __name__ == '__main__': main()