126 lines
4.6 KiB
Python
126 lines
4.6 KiB
Python
"""E5: comparison with a trained user-wise attention receiver.
|
|
|
|
A learned receiver representative of the end-to-end line (per-user query
|
|
attention over the matched-filter outputs, residual skip, unit
|
|
normalization) is trained at the operating SNR over uniformly random
|
|
affinities, then compared against the closed-form LMMSE and DR receivers
|
|
across the affinity sweep. The learned receiver receives no affinity side
|
|
information and must infer the coupling from data, which is the standard
|
|
setting of the learned line.
|
|
|
|
Outputs: data/e5_learned.csv, data/e5_train_log.csv
|
|
(figures come from replot_all.py only)
|
|
"""
|
|
import math
|
|
import os
|
|
import numpy as np
|
|
import torch
|
|
|
|
from semantic_mac import (affinity_matrix, matched_filter,
|
|
sample_latents_isotropic, demux_lmmse, demux_dr,
|
|
demux_sr, demux_sc, metrics)
|
|
|
|
HERE = os.path.dirname(os.path.abspath(__file__))
|
|
FIG = os.path.join(HERE, "..", "fig")
|
|
DATA = os.path.join(HERE, "..", "data")
|
|
|
|
U, D, DC = 4, 64, 16
|
|
SNR_DB = 10
|
|
RHO = 10 ** (SNR_DB / 10)
|
|
STEPS = 3000
|
|
BATCH = 64
|
|
|
|
|
|
class UserWiseAttention(torch.nn.Module):
|
|
"""Per-user query attention over the U matched-filter outputs."""
|
|
|
|
def __init__(self, U, d, dk=16, heads=4):
|
|
super().__init__()
|
|
self.U, self.d, self.dk, self.H = U, d, dk, heads
|
|
self.WK = torch.nn.Linear(d, dk * heads, bias=False)
|
|
self.WV = torch.nn.Linear(d, dk * heads, bias=False)
|
|
self.WO = torch.nn.Linear(dk * heads, d, bias=False)
|
|
self.q = torch.nn.Parameter(torch.randn(U, heads, dk) * 0.1)
|
|
self.log_eta = torch.nn.Parameter(torch.zeros(()))
|
|
|
|
def forward(self, tilde):
|
|
B, Uu, d = tilde.shape
|
|
K = self.WK(tilde).view(B, Uu, self.H, self.dk)
|
|
V = self.WV(tilde).view(B, Uu, self.H, self.dk)
|
|
sc = torch.einsum('uhk,bihk->buih', self.q, K) / math.sqrt(self.dk)
|
|
alpha = torch.softmax(torch.exp(self.log_eta) * sc, dim=2)
|
|
ctx = torch.einsum('buih,bihk->buhk', alpha, V).reshape(B, Uu, -1)
|
|
out = self.WO(ctx) + tilde
|
|
return out / (out.norm(dim=2, keepdim=True) + 1e-12)
|
|
|
|
|
|
def gen_batch(rng, batch, beta):
|
|
a = math.sqrt(beta) * np.ones(U)
|
|
B = affinity_matrix(a)
|
|
z = sample_latents_isotropic(batch, U, D, DC, a, rng)
|
|
tilde, h = matched_filter(z, B, RHO, rng)
|
|
return z, tilde, h, a, B
|
|
|
|
|
|
def main():
|
|
torch.manual_seed(0)
|
|
rng = np.random.default_rng(21)
|
|
net = UserWiseAttention(U, D)
|
|
n_par = sum(p.numel() for p in net.parameters())
|
|
print(f"learned receiver parameters: {n_par}")
|
|
opt = torch.optim.Adam(net.parameters(), lr=1e-3)
|
|
train_log = []
|
|
for it in range(STEPS):
|
|
beta = float(rng.uniform(0.05, 0.9))
|
|
z, tilde, h, a, B = gen_batch(rng, BATCH, beta)
|
|
zt = torch.tensor(z, dtype=torch.float32)
|
|
tt = torch.tensor(tilde, dtype=torch.float32)
|
|
out = net(tt)
|
|
loss = (1.0 - (out * zt).sum(-1)).mean()
|
|
opt.zero_grad()
|
|
loss.backward()
|
|
opt.step()
|
|
if (it + 1) % 50 == 0:
|
|
train_log.append((it + 1, float(loss.item())))
|
|
if (it + 1) % 500 == 0:
|
|
print(f" step {it+1}: loss={loss.item():.4f}")
|
|
# convergence evidence for the fixed training budget
|
|
with open(os.path.join(DATA, "e5_train_log.csv"), "w") as f:
|
|
f.write("step,loss\n")
|
|
for st, lo in train_log:
|
|
f.write(f"{st},{lo}\n")
|
|
|
|
# evaluation across the affinity sweep
|
|
betas = [0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
|
|
Vc = np.eye(D)[:, :DC]
|
|
rows = []
|
|
net.eval()
|
|
rng_e = np.random.default_rng(500)
|
|
for beta in betas:
|
|
z, tilde, h, a, B = gen_batch(rng_e, 4000, beta)
|
|
with torch.no_grad():
|
|
out_l = net(torch.tensor(tilde, dtype=torch.float32)).numpy()
|
|
r = {}
|
|
r["Learned"] = metrics(out_l.astype(np.float64), z)
|
|
lm, _ = demux_lmmse(tilde, B, h, RHO)
|
|
r["LMMSE"] = metrics(lm, z)
|
|
r["DR"] = metrics(demux_dr(tilde, B, h, RHO, a, Vc), z)
|
|
rows.append(r)
|
|
print(f"beta={beta:.2f} " + " ".join(
|
|
f"{k}:cos={v[0]:.3f},ser={v[2]:.3f}" for k, v in r.items()))
|
|
with open(os.path.join(DATA, "e5_learned.csv"), "w") as f:
|
|
keys = ["Learned", "LMMSE", "DR"]
|
|
f.write("beta," + ",".join(f"{k}_cos,{k}_nmse,{k}_ser" for k in keys)
|
|
+ "\n")
|
|
for b, r in zip(betas, rows):
|
|
f.write(f"{b}," + ",".join(
|
|
f"{r[k][0]},{r[k][1]},{r[k][2]}" for k in keys) + "\n")
|
|
|
|
# figures are produced only by the canonical replot_all.py (uniform
|
|
# geometry); experiment scripts write CSVs exclusively.
|
|
print("E5 done. Run replot_all.py to regenerate the figures.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|