Conditional-information-preserving pooling / toy_experiment.py

Mechanism failed

Raw ⬇ ZIP
  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))