1import json
  2import random
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7SEED = 236
  8random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  9torch.set_num_threads(4)
 10DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 11
 12
 13def tv(p, q):
 14    return 0.5 * torch.abs(p - q).sum(dim=-1)
 15
 16
 17def make_data(n, vocab=24, modes=3, hdim=3, noise=0.35, seed=236):
 18    g = np.random.default_rng(seed)
 19    # Distinct recurring continuation distributions, with modest within-mode variation.
 20    centers = g.dirichlet(np.full(vocab, 0.25), size=modes).astype('float32')
 21    labels = g.integers(0, modes, size=n)
 22    p = np.empty((n, vocab), dtype='float32')
 23    for i, z in enumerate(labels):
 24        concentration = 70.0 * centers[z] + 0.25
 25        p[i] = g.dirichlet(concentration).astype('float32')
 26    # Anchor representation is informative but imperfect: one-hot mode signal + Gaussian noise.
 27    h = np.eye(modes, hdim, dtype='float32')[labels]
 28    h += noise * g.normal(size=h.shape).astype('float32')
 29    return torch.tensor(h), torch.tensor(p), torch.tensor(labels), torch.tensor(centers)
 30
 31
 32class SharedProposal(nn.Module):
 33    def __init__(self, vocab):
 34        super().__init__()
 35        self.logits = nn.Parameter(torch.zeros(vocab))
 36    def forward(self, h):
 37        return self.logits.softmax(-1).expand(h.shape[0], -1)
 38
 39
 40class RoutedPrototypes(nn.Module):
 41    def __init__(self, hdim, vocab, k):
 42        super().__init__()
 43        self.prototypes = nn.Parameter(torch.randn(k, vocab) * 0.1)
 44        self.router = nn.Linear(hdim, k)
 45    def forward(self, h):
 46        q = self.prototypes.softmax(-1)
 47        pi = self.router(h).softmax(-1)
 48        mixture = pi @ q
 49        return mixture, q, pi
 50
 51
 52def train_shared(h, p, steps=900, lr=0.08):
 53    m = SharedProposal(p.shape[1]).to(DEVICE)
 54    opt = torch.optim.Adam(m.parameters(), lr=lr)
 55    h, p = h.to(DEVICE), p.to(DEVICE)
 56    for _ in range(steps):
 57        loss = tv(p, m(h)).mean()
 58        opt.zero_grad(); loss.backward(); opt.step()
 59    return m
 60
 61
 62def train_routed(h, p, k, steps=1200, lr=0.06, entropy=0.002):
 63    m = RoutedPrototypes(h.shape[1], p.shape[1], k).to(DEVICE)
 64    opt = torch.optim.Adam(m.parameters(), lr=lr)
 65    h, p = h.to(DEVICE), p.to(DEVICE)
 66    for _ in range(steps):
 67        mixture, q, pi = m(h)
 68        # The paper's soft surrogate, with a small positive entropy penalty.
 69        distances = 0.5 * torch.abs(p[:, None, :] - q[None, :, :]).sum(-1)
 70        loss = (pi * distances).sum(-1).mean() + entropy * (-(pi * (pi + 1e-8).log()).sum(-1).mean())
 71        opt.zero_grad(); loss.backward(); opt.step()
 72    return m
 73
 74
 75def evaluate(shared, routed, h, p, labels):
 76    hdev, pdev = h.to(DEVICE), p.to(DEVICE)
 77    with torch.no_grad():
 78        qb = shared(hdev)
 79        mix, q, pi = routed(hdev)
 80        d = 0.5 * torch.abs(pdev[:, None, :] - q[None, :, :]).sum(-1)
 81        oracle = d.min(-1).values
 82        selected = d.gather(1, pi.argmax(-1, keepdim=True)).squeeze(1)
 83        actual_mix = tv(pdev, mix)
 84        # A simple rejection/verification proxy: maximal coupling acceptance.
 85        out = {
 86            'shared_tv': tv(pdev, qb).mean().item(),
 87            'oracle_prototype_tv': oracle.mean().item(),
 88            'routed_selected_tv': selected.mean().item(),
 89            'mixture_tv': actual_mix.mean().item(),
 90            'oracle_reduction_vs_shared': (tv(pdev, qb).mean() - oracle.mean()).item(),
 91            'routed_reduction_vs_shared': (tv(pdev, qb).mean() - selected.mean()).item(),
 92            'mixture_reduction_vs_shared': (tv(pdev, qb).mean() - actual_mix.mean()).item(),
 93            'shared_accept_proxy': (1-tv(pdev, qb)).mean().item(),
 94            'routed_accept_proxy': (1-selected).mean().item(),
 95            'mixture_accept_proxy': (1-actual_mix).mean().item(),
 96            'router_accuracy': (pi.argmax(-1).cpu() == labels).float().mean().item(),
 97            'prototype_count': q.shape[0],
 98        }
 99    return out
100
101
102def exact_math_check():
103    # Two sharply separated distributions: K=1 has a nonzero information floor,
104    # while two prototypes attain zero oracle distance.
105    p = torch.tensor([[0.92, 0.04, 0.02, 0.02], [0.03, 0.03, 0.04, 0.90]])
106    q_shared = p.mean(0, keepdim=True).expand_as(p)
107    q_proto = p.clone()
108    l1 = tv(p, q_shared).mean().item()
109    lk = tv(p, q_proto).min(-1).values.mean().item()
110    return {'two_mode_L1': l1, 'two_mode_L2_oracle': lk, 'reduction': l1-lk,
111            'bound_sanity': bool(l1 >= lk - 1e-8 and lk < 1e-8)}
112
113
114def main():
115    math = exact_math_check()
116    htr, ptr, ytr, _ = make_data(1800, seed=SEED)
117    hte, pte, yte, _ = make_data(900, seed=SEED+1)
118    # Use the same underlying centers for a fair held-out test by generating test labels
119    # from training centers with fresh within-mode samples.
120    rng = np.random.default_rng(SEED+1)
121    centers = np.array(make_data(1, seed=SEED)[3])
122    labels = rng.integers(0, 3, size=900)
123    ptest = np.array([rng.dirichlet(70*centers[z] + .25) for z in labels]).astype('float32')
124    htest = np.eye(3, 3, dtype='float32')[labels] + .35*rng.normal(size=(900,3)).astype('float32')
125    hte, pte, yte = torch.tensor(htest), torch.tensor(ptest), torch.tensor(labels)
126    shared = train_shared(htr, ptr)
127    results = {'device': DEVICE, 'math_check': math, 'train_n': len(htr), 'test_n': len(hte)}
128    # Train on identical data and compare K=2 and K=3 to the standard single head.
129    for k in (2, 3):
130        routed = train_routed(htr, ptr, k)
131        results[f'K{k}'] = evaluate(shared, routed, hte, pte, yte)
132    # Parameter/FLOP accounting for the final vocabulary heads (shared hidden compute).
133    results['head_parameter_multiplier_K3'] = 3.0
134    results['note'] = 'Output-head-only multiplier; shared hidden/router compute is not included in this synthetic proxy.'
135    print(json.dumps(results, indent=2, sort_keys=True))
136
137if __name__ == '__main__':
138    try:
139        main()
140    except (RuntimeError, torch.cuda.OutOfMemoryError) as e:
141        if DEVICE == 'cuda':
142            print(json.dumps({'cuda_error': str(e), 'fallback': 'rerun with CUDA_VISIBLE_DEVICES empty'}))
143            raise
144        raise