Conditional-information-preserving pooling / toy_experiment.py
Mechanism failed
1import json, math, random
2from itertools import product
3import numpy as np
4
5SEED = 3040
6random.seed(SEED); np.random.seed(SEED)
7
8
9def majority(x):
10 return 1 if sum(x) >= 0 else -1
11
12
13def mutual_info(d, mapping):
14 q = {}
15 for (l, r), v in d.items():
16 x = mapping(l); q[x, r] = q.get((x, r), 0.0) + v
17 px, py = {}, {}
18 for (x, y), v in q.items():
19 px[x] = px.get(x, 0) + v; py[y] = py.get(y, 0) + v
20 return sum(v * math.log2(v / (px[x] * py[y])) for (x, y), v in q.items())
21
22
23def exact_check(p=.15):
24 # Z is a latent remote spin; L is a three-spin local block and R is a
25 # distant noisy observation. This is the paper's majority failure mode.
26 d = {}
27 for z in (-1, 1):
28 for l in product((-1, 1), repeat=3):
29 for r in (-1, 1):
30 v = .5
31 v *= np.prod([(1-p) if x == z else p for x in l])
32 v *= (1-p) if r == z else p
33 d[l, r] = d.get((l, r), 0) + v
34 full = mutual_info(d, lambda x: x)
35 maj = mutual_info(d, majority)
36 # Deterministic pooling gives I(L:R|M)=I(L:R)-I(M:R).
37 post = {}
38 for l in ((1, 1, 1), (1, 1, -1)):
39 a = (1-p)**sum(x == 1 for x in l) * p**sum(x == -1 for x in l)
40 b = p**sum(x == 1 for x in l) * (1-p)**sum(x == -1 for x in l)
41 post[str(l)] = a / (a+b)
42 # Unequal positional weights are injective over the eight binary blocks.
43 code = lambda x: round(sum(w*y for w, y in zip((.7, .2, .1), x)), 8)
44 return {'p': p, 'I(L:R)': full, 'I(majority(L):R)': maj,
45 'discarded_conditional_information_bits': full-maj,
46 'I(injective_pool(L):R)': mutual_info(d, code),
47 'P(Z=+1|L)': post}
48
49
50def neural_demo(p=.15, n=16000):
51 import torch
52 torch.manual_seed(SEED)
53 requested = 'cuda' if torch.cuda.is_available() else 'cpu'
54 try:
55 dev = torch.device(requested)
56 z = torch.where(torch.rand(n, device=dev) < .5, -torch.ones(n, device=dev), torch.ones(n, device=dev))
57 l = torch.where(torch.rand(n, 3, device=dev) < p, -z[:, None], z[:, None]).float()
58 r = torch.where(torch.rand(n, device=dev) < p, -z, z).float()
59 cut = n * 3 // 4
60 tr, te = torch.arange(cut, device=dev), torch.arange(cut, n, device=dev)
61
62 def train_pool(kind):
63 if kind == 'mean':
64 s = l.mean(1, keepdim=True); params = []
65 elif kind == 'majority':
66 s = torch.sign(l.sum(1)).reshape(-1, 1); s[s == 0] = 1; params = []
67 else:
68 # Slightly asymmetric initialization breaks the permutation
69 # symmetry; unlike majority, the scalar can retain all 3 bits.
70 logits = torch.nn.Parameter(torch.log(torch.tensor([.7,.2,.1], device=dev)))
71 params = [logits]
72 s = None
73 head = torch.nn.Sequential(torch.nn.Linear(1, 24), torch.nn.Tanh(),
74 torch.nn.Linear(24, 1)).to(dev)
75 params += list(head.parameters())
76 opt = torch.optim.Adam(params, lr=.025)
77 y = (r + 1) / 2
78 for _ in range(900):
79 if kind == 'learned':
80 a = torch.softmax(logits, 0); s = (l*a).sum(1, keepdim=True)
81 loss = torch.nn.functional.binary_cross_entropy_with_logits(head(s[tr]).squeeze(1), y[tr])
82 opt.zero_grad(); loss.backward(); opt.step()
83 with torch.no_grad():
84 if kind == 'learned':
85 a = torch.softmax(logits, 0); s = (l*a).sum(1, keepdim=True)
86 pred = head(s[te]).squeeze(1)
87 nll = torch.nn.functional.binary_cross_entropy_with_logits(pred, y[te]).item()
88 acc = ((pred > 0) == (r[te] > 0)).float().mean().item()
89 result = {'heldout_remote_accuracy': acc, 'heldout_remote_nll': nll}
90 if kind == 'learned': result['weights'] = a.detach().cpu().numpy().tolist()
91 return result
92
93 # q_phi has the microscopic source block; q_psi has only the pooled S.
94 # Their held-out NLL difference is the empirical information gap.
95 aux = torch.nn.Sequential(torch.nn.Linear(3, 24), torch.nn.Tanh(), torch.nn.Linear(24,1)).to(dev)
96 opt = torch.optim.Adam(aux.parameters(), lr=.025); y=(r+1)/2
97 for _ in range(900):
98 loss=torch.nn.functional.binary_cross_entropy_with_logits(aux(l[tr]).squeeze(1), y[tr])
99 opt.zero_grad(); loss.backward(); opt.step()
100 with torch.no_grad():
101 nll_full=torch.nn.functional.binary_cross_entropy_with_logits(aux(l[te]).squeeze(1),y[te]).item()
102 out = {k: train_pool(k) for k in ('mean','majority','learned')}
103 out['full_source_predictor_nll'] = nll_full
104 out['mean_predictor_gap_nll'] = out['mean']['heldout_remote_nll'] - nll_full
105 out['majority_predictor_gap_nll'] = out['majority']['heldout_remote_nll'] - nll_full
106 out['learned_predictor_gap_nll'] = out['learned']['heldout_remote_nll'] - nll_full
107 out['device'] = str(dev)
108 return out
109 except Exception as e:
110 return {'device': 'cpu-fallback', 'error': repr(e)}
111
112
113if __name__ == '__main__':
114 print(json.dumps({'exact_check': exact_check(), 'neural_demo': neural_demo()}, indent=2))