Reproducibility package: UWCA semantic multiple access (TWC submission)

legacy/revision_ablation.py

@@ -0,0 +1,124 @@
+"""
+Ablation / sensitivity studies for the TWC revision.
+  Part 1 (torch): H (heads), K (number of SNR meta-tasks), S (inner steps)
+                  sensitivity -> SER@10dB, cos@10dB, |rho_off|@10dB on HIGH.
+  Part 2 (numpy): refined scalability -> mean cosine + top-k retention vs U.
+"""
+import types
+import numpy as np
+import torch
+import maml_semantic as M
+
+torch.manual_seed(0)
+
+class _Sim:                      # minimal stand-ins (avoid importing the heavy sim module)
+    BATCH = 64
+    @staticmethod
+    def _norm(E):
+        return E / (np.linalg.norm(E, axis=-1, keepdims=True) + 1e-8)
+    @staticmethod
+    def cos_sim(Eh, Eg):
+        return (Eh * Eg).sum(-1)
+sim = _Sim()
+
+def make_cfg(**kw):
+    base = dict(d=64, U=4, H=4, tau=0.45, lam=0.1, snr_min=0.0, snr_max=20.0,
+                snr_step=2.0, inner_lr=0.01, inner_steps=5, outer_lr=1e-3,
+                meta_epochs=70, joint_epochs=70, batch=64, n_mc=60, seed=42,
+                scenario='HIGH', decoder_only=True)
+    base.update(kw)
+    return types.SimpleNamespace(**base)
+
+def train_eval(cfg):
+    device = 'cpu'
+    rng = np.random.default_rng(cfg.seed)
+    scen = M.SCENARIO_CONFIGS[cfg.scenario]
+    model = M.SemanticCommSystem(cfg.d, cfg.U, cfg.H, decoder_only=cfg.decoder_only).to(device)
+    tr = M.MAMLTrainer(model, cfg, device, rng, scen)
+    tr.train()
+    res = M.evaluate_model(model, cfg, device, rng, "rayleigh", scen)
+    snr = np.arange(cfg.snr_min, cfg.snr_max + 1e-6, cfg.snr_step)
+    i10 = int(np.argmin(np.abs(snr - 10)))
+    rho = res['rho'][i10]
+    mask = ~np.eye(cfg.U, dtype=bool)
+    return res['ser'][i10], res['cos'][i10], float(np.abs(rho[mask]).mean()), len(snr)
+
+def part1():
+    print("\n=== ABLATION (HIGH scenario, decoder-only, SNR=10 dB) ===")
+    print("\n-- Attention heads H (d=64) --")
+    for H in [1, 2, 4, 8]:
+        ser, cos, rho, _ = train_eval(make_cfg(H=H))
+        print(f"  H={H}: dk={64//H:2d}  SER={ser:.3f}  cos={cos:.3f}  |rho_off|={rho:.3f}")
+    print("\n-- Number of SNR meta-tasks K (via snr_step) --")
+    for step in [20.0, 10.0, 4.0, 2.0, 1.0]:
+        ser, cos, rho, K = train_eval(make_cfg(snr_step=step))
+        print(f"  K={K:2d} (step={step:>4}): SER={ser:.3f}  cos={cos:.3f}  |rho_off|={rho:.3f}")
+    print("\n-- Inner-loop steps S --")
+    for S in [1, 3, 5, 10]:
+        ser, cos, rho, _ = train_eval(make_cfg(inner_steps=S))
+        print(f"  S={S:2d}: SER={ser:.3f}  cos={cos:.3f}  |rho_off|={rho:.3f}")
+
+# ---- Part 2: refined scalability (cosine + top-k retention) ----
+RNG = np.random.default_rng(11)
+
+def gen_clustered(n, U, D, g, beta):
+    nc = U // g
+    scenes = [RNG.standard_normal(D) for _ in range(nc)]
+    scenes = [s / np.linalg.norm(s) for s in scenes]
+    embs = []
+    for u in range(U):
+        s = scenes[u // g]
+        priv = RNG.standard_normal((n, D)); priv /= np.linalg.norm(priv, axis=-1, keepdims=True) + 1e-8
+        e = np.sqrt(1 - beta**2) * priv + beta * s[None, :]
+        e /= np.linalg.norm(e, axis=-1, keepdims=True) + 1e-8
+        embs.append(e)
+    return np.stack(embs, 1)
+
+def masks_for(U, D):
+    dpu = D // U; Mk = np.zeros((U, D))
+    for u in range(U):
+        Mk[u, u*dpu:(u+1)*dpu] = 1.0
+    return Mk
+
+def se_chan(E, snr_db, Mk):
+    n, U, D = E.shape
+    Ytx = (E * Mk[None]).sum(1)
+    h = np.sqrt(RNG.standard_normal((n, U, 1))**2 + RNG.standard_normal((n, U, 1))**2) * np.sqrt(0.5)
+    nstd = np.sqrt(float(np.mean(Ytx**2)) / (10**(snr_db/10)))
+    return h * Ytx[:, None, :] + RNG.standard_normal((n, U, D)) * nstd
+
+def attn(Yrx, Mk, bm, topk=None):
+    n, U, D = Yrx.shape
+    R = Yrx[:, :, None, :] * Mk[None, None]
+    a = bm.copy(); np.fill_diagonal(a, 1.0)
+    if topk is not None and topk < U:
+        for u in range(U):
+            order = np.argsort(-a[u]); keep = set(order[:topk]) | {u}
+            for v in range(U):
+                if v not in keep: a[u, v] = 0.0
+    a /= a.sum(1, keepdims=True) + 1e-8
+    ctx = np.einsum('ui,buid->bud', a, R)
+    return np.stack([sim._norm(ctx[:, u, :]) for u in range(U)], 1)
+
+def part2():
+    print("\n=== SCALABILITY: cosine + top-k retention (g=4, beta=0.65, SNR=20 dB) ===")
+    g, beta, snr = 4, 0.65, 20.0
+    for U in [4, 8, 16, 32]:
+        D = 16 * U; Mk = masks_for(U, D)
+        bm = np.zeros((U, U))
+        for i in range(U):
+            for j in range(U):
+                if i // g == j // g: bm[i, j] = beta*beta
+        nmc = 150; cf = ct = 0.0
+        for _ in range(nmc):
+            E = gen_clustered(sim.BATCH, U, D, g, beta)
+            cf += sim.cos_sim(attn(se_chan(E, snr, Mk), Mk, bm), E).mean()
+            ct += sim.cos_sim(attn(se_chan(E, snr, Mk), Mk, bm, topk=g), E).mean()
+        cf /= nmc; ct /= nmc
+        print(f"  U={U:3d}  cos_full={cf:.3f}  cos_topk(k={g})={ct:.3f}  "
+              f"retention={100*ct/cf:5.1f}%  ops full={U*U} topk={U*g} ({U//g}x)")
+
+if __name__ == '__main__':
+    part1()
+    part2()
+    print("\nDONE.")