1import json, random
  2import numpy as np
  3import torch
  4
  5SEED = 2303
  6np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  7EPS = 1e-12
  8
  9def normalize(x):
 10    x = np.maximum(np.asarray(x, dtype=float), EPS)
 11    return x / x.sum()
 12
 13def ait(p, q):
 14    p, q = normalize(p), normalize(q)
 15    D = len(p)
 16    lr = np.log(p[:, None] / p[None, :]) - np.log(q[:, None] / q[None, :])
 17    return float((lr * lr).sum() / (2 * D))
 18
 19def entropy(p):
 20    p = normalize(p)
 21    return float(-(p * np.log(p)).sum())
 22
 23def hbin(s):
 24    return float(-s*np.log(s) - (1-s)*np.log(1-s))
 25
 26def row(s, c):
 27    return np.r_[s, (1-s)*normalize(c)]
 28
 29def clr(c):
 30    z = np.log(normalize(c)); return z-z.mean()
 31
 32def checks():
 33    # Prediction 1: fixed content gives an exact quadratic sink dependence.
 34    # Direct expansion of the supplied pairwise definition gives m/D*b^2,
 35    # not m/(2D)*b^2 as stated in the idea.
 36    m, D = 4, 5
 37    c = normalize([.50, .30, .15, .05])
 38    s = .80
 39    sink_rows = []
 40    for t in np.linspace(.08, .92, 18):
 41        b = np.log(s*(1-t)/(t*(1-s)))
 42        actual = ait(row(s,c), row(t,c))
 43        sink_rows.append((float(b*b), actual, float(b)))
 44    x = np.array([v[0] for v in sink_rows]); y = np.array([v[1] for v in sink_rows])
 45    coef = float(np.dot(x,y)/np.dot(x,x))
 46    correct = m/D
 47    claimed = m/(2*D)
 48
 49    # Prediction 2: the implemented content target is exactly quadratic in a
 50    # clr displacement, independent of sink mass.
 51    base = normalize([.50, .30, .15, .05])
 52    content_rows = []
 53    for a in np.linspace(0, 2.5, 12):
 54        c1 = normalize(np.exp(np.log(base) + a*clr(base)))
 55        c2 = normalize(np.exp(np.log(base) - a*clr(base)))
 56        implemented = float(np.sum((clr(c1)-clr(c2))**2)/m)
 57        expected = float((2*a)**2*np.sum(clr(base)**2)/m)
 58        content_rows.append((float(a), implemented, expected))
 59    content_max_err = max(abs(r[1]-r[2]) for r in content_rows)
 60
 61    # Prediction 3: entropy decomposition is exact.
 62    ent_err = []
 63    for _ in range(20):
 64        z = np.random.dirichlet(np.ones(D)); s0 = z[0]
 65        ent_err.append(abs(entropy(z) - (hbin(s0) + (1-s0)*entropy(z[1:]/(1-s0)))))
 66
 67    # Explicit counterexample to the claimed full decomposition: same content
 68    # makes the discrepancy especially transparent.
 69    s2, t2 = .82, .63
 70    b2 = np.log(s2*(1-t2)/(t2*(1-s2)))
 71    actual2 = ait(row(s2,base), row(t2,base))
 72    claimed2 = m/(2*D)*b2*b2
 73    return {'prediction_1_sink_quadratic': {
 74                'observed_coefficient': coef, 'predicted_from_definition_m_over_D': correct,
 75                'paper_claim_m_over_2D': claimed, 'relative_error_correct': abs(coef-correct)/correct,
 76                'relative_error_paper_claim': abs(coef-claimed)/claimed},
 77            'prediction_2_content_quadratic': {
 78                'max_abs_error': content_max_err,
 79                'mean_ratio_implemented_to_expected': float(np.mean([r[1]/r[2] for r in content_rows[1:]]))},
 80            'prediction_3_entropy_identity_max_abs_error': max(ent_err),
 81            'decomposition_counterexample_same_content': {
 82                'actual_pairwise_aitchison': actual2, 'paper_formula': claimed2,
 83                'ratio_actual_to_paper': actual2/claimed2},
 84            'sink_sweep': sink_rows, 'content_sweep': content_rows}
 85
 86def toy_distill():
 87    # Sink-heavy teacher and deliberately low-capacity student (one shared row).
 88    torch.manual_seed(SEED)
 89    R, M = 64, 7
 90    sink = torch.full((R,1), .985)
 91    tc = torch.softmax(torch.randn(R,M)*1.1, dim=-1)
 92    teacher = torch.cat([sink, (1-sink)*tc], dim=1)
 93    def run(kind):
 94        sink_logit = torch.tensor(0., requires_grad=True)
 95        content_logits = torch.zeros(M, requires_grad=True)
 96        opt = torch.optim.SGD([sink_logit, content_logits], lr=.12)
 97        grad_max = 0.
 98        for _ in range(300):
 99            t = torch.sigmoid(sink_logit); c = torch.softmax(content_logits, dim=-1)
100            student = torch.cat([t.expand(R,1), ((1-t)*c).expand(R,M)], dim=1)
101            if kind == 'kl':
102                loss = (teacher * (teacher.clamp_min(1e-12).log()-student.clamp_min(1e-12).log())).sum(1).mean()
103            else:
104                sinkloss = (torch.log(t/(1-t))-torch.log(teacher[:,0]/(1-teacher[:,0]))).square().mean()
105                tz = torch.log(teacher[:,1:]/(1-teacher[:,0:1])); tz -= tz.mean(1, keepdim=True)
106                sz = torch.log(c); sz -= sz.mean()
107                contentloss = (tz-sz).square().sum(1).mean()/M
108                loss = sinkloss + contentloss
109            opt.zero_grad(); loss.backward()
110            grad_max = max(grad_max, float(max(abs(sink_logit.grad.item()), content_logits.grad.abs().max().item())))
111            opt.step()
112        with torch.no_grad():
113            t = torch.sigmoid(sink_logit); c = torch.softmax(content_logits,0)
114            st = torch.cat([t.expand(R,1), ((1-t)*c).expand(R,M)],1)
115            sink_err = float((st[:,0]-teacher[:,0]).abs().mean())
116            cz = torch.log(teacher[:,1:]/(1-teacher[:,0:1])); cz-=cz.mean(1,keepdim=True)
117            sz = torch.log(c); sz-=sz.mean()
118            cd = float(((cz-sz)**2).sum(1).mean()/M)
119            kl = float((teacher*(teacher.log()-st.log())).sum(1).mean())
120        return {'final_kl':kl,'sink_abs_error':sink_err,'content_aitchison_squared':cd,'max_gradient':grad_max}
121    return {'kl_baseline':run('kl'), 'sink_content_aitchison':run('idea')}
122
123def main():
124    out = {'seed': SEED, 'math_checks': checks(), 'toy_distillation': toy_distill()}
125    with open('results.json','w') as f: json.dump(out,f,indent=2)
126    print(json.dumps(out, indent=2))
127
128if __name__ == '__main__': main()