v2 design: independent masks + affinity-aware Wiener demultiplexer
Redesign after the independent-mask dominance finding: the affinity now parameterizes the receiver (closed-form Wiener) instead of the mask ensemble. New Theorem 1 (spectral closed form), floors sqrt(1-b^2)/2 vs 1/2, full-cooperation bound with equality at b=1. GPU (torch) Monte Carlo backend, decision-directed SIC baseline, TikZ block diagram source, verification suite V1-V11.
This commit is contained in:
KiHoLee
+140
-119
@@ -1,23 +1,23 @@
|
||||
"""
|
||||
Capacity-matched EDMA refinement (parameter budget equal to the
|
||||
attention scheme: 4 d^2 = 2.36M at d = 768).
|
||||
Refinement stage for the v2 (affinity-aware Wiener) EDMA receiver.
|
||||
================================================================
|
||||
Four-head averaged gated refinement applied to the closed-form
|
||||
demultiplexer output:
|
||||
Trains the single-gate refinement operator
|
||||
|
||||
out = (1/4) sum_k D softmax(Q_k x / sqrt(D)) .* x,
|
||||
out = D softmax(W z / sqrt(D)) .* z, z = Re(e_hat),
|
||||
|
||||
with Q_1..Q_4 in R^{D x D} (4 d^2 parameters, exactly the
|
||||
attention scheme's budget). The single-gate 0.59M refiner is the
|
||||
special case of four identical heads, so the family contains it
|
||||
by construction. Same training recipe: demux outputs from
|
||||
parametric pairs at beta = 0.028, Haar pool 32, Rayleigh
|
||||
channels, complex noise, training SNR uniform in [5, 25] dB,
|
||||
Adam 5e-4 with gradient clipping, batch 48, 200 epochs.
|
||||
(0.59M parameters at d = 768) on aware-demultiplexer outputs, then
|
||||
warm-starts a capacity-check variant with four heads (4 d^2 = 2.36M)
|
||||
and fine-tunes it, so the family contains the single gate by
|
||||
construction. Training data: parametric pairs at beta = 0.028, a
|
||||
fixed pool of 32 independent Haar mask pairs, Rayleigh channels,
|
||||
complex noise, training SNR uniform in [5, 25] dB, Adam 5e-4 with
|
||||
gradient clipping, batch 48, 220 epochs (stage 2 from epoch 120 at
|
||||
lr 2e-4).
|
||||
|
||||
Evaluation on the real BERT/ViT pairs with fresh Haar masks and
|
||||
200 fading draws per pair. Appends column `edma_ref2` to
|
||||
data/bertvit_merged.csv and prints all-curve numbers.
|
||||
NFADE fading draws per pair. Appends columns `edma_ref` (single
|
||||
gate) and `edma_ref2` (four heads) to data/bertvit_merged.csv.
|
||||
Requires torch (run under WSL with CUDA if available).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
import csv
|
||||
@@ -26,161 +26,182 @@ import time
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from fig_real_merged import load_pairs, cosine, SNRS, NFADE, D, DATA
|
||||
from fig_real_merged import load_pairs, SNRS, NFADE, D, DATA
|
||||
|
||||
SEED = 2026
|
||||
rng = np.random.default_rng(SEED + 31)
|
||||
torch.manual_seed(SEED + 31)
|
||||
DEV = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
BETA0 = 0.028
|
||||
G0 = 1.0 - BETA0**2
|
||||
print(f"[refine] device = {DEV}")
|
||||
|
||||
|
||||
def haar_t(gen):
|
||||
Q, R = torch.linalg.qr(torch.randn(D, D, generator=gen))
|
||||
return Q * torch.sign(torch.diagonal(R))
|
||||
def haar_t(n, gen):
|
||||
G = torch.randn(n, D, D, generator=gen, device=DEV)
|
||||
Q, R = torch.linalg.qr(G)
|
||||
return Q * torch.sign(torch.diagonal(R, dim1=-2, dim2=-1)).unsqueeze(-2)
|
||||
|
||||
|
||||
def train_refiner2(epochs=220, steps=20, batch=48, lr=5e-4,
|
||||
l2=0.5, l3=0.5, pool=32):
|
||||
"""Stage 1 trains a single gate (the proven 0.59M recipe); stage 2
|
||||
warm-starts four heads from it plus small perturbations and
|
||||
fine-tunes at a reduced learning rate, so the capacity-matched
|
||||
family starts at the single-gate solution it contains."""
|
||||
print(f"=== training capacity-matched refinement (4-head gate, "
|
||||
f"4d^2 = {4*D*D/1e6:.2f}M params, warm-started) ===",
|
||||
flush=True)
|
||||
gen = torch.Generator().manual_seed(SEED + 31)
|
||||
masks = []
|
||||
for _ in range(pool):
|
||||
U1, U2 = haar_t(gen), haar_t(gen)
|
||||
masks.append((U1.numpy(), (BETA0 * U1
|
||||
+ math.sqrt(G0) * U2).numpy()))
|
||||
Q0 = torch.nn.Parameter(torch.randn(D, D, generator=gen)
|
||||
/ math.sqrt(D))
|
||||
params = [Q0]
|
||||
def aware_t(t1, Q, beta, c1, nvar):
|
||||
"""Batched affinity-aware Wiener demux in torch (complex)."""
|
||||
b = t1.shape[0]
|
||||
g = 1.0 - beta * beta
|
||||
rho = g * (c1.abs()**2) / D + nvar # (b,)
|
||||
A = torch.eye(D, device=DEV, dtype=torch.cfloat).expand(b, D, D) \
|
||||
+ beta * c1.view(b, 1, 1) * Q.to(torch.cfloat)
|
||||
S = A @ A.mH / D + rho.view(b, 1, 1) \
|
||||
* torch.eye(D, device=DEV, dtype=torch.cfloat)
|
||||
x = torch.linalg.solve(S, t1.unsqueeze(-1))
|
||||
return (A.mH @ x).squeeze(-1) / D
|
||||
|
||||
|
||||
def train_batch(masks, Qs, gen, batch):
|
||||
"""Generate one training batch of aware-demux outputs (user 1)."""
|
||||
e1 = torch.nn.functional.normalize(
|
||||
torch.randn(batch, D, generator=gen, device=DEV), dim=1)
|
||||
w = torch.randn(batch, D, generator=gen, device=DEV)
|
||||
w = w - (w * e1).sum(1, keepdim=True) * e1
|
||||
w = torch.nn.functional.normalize(w, dim=1)
|
||||
e2 = BETA0 * e1 + math.sqrt(G0) * w
|
||||
sel = torch.randint(len(masks), (batch,), generator=gen, device=DEV)
|
||||
M1 = masks[0][sel]; M2 = masks[1][sel]; Q = Qs[sel]
|
||||
snr = 5.0 + 20.0 * torch.rand(batch, generator=gen, device=DEV)
|
||||
sig = 10 ** (-snr / 20.0)
|
||||
h = (torch.randn(batch, 2, generator=gen, device=DEV)
|
||||
+ 1j * torch.randn(batch, 2, generator=gen, device=DEV)) \
|
||||
/ math.sqrt(2)
|
||||
n = (torch.randn(batch, D, generator=gen, device=DEV)
|
||||
+ 1j * torch.randn(batch, D, generator=gen, device=DEV)) \
|
||||
/ math.sqrt(2)
|
||||
r = h[:, :1] * (M1 @ e1.unsqueeze(-1)).squeeze(-1).to(torch.cfloat) \
|
||||
+ h[:, 1:2] * (M2 @ e2.unsqueeze(-1)).squeeze(-1).to(torch.cfloat) \
|
||||
+ sig.view(-1, 1) * n
|
||||
t1 = (M1.transpose(-1, -2).to(torch.cfloat)
|
||||
@ r.unsqueeze(-1)).squeeze(-1) / h[:, :1]
|
||||
c1 = h[:, 1] / h[:, 0]
|
||||
nvar = sig**2 / h[:, 0].abs()**2
|
||||
g1 = aware_t(t1, Q, BETA0, c1, nvar)
|
||||
return g1.real.float(), e1
|
||||
|
||||
|
||||
def train_refiners(epochs=220, steps=20, batch=48, lr=5e-4,
|
||||
l2=0.5, l3=0.5, pool=32, stage2_at=120):
|
||||
print(f"=== training refinement (single gate {D*D/1e6:.2f}M, "
|
||||
f"then 4-head warm start {4*D*D/1e6:.2f}M) ===", flush=True)
|
||||
gen = torch.Generator(device=DEV).manual_seed(SEED + 31)
|
||||
U1 = haar_t(pool, gen); U2 = haar_t(pool, gen)
|
||||
masks = (U1, U2)
|
||||
Qs = U1.transpose(-1, -2) @ U2
|
||||
params = [torch.nn.Parameter(
|
||||
torch.randn(D, D, generator=gen, device=DEV) / math.sqrt(D))]
|
||||
opt = torch.optim.Adam(params, lr=lr)
|
||||
stage2_at = 120 # epochs of single-gate pre-training
|
||||
P_single = None
|
||||
|
||||
def forward(x):
|
||||
def forward(x, ps):
|
||||
outs = [D * torch.softmax((x @ Qk.T) / math.sqrt(D), dim=1) * x
|
||||
for Qk in params]
|
||||
return sum(outs) / len(params)
|
||||
for Qk in ps]
|
||||
return sum(outs) / len(ps)
|
||||
|
||||
t0 = time.time()
|
||||
for ep in range(epochs):
|
||||
if ep == stage2_at:
|
||||
P_single = params[0].detach().clone()
|
||||
base = params[0].detach()
|
||||
params = [torch.nn.Parameter(
|
||||
base.clone() + 0.02 * torch.randn(D, D, generator=gen)
|
||||
base.clone() + 0.02 * torch.randn(D, D, generator=gen,
|
||||
device=DEV)
|
||||
/ math.sqrt(D)) for _ in range(4)]
|
||||
opt = torch.optim.Adam(params, lr=2e-4)
|
||||
print(f" [warm start] 4 heads initialised from the trained "
|
||||
f"gate at epoch {ep}", flush=True)
|
||||
print(f" [warm start] 4 heads at epoch {ep}", flush=True)
|
||||
for _ in range(steps):
|
||||
xs, ts = [], []
|
||||
for _ in range(batch):
|
||||
e1 = torch.nn.functional.normalize(
|
||||
torch.randn(D, generator=gen), dim=0).numpy()
|
||||
w = torch.randn(D, generator=gen).numpy()
|
||||
w = w - (w @ e1) * e1
|
||||
w = w / np.linalg.norm(w)
|
||||
e2 = BETA0 * e1 + math.sqrt(G0) * w
|
||||
M1, M2 = masks[int(torch.randint(pool, (1,),
|
||||
generator=gen))]
|
||||
snr = float(5.0 + 20.0 * torch.rand(1, generator=gen))
|
||||
sig = 10 ** (-snr / 20.0)
|
||||
h = (torch.randn(2, generator=gen).numpy()
|
||||
+ 1j * torch.randn(2, generator=gen).numpy()) \
|
||||
/ math.sqrt(2)
|
||||
nc = (torch.randn(D, generator=gen).numpy()
|
||||
+ 1j * torch.randn(D, generator=gen).numpy()) \
|
||||
/ math.sqrt(2)
|
||||
rc = h[0] * (M1 @ e1) + h[1] * (M2 @ e2) + sig * nc
|
||||
t1 = M1.T @ rc / h[0]
|
||||
t2 = M2.T @ rc / h[1]
|
||||
g1 = (t1 - BETA0 * (h[1] / h[0]) * t2) / G0
|
||||
xs.append(torch.tensor(np.real(g1), dtype=torch.float32))
|
||||
ts.append(torch.tensor(e1, dtype=torch.float32))
|
||||
x = torch.stack(xs); t = torch.stack(ts)
|
||||
out = forward(x)
|
||||
mse = ((out - t)**2).mean()
|
||||
cs = torch.nn.functional.cosine_similarity(out, t, dim=1).mean()
|
||||
with torch.no_grad():
|
||||
x, tgt = train_batch(masks, Qs, gen, batch)
|
||||
out = forward(x, params)
|
||||
mse = ((out - tgt)**2).mean()
|
||||
cs = torch.nn.functional.cosine_similarity(out, tgt, dim=1).mean()
|
||||
loss = l2 * mse + l3 * (1.0 - cs)
|
||||
opt.zero_grad(); loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(params, 1.0)
|
||||
opt.step()
|
||||
if (ep + 1) % 50 == 0:
|
||||
if (ep + 1) % 40 == 0:
|
||||
print(f" epoch {ep+1}: loss {float(loss.detach()):.4f} "
|
||||
f"(cos {float(cs.detach()):.3f})", flush=True)
|
||||
print(f" trained in {time.time()-t0:.0f}s")
|
||||
return [p.detach().numpy() for p in params]
|
||||
return P_single, [p.detach() for p in params]
|
||||
|
||||
|
||||
def refine2(P, g):
|
||||
x = np.real(g)
|
||||
|
||||
def gate(Q, v):
|
||||
sc = (Q @ v) / math.sqrt(D)
|
||||
sc = sc - sc.max()
|
||||
w = np.exp(sc); w /= w.sum()
|
||||
return D * w * v
|
||||
|
||||
return sum(gate(Qk, x) for Qk in P) / 4.0
|
||||
def refine_apply(ps, z):
|
||||
"""z: (b, D) real torch tensor; ps: list of gates."""
|
||||
outs = [D * torch.softmax((z @ Qk.T) / math.sqrt(D), dim=1) * z
|
||||
for Qk in ps]
|
||||
return sum(outs) / len(ps)
|
||||
|
||||
|
||||
def main():
|
||||
A, B, betas = load_pairs()
|
||||
P = train_refiner2()
|
||||
ref = np.zeros(len(SNRS)); cnt = 0
|
||||
P1, P4 = train_refiners()
|
||||
gen = torch.Generator(device=DEV).manual_seed(SEED + 77)
|
||||
ref1 = np.zeros(len(SNRS)); ref4 = np.zeros(len(SNRS)); cnt = 0
|
||||
t0 = time.time()
|
||||
At = torch.tensor(A, dtype=torch.float32, device=DEV)
|
||||
Bt = torch.tensor(B, dtype=torch.float32, device=DEV)
|
||||
for i in range(len(A)):
|
||||
e1, e2, bi = A[i], B[i], float(betas[i])
|
||||
gi = 1.0 - bi**2
|
||||
bi = float(betas[i])
|
||||
e1 = At[i]; e2 = Bt[i]
|
||||
for f in range(NFADE):
|
||||
G1 = rng.standard_normal((D, D))
|
||||
Qh, Rh = np.linalg.qr(G1)
|
||||
U1 = Qh * np.sign(np.diag(Rh))
|
||||
G2 = rng.standard_normal((D, D))
|
||||
Qh, Rh = np.linalg.qr(G2)
|
||||
U2 = Qh * np.sign(np.diag(Rh))
|
||||
M1 = U1
|
||||
M2 = bi * U1 + math.sqrt(gi) * U2
|
||||
h = (rng.standard_normal(2) + 1j * rng.standard_normal(2)) \
|
||||
M = haar_t(2, gen)
|
||||
M1, M2 = M[0], M[1]
|
||||
Q = M1.T @ M2
|
||||
h = (torch.randn(2, generator=gen, device=DEV)
|
||||
+ 1j * torch.randn(2, generator=gen, device=DEV)) \
|
||||
/ math.sqrt(2)
|
||||
h1, h2 = h
|
||||
r0 = h1 * (M1 @ e1) + h2 * (M2 @ e2)
|
||||
n = (rng.standard_normal(D) + 1j * rng.standard_normal(D)) \
|
||||
n = (torch.randn(D, generator=gen, device=DEV)
|
||||
+ 1j * torch.randn(D, generator=gen, device=DEV)) \
|
||||
/ math.sqrt(2)
|
||||
for k, s in enumerate(SNRS):
|
||||
sig = 10 ** (-s / 20.0)
|
||||
r = r0 + sig * n
|
||||
t1 = M1.T @ r / h1; t2 = M2.T @ r / h2
|
||||
g1 = (t1 - bi * (h2 / h1) * t2) / gi
|
||||
g2 = (t2 - bi * (h1 / h2) * t1) / gi
|
||||
ref[k] += 0.5 * (cosine(refine2(P, g1), e1)
|
||||
+ cosine(refine2(P, g2), e2))
|
||||
r0 = h[0] * (M1 @ e1).to(torch.cfloat) \
|
||||
+ h[1] * (M2 @ e2).to(torch.cfloat)
|
||||
sigs = torch.tensor(10 ** (-SNRS / 20.0), device=DEV,
|
||||
dtype=torch.float32)
|
||||
nb = len(SNRS)
|
||||
r = r0.unsqueeze(0) + sigs.view(-1, 1) * n.unsqueeze(0)
|
||||
t1 = (M1.T.to(torch.cfloat) @ r.unsqueeze(-1)).squeeze(-1) / h[0]
|
||||
t2 = (M2.T.to(torch.cfloat) @ r.unsqueeze(-1)).squeeze(-1) / h[1]
|
||||
c1 = (h[1] / h[0]).expand(nb)
|
||||
c2 = (h[0] / h[1]).expand(nb)
|
||||
v1 = (sigs**2 / h[0].abs()**2)
|
||||
v2 = (sigs**2 / h[1].abs()**2)
|
||||
g1 = aware_t(t1, Q.expand(nb, D, D), bi, c1, v1).real.float()
|
||||
g2 = aware_t(t2, Q.T.expand(nb, D, D), bi, c2, v2).real.float()
|
||||
with torch.no_grad():
|
||||
for P, acc in ((([P1]), ref1), ((P4), ref4)):
|
||||
o1 = refine_apply(P, g1)
|
||||
o2 = refine_apply(P, g2)
|
||||
cs1 = torch.nn.functional.cosine_similarity(
|
||||
o1, e1.unsqueeze(0), dim=1).abs()
|
||||
cs2 = torch.nn.functional.cosine_similarity(
|
||||
o2, e2.unsqueeze(0), dim=1).abs()
|
||||
acc += (0.5 * (cs1 + cs2)).cpu().numpy()
|
||||
cnt += 1
|
||||
print(f" pair {i+1}/{len(A)} done ({time.time()-t0:.0f}s)",
|
||||
flush=True)
|
||||
ref /= cnt
|
||||
ref1 /= cnt; ref4 /= cnt
|
||||
|
||||
rows = list(csv.DictReader(open(DATA / "bertvit_merged.csv")))
|
||||
names = list(rows[0].keys())
|
||||
if "edma_ref2" not in names:
|
||||
names.append("edma_ref2")
|
||||
for col in ("edma_ref", "edma_ref2"):
|
||||
if col not in names:
|
||||
names.append(col)
|
||||
for k, r in enumerate(rows):
|
||||
r["edma_ref2"] = f"{ref[k]}"
|
||||
r["edma_ref"] = f"{ref1[k]}"
|
||||
r["edma_ref2"] = f"{ref4[k]}"
|
||||
with open(DATA / "bertvit_merged.csv", "w", newline="") as f:
|
||||
w = csv.DictWriter(f, fieldnames=names)
|
||||
w.writeheader(); w.writerows(rows)
|
||||
print("[OK] appended edma_ref2 to bertvit_merged.csv")
|
||||
print("[OK] appended edma_ref / edma_ref2 to bertvit_merged.csv")
|
||||
for k, r in enumerate(rows):
|
||||
print(f" {float(r['snr_db']):4.0f} dB "
|
||||
f"EDMA {float(r['edma']):.3f} "
|
||||
f"ref(0.59M) {float(r['edma_ref']):.3f} "
|
||||
f"ref2(2.36M) {ref[k]:.3f} "
|
||||
f"ATT(2.36M) {float(r['att']):.3f} "
|
||||
f"genie {float(r['genie']):.3f}")
|
||||
print(f" {float(r['snr_db']):4.1f} dB "
|
||||
f"EDMA {float(r['edma']):.3f} ref {ref1[k]:.3f} "
|
||||
f"ref2 {ref4[k]:.3f} genie {float(r['genie']):.3f}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user