1import json, math, random
  2import numpy as np
  3import torch
  4from torch import nn
  5from sklearn.datasets import load_digits
  6from sklearn.model_selection import train_test_split
  7
  8SEED = 371
  9random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
 10try:
 11    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 12    if device.type == 'cuda':
 13        torch.cuda.get_device_properties(0)
 14except Exception:
 15    device = torch.device('cpu')
 16
 17
 18def binomial_probs(x, n):
 19    """Enumerated P(Y=y|X=x), with a differentiable torch implementation."""
 20    x = x.clamp(1e-6, 1-1e-6)
 21    y = torch.arange(n + 1, device=x.device, dtype=x.dtype)
 22    logc = torch.lgamma(torch.tensor(float(n + 1), device=x.device, dtype=x.dtype)) \
 23        - torch.lgamma(y + 1) - torch.lgamma(torch.tensor(float(n), device=x.device, dtype=x.dtype) - y + 1)
 24    return torch.exp(logc + x[..., None] * 0 + y * torch.log(x[..., None])
 25                     + (n-y) * torch.log1p(-x[..., None]))
 26
 27
 28def beta_binomial_prior(n):
 29    y = np.arange(n + 1, dtype=float)
 30    # comb(n,y) B(y+1/2,n-y+1/2)/B(1/2,1/2), evaluated stably
 31    logq = (math.lgamma(n+1) - np.array([math.lgamma(v+1) for v in y])
 32            - np.array([math.lgamma(n-v+1) for v in y])
 33            + np.array([math.lgamma(v+.5) + math.lgamma(n-v+.5) - math.lgamma(n+1)
 34                        for v in y]) - math.log(math.pi))
 35    q = np.exp(logq); return q / q.sum()
 36
 37
 38def entropy(p):
 39    return -(p.clamp_min(1e-12) * p.clamp_min(1e-12).log()).sum(-1)
 40
 41
 42def bottleneck_terms(x, n):
 43    p = binomial_probs(x, n)
 44    q = p.mean(0)
 45    I = (entropy(q) - entropy(p).mean(0)).mean()
 46    qr = torch.tensor(beta_binomial_prior(n), device=x.device, dtype=x.dtype)
 47    kl = (q.clamp_min(1e-12) * (q.clamp_min(1e-12).log() - qr.log())).sum(-1).mean()
 48    return p, I, kl
 49
 50
 51def derivative_check():
 52    # Check Lemma 7 for several smooth functions and interior probabilities.
 53    n, x = 11, .37
 54    vals = []
 55    for f in [lambda y: float(y*y), lambda y: math.exp(.08*y), lambda y: math.sin(.4*y)]:
 56        def expectation(z, trials):
 57            return sum(math.comb(trials, y)*z**y*(1-z)**(trials-y)*f(y) for y in range(trials+1))
 58        h = 1e-5
 59        lhs = (expectation(x+h,n)-expectation(x-h,n))/(2*h)
 60        rhs = n * sum(math.comb(n-1,y)*x**y*(1-x)**(n-1-y)*(f(y+1)-f(y)) for y in range(n))
 61        vals.append(abs(lhs-rhs))
 62    q = beta_binomial_prior(n)
 63    return {"max_abs_derivative_error": float(max(vals)), "prior_sum": float(q.sum()),
 64            "prior_endpoint_mass": float(q[0]+q[-1]), "prior": q.tolist()}
 65
 66
 67class Net(nn.Module):
 68    def __init__(self, idea=False, n=8):
 69        super().__init__(); self.idea=idea; self.n=n
 70        self.enc=nn.Sequential(nn.Linear(64,32), nn.ReLU(), nn.Linear(32,8))
 71        self.head=nn.Sequential(nn.Linear(8,16), nn.ReLU(), nn.Linear(16,10))
 72    def forward(self, z):
 73        x=torch.sigmoid(self.enc(z))
 74        if self.idea:
 75            p,I,KL=bottleneck_terms(x,self.n)
 76            return self.head(x), I, KL, x, p
 77        return self.head(x), None, None, x, None
 78
 79
 80def run(idea, epochs=45, n=8):
 81    X,y=load_digits(return_X_y=True)
 82    X=X.astype('float32')/16.0
 83    xtr,xte,ytr,yte=train_test_split(X,y,test_size=.25,random_state=SEED,stratify=y)
 84    tr=torch.tensor(xtr,device=device); ty=torch.tensor(ytr,device=device,dtype=torch.long)
 85    te=torch.tensor(xte,device=device); ey=torch.tensor(yte,device=device,dtype=torch.long)
 86    model=Net(idea,n).to(device); opt=torch.optim.Adam(model.parameters(),lr=2e-3,weight_decay=1e-4)
 87    bs=96; history=[]
 88    for ep in range(epochs):
 89        model.train(); perm=torch.randperm(len(tr),device=device); total=0
 90        for ix in perm.split(bs):
 91            out,I,KL,_,_=model(tr[ix]); ce=nn.functional.cross_entropy(out,ty[ix])
 92            # modest shaping: maximize exact batch MI and match the prescribed prior
 93            loss=ce + ((.04*KL - .01*I) if idea else 0)
 94            opt.zero_grad(); loss.backward(); opt.step(); total += ce.item()*len(ix)
 95        if ep in (epochs//2,epochs-1): history.append(total/len(tr))
 96    model.eval()
 97    with torch.no_grad():
 98        out,I,KL,z,p=model(te); pred=out.argmax(1); acc=(pred==ey).float().mean().item()
 99        # evaluate information and prior divergence over the full test set in batches
100        ps=[]
101        for ix in torch.arange(len(te),device=device).split(bs): ps.append(binomial_probs(z[ix],n))
102        pp=torch.cat(ps); q=pp.mean(0); qr=torch.tensor(beta_binomial_prior(n),device=device,dtype=pp.dtype)
103        info=(entropy(q)-entropy(pp).mean(0)).mean().item(); kl=(q*(q.clamp_min(1e-12).log()-qr.log())).sum(-1).mean().item()
104        endpoint=(q[:,0]+q[:,-1]).mean().item()
105        sampled_acc = None
106        if idea:
107            # Discrete deployment path: sample one count per latent coordinate.
108            ys = torch.multinomial(pp.reshape(-1, n+1), 1).reshape(len(te), 8).float() / n
109            sampled_acc = (model.head(ys).argmax(1) == ey).float().mean().item()
110    return {"accuracy":acc,"sampled_count_accuracy":sampled_acc,
111            "test_I_nats":info,"test_KL_to_beta_binomial":kl,
112            "test_endpoint_mass":endpoint,"train_ce_checkpoints":history,
113            "activation_storage_bits_per_coordinate":(math.log2(n+1) if idea else 32.0)}
114
115
116if __name__ == '__main__':
117    # Keep the typo-prone implementation path above explicit and fail-safe by running checks first.
118    result={"device":str(device),"math_check":derivative_check()}
119    result["baseline_continuous"] = run(False)
120    result["capacity_shaped_binomial"] = run(True)
121    print(json.dumps(result, indent=2))