126 lines
4.6 KiB
Python
Executable File
126 lines
4.6 KiB
Python
Executable File
"""E8 v2 — End-to-end joint encoder-decoder training (R1.5), source-anchored.
|
|
|
|
All fidelity metrics are measured against the SOURCE embedding normalize(x),
|
|
never against the trainable encoder output (a moving target that makes
|
|
collapse look like success).
|
|
|
|
Configs (HIGH):
|
|
frozen : identity encoder, decoder trained (paper reference)
|
|
e2e_moving : trainable encoder, loss vs f_phi(x) (collapse demo)
|
|
e2e_anchored : trainable encoder, loss vs normalize(x) (fixed anchor)
|
|
e2e_vicreg : trainable encoder, loss vs f_phi(x) + VICReg anti-collapse
|
|
Metrics: SER/cos vs normalize(x); batch nearest-neighbor retrieval accuracy;
|
|
encoder-output effective rank and off-diagonal correlation.
|
|
"""
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
import lib
|
|
from lib import (SCENARIOS, SNR_GRID, UWCA, DEVICE, Encoder, block_masks,
|
|
channel, gen_embeddings, save_json, semantic_loss, set_seed)
|
|
|
|
rng = set_seed(42)
|
|
d, U, H = 64, 4, 4
|
|
masks = block_masks(U, d)
|
|
scen = SCENARIOS["HIGH"]
|
|
snrs = [0.0, 4.0, 8.0, 12.0, 16.0, 20.0]
|
|
|
|
|
|
def gen(n):
|
|
return gen_embeddings(n, d, U, rng, scen).to(DEVICE)
|
|
|
|
|
|
def vicreg_reg(Z):
|
|
Zc = Z - Z.mean(0, keepdim=True)
|
|
std = (Zc.var(0) + 1e-4).sqrt()
|
|
v = F.relu(1.0 / d ** 0.5 - std).mean()
|
|
C = (Zc.T @ Zc) / (Z.shape[0] - 1)
|
|
off = C - torch.diag(torch.diag(C))
|
|
c = (off ** 2).sum() / d
|
|
return 25.0 * v + 100.0 * c
|
|
|
|
|
|
def train(mode, epochs=300):
|
|
model = UWCA(d, U, H).to(DEVICE)
|
|
enc = Encoder(d).to(DEVICE) if mode != "frozen" else None
|
|
mask_p = [p for nm, p in model.named_parameters() if "mask_logits" in nm]
|
|
other = [p for nm, p in model.named_parameters() if "mask_logits" not in nm]
|
|
groups = [{"params": other, "lr": 1e-3}, {"params": mask_p, "lr": 0.1}]
|
|
if enc is not None:
|
|
groups.append({"params": enc.parameters(), "lr": 1e-3})
|
|
opt = torch.optim.Adam(groups)
|
|
for ep in range(1, epochs + 1):
|
|
opt.zero_grad()
|
|
loss = 0.0
|
|
for s in snrs:
|
|
X = gen(64)
|
|
n = X.shape[0]
|
|
E = enc(X.reshape(-1, d)).reshape(n, U, d) if enc is not None else X
|
|
target = E if mode in ("frozen", "e2e_moving", "e2e_vicreg") else X
|
|
ch = channel(E, snr_db=s)
|
|
Eh = model(ch["yI"], ch["yQ"])
|
|
L = semantic_loss(Eh, target.detach()
|
|
if mode == "e2e_moving_detach" else target, 0.1)
|
|
if mode == "e2e_vicreg":
|
|
L = L + vicreg_reg(E.reshape(-1, d))
|
|
loss = loss + L
|
|
(loss / len(snrs)).backward()
|
|
torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
|
|
opt.step()
|
|
if ep % 75 == 0:
|
|
print(f" [E8v2-{mode}] ep {ep}/{epochs} "
|
|
f"loss={float(loss)/len(snrs):.4f}", flush=True)
|
|
return model, enc
|
|
|
|
|
|
@torch.no_grad()
|
|
def collapse_metrics(enc):
|
|
X = gen(1024).reshape(-1, d)
|
|
Z = enc(X) if enc is not None else F.normalize(X, dim=-1)
|
|
Zc = Z - Z.mean(0, keepdim=True)
|
|
C = (Zc.T @ Zc) / (Z.shape[0] - 1)
|
|
ev = torch.linalg.eigvalsh(C).clamp(min=1e-12)
|
|
p = ev / ev.sum()
|
|
return float(torch.exp(-(p * p.log()).sum()))
|
|
|
|
|
|
@torch.no_grad()
|
|
def curves(model, enc):
|
|
ss, cc, rr = [], [], []
|
|
for snr in SNR_GRID:
|
|
a = b = r = 0.0
|
|
n_mc = 150
|
|
for _ in range(n_mc):
|
|
X = gen(64)
|
|
n = X.shape[0]
|
|
E = enc(X.reshape(-1, d)).reshape(n, U, d) if enc is not None else X
|
|
ch = channel(E, snr_db=float(snr))
|
|
Eh = model(ch["yI"], ch["yQ"])
|
|
# all metrics vs the SOURCE
|
|
cos = (Eh * X).sum(-1)
|
|
a += float((cos < 0.45).float().mean())
|
|
b += float(cos.mean())
|
|
# batch retrieval: nearest ENCODED gallery entry (collapse makes
|
|
# the gallery indistinguishable and drives accuracy to chance)
|
|
q = Eh.reshape(-1, d)
|
|
g = E.reshape(-1, d)
|
|
sim = q @ g.T
|
|
r += float((sim.argmax(1) == torch.arange(q.shape[0],
|
|
device=q.device))
|
|
.float().mean())
|
|
ss.append(a / n_mc); cc.append(b / n_mc); rr.append(r / n_mc)
|
|
return ss, cc, rr
|
|
|
|
|
|
out = {"snr": SNR_GRID.tolist(), "configs": {}}
|
|
for mode in ["frozen", "e2e_moving", "e2e_anchored", "e2e_vicreg"]:
|
|
model, enc = train(mode)
|
|
erank = collapse_metrics(enc)
|
|
ss, cc, rr = curves(model, enc)
|
|
out["configs"][mode] = {"ser": ss, "cos": cc, "retr": rr, "erank": erank}
|
|
print(f"[E8v2] {mode}: erank={erank:.1f} SER@10={ss[5]:.3f} "
|
|
f"cos@10={cc[5]:.3f} retr@10={rr[5]:.3f}", flush=True)
|
|
|
|
save_json("e8_v2_e2e.json", out)
|