Ki-Ho Lee
248e637f55
Semantic Multiplexing Gain in Wireless Systems via Expanded Embeddings: A BERT Case Study. Includes the shared library, all experiment scripts (training with and without SNR-aware MAML, the token-domain comparison, the K sweep, and DistilBERT), the replot script that regenerates every figure from the stored results, the supplementary probe-versus-cosine analysis, and the raw results behind every figure in the letter.
189 lines
7.8 KiB
Python
Executable File
189 lines
7.8 KiB
Python
Executable File
# =========================================================
|
|
# cl_maml_all.py — promote SNR-aware MAML to the default
|
|
# training procedure for all reported configurations.
|
|
#
|
|
# Adds MAML-trained counterparts of the load sweep, the
|
|
# conventional orthogonal scheme (fairness: both sides at
|
|
# their best), and the random-mask variant. The K sweep and
|
|
# DistilBERT stay joint-trained (sensitivity studies).
|
|
#
|
|
# Held-out evaluation identical to cl_experiments.py.
|
|
# Output: fig_cl/cl_results_maml.json (+ convergence CSV)
|
|
# =========================================================
|
|
|
|
import argparse, os, json, csv, time, random
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from bert_semcom import (
|
|
BertSemComMux, semantic_loss, split_params,
|
|
gather_inner_params, apply_inner_update
|
|
)
|
|
from cl_experiments import (
|
|
load_agnews_labeled, Extractor, SplitCache,
|
|
final_eval, train_probe, probe_accuracy, set_seed,
|
|
EVAL_SNRS, TRAIN_SNRS, measure_latency
|
|
)
|
|
|
|
|
|
def train_maml(cache, U, d_bert, K, device, freeze_masks=False,
|
|
epochs=200, steps=300, lr=1e-3, inner_lr=5e-4,
|
|
meta_batch=4, lam=0.5, channel="rayleigh",
|
|
conv_trials=50, label="", seed=42):
|
|
set_seed(seed)
|
|
model = BertSemComMux(U, d_bert, d_bert * K, 512).to(device)
|
|
if freeze_masks:
|
|
model.user_mask.weight.requires_grad_(False)
|
|
inner_keys = ["tx_proj", "user_query", "rx_proj"]
|
|
else:
|
|
inner_keys = ["tx_proj", "user_mask", "user_query", "rx_proj"]
|
|
opt = torch.optim.Adam(
|
|
[p for p in model.parameters() if p.requires_grad], lr=lr)
|
|
conv = []
|
|
t0 = time.time()
|
|
for ep in range(1, epochs + 1):
|
|
model.train()
|
|
loss_sum = 0.0
|
|
for _ in range(steps):
|
|
snr_tasks = [float(x) for x in
|
|
random.sample(TRAIN_SNRS, meta_batch)]
|
|
base_params = split_params(model)
|
|
meta_loss = torch.tensor(0.0, device=device)
|
|
for snr in snr_tasks:
|
|
fp = {k: {n: p for n, p in v.items()}
|
|
for k, v in base_params.items()}
|
|
b_s, _ = cache.sample_train(U, device)
|
|
bh = model(b_s, snr, channel, params=fp)
|
|
ls, _, _ = semantic_loss(b_s, bh, lam)
|
|
fl, mi = gather_inner_params(fp, inner_keys)
|
|
grads = torch.autograd.grad(ls, fl)
|
|
fp = apply_inner_update(fp, inner_keys, mi, grads,
|
|
inner_lr, "first")
|
|
b_q, _ = cache.sample_train(U, device)
|
|
bh_q = model(b_q, snr, channel, params=fp)
|
|
ql, _, _ = semantic_loss(b_q, bh_q, lam)
|
|
meta_loss = meta_loss + ql
|
|
meta_loss = meta_loss / float(meta_batch)
|
|
opt.zero_grad(set_to_none=True)
|
|
meta_loss.backward()
|
|
opt.step()
|
|
loss_sum += meta_loss.item()
|
|
model.eval()
|
|
with torch.no_grad():
|
|
c = 0.0
|
|
for _ in range(conv_trials):
|
|
b, _ = cache.sample_test(U, device)
|
|
bh = model(b, 10, channel)
|
|
c += F.cosine_similarity(bh, b, dim=-1).mean().item()
|
|
conv.append({"epoch": ep, "loss": loss_sum / steps,
|
|
"cos10": c / conv_trials})
|
|
if ep % 20 == 0 or ep == 1 or ep == epochs:
|
|
print(f" [{label} Ep {ep:03d}/{epochs}] "
|
|
f"loss={loss_sum/steps:.4f} "
|
|
f"cos@10dB(test)={c/conv_trials:.4f}", flush=True)
|
|
return model, conv, time.time() - t0
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--save-dir", default="fig_cl")
|
|
ap.add_argument("--epochs", type=int, default=200)
|
|
ap.add_argument("--steps", type=int, default=300)
|
|
ap.add_argument("--trials", type=int, default=500)
|
|
args = ap.parse_args()
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
os.makedirs(args.save_dir, exist_ok=True)
|
|
print(f"[Device: {device}]", flush=True)
|
|
train_items, test_items = load_agnews_labeled()
|
|
bert = Extractor("bert-base-uncased", device)
|
|
d_bert = bert.embed_dim
|
|
cache = SplitCache(bert, train_items, test_items)
|
|
del bert
|
|
torch.cuda.empty_cache()
|
|
|
|
probe, acc_clean = train_probe(cache, device)
|
|
R = {"probe_clean_acc": acc_clean}
|
|
PROBE_SNRS = list(EVAL_SNRS)
|
|
conv_rows = []
|
|
|
|
configs = [
|
|
("mamlB_U1_K1", dict(U=1, K=1)),
|
|
("mamlP_U1_K4", dict(U=1, K=4)),
|
|
("mamlP_U2_K4", dict(U=2, K=4)),
|
|
("mamlP_U3_K4", dict(U=3, K=4)),
|
|
("mamlP_U5_K4", dict(U=5, K=4)),
|
|
("mamlP_U6_K4", dict(U=6, K=4)),
|
|
("mamlR_U4_K4", dict(U=4, K=4, freeze_masks=True)),
|
|
]
|
|
masks_store = {}
|
|
for name, kw in configs:
|
|
print(f"\n=== {name} ===", flush=True)
|
|
U, K = kw.pop("U"), kw.pop("K")
|
|
model, conv, ttime = train_maml(
|
|
cache, U, d_bert, K, device,
|
|
epochs=args.epochs, steps=args.steps, label=name, **kw)
|
|
collect = PROBE_SNRS if name in (
|
|
"mamlB_U1_K1", "mamlR_U4_K4") else None
|
|
sweep, collected = final_eval(model, cache, U, device,
|
|
trials=args.trials,
|
|
collect_at=collect)
|
|
entry = {"U": U, "K": K, "train_s": ttime,
|
|
"snr": {str(s): sweep[s] for s in EVAL_SNRS}}
|
|
if collect:
|
|
entry["probe_acc"] = {str(s): a for s, a in
|
|
probe_accuracy(probe, collected,
|
|
cache, device).items()}
|
|
if name in ("mamlP_U2_K4", "mamlP_U3_K4"):
|
|
m = model.user_mask.weight.detach().cpu()
|
|
mn = F.normalize(m, p=2, dim=1)
|
|
masks_store[name] = (mn @ mn.T).numpy().tolist()
|
|
R[name] = entry
|
|
for c in conv:
|
|
conv_rows.append([name, c["epoch"], c["loss"], c["cos10"]])
|
|
R["mask_corr"] = masks_store
|
|
with open(os.path.join(args.save_dir,
|
|
"cl_results_maml.json"), "w") as f:
|
|
json.dump(R, f, indent=1)
|
|
|
|
# probe accuracies for the already-trained maml_U4_K4 are collected
|
|
# by re-training? No — retrain U=4 MAML for probe collection and
|
|
# mask correlation so every reported number comes from one protocol.
|
|
print("\n=== mamlP_U4_K4 (retrain for probe/masks) ===", flush=True)
|
|
model, conv, ttime = train_maml(cache, 4, d_bert, 4, device,
|
|
epochs=args.epochs, steps=args.steps,
|
|
label="mamlP_U4_K4")
|
|
sweep, collected = final_eval(model, cache, 4, device,
|
|
trials=args.trials,
|
|
collect_at=PROBE_SNRS)
|
|
entry = {"U": 4, "K": 4, "train_s": ttime,
|
|
"snr": {str(s): sweep[s] for s in EVAL_SNRS}}
|
|
entry["probe_acc"] = {str(s): a for s, a in
|
|
probe_accuracy(probe, collected,
|
|
cache, device).items()}
|
|
m = model.user_mask.weight.detach().cpu()
|
|
mn = F.normalize(m, p=2, dim=1)
|
|
masks_store["mamlP_U4_K4"] = (mn @ mn.T).numpy().tolist()
|
|
ms_gpu, n_params = measure_latency(model, cache, 4, device)
|
|
entry["lat_gpu_ms"] = ms_gpu
|
|
entry["params"] = n_params
|
|
R["mamlP_U4_K4"] = entry
|
|
R["mask_corr"] = masks_store
|
|
for c in conv:
|
|
conv_rows.append(["mamlP_U4_K4", c["epoch"], c["loss"],
|
|
c["cos10"]])
|
|
|
|
with open(os.path.join(args.save_dir, "cl_results_maml.json"),
|
|
"w") as f:
|
|
json.dump(R, f, indent=1)
|
|
with open(os.path.join(args.save_dir, "cl_convergence_maml.csv"),
|
|
"w", newline="") as f:
|
|
w = csv.writer(f)
|
|
w.writerow(["config", "epoch", "loss", "cos10_test"])
|
|
w.writerows(conv_rows)
|
|
print("\nAll MAML-default experiments complete.", flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|