Initial release: code for WCL2026-1544 (context-aware embedding masking via DRL)
This commit is contained in:
Ki-Ho Lee
@@ -0,0 +1,224 @@
|
||||
#!/usr/bin/env python3
|
||||
# ---------------------------------------------------------------
|
||||
# Task-oriented evaluation for the WCL revision.
|
||||
#
|
||||
# Adds two things requested by the reviewers:
|
||||
# (R1.6) a task-oriented metric beyond CosSim: top-1 semantic
|
||||
# retrieval accuracy over the C=8,000 AG News pool
|
||||
# (chance = 1/C), i.e. each reconstructed user embedding
|
||||
# is classified against the full sentence pool by cosine.
|
||||
# (R2.6) an INDEPENDENT baseline that does not share the proposed
|
||||
# mask-learning machinery: FIXED mutually-orthogonal masks
|
||||
# (non-learned signatures, O(M)=0 by construction) with a
|
||||
# transceiver trained to them. This is the classical
|
||||
# "orthogonal-signature" reference point.
|
||||
#
|
||||
# Methods (all at U=4, matched 100-epoch budget, seed 0):
|
||||
# - Proposed DRL : load existing checkpoint (trx+actor), det. mean.
|
||||
# - Fixed-Orth : fixed orthogonal masks + trained transceiver.
|
||||
# - Static (Sem) : free masks + transceiver, MSE+CosSim loss.
|
||||
# - Static (CE) : free masks + transceiver, symbol-CE loss.
|
||||
#
|
||||
# Outputs: ../results_sweeps/task_oriented/taskmetric_sweep.csv
|
||||
# ---------------------------------------------------------------
|
||||
import os, math, csv, random, argparse
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from drl_mask_policy import (MaskedTransceiver, MaskActor, apply_channel,
|
||||
encode_state, orthogonality_penalty)
|
||||
|
||||
SNRS = list(range(0, 31, 5))
|
||||
|
||||
|
||||
def set_seed(s):
|
||||
torch.manual_seed(s); np.random.seed(s); random.seed(s)
|
||||
|
||||
|
||||
def load_emb(path, d_bert, pool_size, device):
|
||||
if path and os.path.exists(path):
|
||||
emb = torch.load(path, map_location="cpu")
|
||||
emb = F.normalize(emb - emb.mean(dim=0, keepdim=True), p=2, dim=-1)
|
||||
print(f"[INFO] Loaded {tuple(emb.shape)} embeddings from {path}")
|
||||
else:
|
||||
emb = F.normalize(torch.randn(pool_size, d_bert), p=2, dim=-1)
|
||||
print("[INFO] Using synthetic embeddings.")
|
||||
return emb.to(device)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def sweep_metrics(trx, emb, mask_fn, device, U, eval_trials, eval_seed=12345,
|
||||
channel="rayleigh"):
|
||||
"""Return per-SNR (cossim, orthogonality, top1_accuracy).
|
||||
|
||||
mask_fn(snr_db) -> (U, d_s) mask matrix for that context.
|
||||
A fixed eval_seed makes the channel draws identical across
|
||||
methods, giving a paired comparison.
|
||||
"""
|
||||
trx.eval()
|
||||
C = emb.size(0)
|
||||
rows = []
|
||||
for snr_db in SNRS:
|
||||
set_seed(eval_seed) # identical channel/idx draws per method
|
||||
coss, orths, accs = [], [], []
|
||||
for _ in range(eval_trials):
|
||||
idx = torch.randint(0, C, (U,), device=device)
|
||||
b = emb[idx]
|
||||
m = mask_fn(snr_db)
|
||||
b_hat, _ = trx(b, m, snr_db, channel)
|
||||
coss.append(F.cosine_similarity(b_hat, b, dim=-1).mean().item())
|
||||
orths.append(orthogonality_penalty(m).item())
|
||||
# top-1 retrieval over the full pool
|
||||
bn = F.normalize(b_hat, p=2, dim=-1)
|
||||
sim = bn @ emb.T # (U, C)
|
||||
pred = sim.argmax(dim=-1) # (U,)
|
||||
accs.append((pred == idx).float().mean().item())
|
||||
rows.append((snr_db, float(np.mean(coss)), float(np.mean(orths)),
|
||||
float(np.mean(accs))))
|
||||
print(f" [SNR {snr_db:>2d}dB] CosSim={rows[-1][1]:.4f} "
|
||||
f"Orth={rows[-1][2]:.5f} Top1Acc={rows[-1][3]:.4f}")
|
||||
return rows
|
||||
|
||||
|
||||
def train_transceiver(emb, args, device, mode, fixed_masks=None):
|
||||
"""Train a transceiver (+ free masks unless fixed_masks given).
|
||||
mode in {'sem','ce','fixed'}. Returns (trx, mask_tensor_fn)."""
|
||||
set_seed(args.seed)
|
||||
U, d_bert, d_s = args.users, args.d_bert, args.d_bert * args.mux_factor
|
||||
C = emb.size(0)
|
||||
trx = MaskedTransceiver(U, d_bert, d_s).to(device)
|
||||
|
||||
if fixed_masks is not None:
|
||||
masks_param = None
|
||||
m_fixed = fixed_masks.to(device)
|
||||
params = list(trx.parameters())
|
||||
else:
|
||||
masks_param = nn.Embedding(U, d_s).to(device)
|
||||
nn.init.normal_(masks_param.weight, std=0.5)
|
||||
params = list(trx.parameters()) + list(masks_param.parameters())
|
||||
opt = torch.optim.Adam(params, lr=args.lr)
|
||||
ar = torch.arange(U, device=device)
|
||||
tau = float(args.ce_tau)
|
||||
|
||||
for ep in range(1, args.epochs + 1):
|
||||
for _ in range(args.steps_per_epoch):
|
||||
snr_db = random.choice(args.train_snr)
|
||||
idx = torch.randint(0, C, (U,), device=device)
|
||||
b = emb[idx]
|
||||
m = m_fixed if fixed_masks is not None else masks_param(ar)
|
||||
b_hat, _ = trx(b, m, snr_db, args.channel)
|
||||
if mode == "ce":
|
||||
b_hat_n = F.normalize(b_hat, p=2, dim=-1)
|
||||
loss = F.cross_entropy(tau * (b_hat_n @ emb.T), idx)
|
||||
else: # 'sem' and 'fixed' both use MSE+CosSim
|
||||
cos = F.cosine_similarity(b_hat, b, dim=-1).mean()
|
||||
loss = 0.5 * F.mse_loss(b_hat, b) + 0.5 * (1.0 - cos)
|
||||
opt.zero_grad(set_to_none=True)
|
||||
loss.backward()
|
||||
nn.utils.clip_grad_norm_(params, 1.0)
|
||||
opt.step()
|
||||
if ep % 20 == 0 or ep == 1:
|
||||
print(f" [{mode} ep {ep:03d}/{args.epochs}] loss={loss.item():.4f}")
|
||||
|
||||
if fixed_masks is not None:
|
||||
return trx, (lambda snr: m_fixed)
|
||||
masks_const = masks_param(ar).detach()
|
||||
return trx, (lambda snr: masks_const)
|
||||
|
||||
|
||||
def build_fixed_orthogonal_masks(U, d_s, seed=0):
|
||||
"""U mutually-orthogonal unit-RMS signature masks (O(M)=0)."""
|
||||
g = torch.Generator().manual_seed(seed)
|
||||
A = torch.randn(d_s, U, generator=g)
|
||||
Q, _ = torch.linalg.qr(A) # (d_s, U) orthonormal columns
|
||||
M = Q.T # (U, d_s), orthonormal rows
|
||||
M = M * math.sqrt(d_s) # scale to unit RMS per row
|
||||
return M
|
||||
|
||||
|
||||
def load_drl(ckpt_path, args, device):
|
||||
ck = torch.load(ckpt_path, map_location=device, weights_only=False)
|
||||
U, d_bert, d_s = args.users, args.d_bert, args.d_bert * args.mux_factor
|
||||
trx = MaskedTransceiver(U, d_bert, d_s).to(device)
|
||||
trx.load_state_dict(ck["trx"])
|
||||
actor = MaskActor(U, d_s, hidden=args.hidden, rank=args.rank).to(device)
|
||||
actor.load_state_dict(ck["actor"])
|
||||
actor.eval()
|
||||
|
||||
@torch.no_grad()
|
||||
def mask_fn(snr_db):
|
||||
st = encode_state(snr_db, U, args.users_max, device=device)
|
||||
mean, _ = actor(st)
|
||||
return mean
|
||||
return trx, mask_fn
|
||||
|
||||
|
||||
def main():
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--users", type=int, default=4)
|
||||
p.add_argument("--users-max", type=int, default=8)
|
||||
p.add_argument("--mux-factor", type=int, default=4)
|
||||
p.add_argument("--d-bert", type=int, default=768)
|
||||
p.add_argument("--hidden", type=int, default=256)
|
||||
p.add_argument("--rank", type=int, default=64)
|
||||
p.add_argument("--embed-file", type=str, default="../bert_agnews_8000.pt")
|
||||
p.add_argument("--pool-size", type=int, default=8000)
|
||||
p.add_argument("--channel", default="rayleigh")
|
||||
p.add_argument("--train-snr", type=float, nargs="+",
|
||||
default=[0, 5, 10, 15, 20, 25])
|
||||
p.add_argument("--lr", type=float, default=1e-3)
|
||||
p.add_argument("--ce-tau", type=float, default=16.0)
|
||||
p.add_argument("--epochs", type=int, default=100)
|
||||
p.add_argument("--steps-per-epoch", type=int, default=200)
|
||||
p.add_argument("--eval-trials", type=int, default=200)
|
||||
p.add_argument("--seed", type=int, default=0)
|
||||
p.add_argument("--drl-ckpt", type=str,
|
||||
default="../results_sweeps/drl_U4_100ep_s0/drl_mask_policy.pt")
|
||||
p.add_argument("--out-dir", type=str,
|
||||
default="../results_sweeps/task_oriented")
|
||||
args = p.parse_args()
|
||||
|
||||
device = torch.device("mps" if torch.backends.mps.is_available()
|
||||
else ("cuda" if torch.cuda.is_available() else "cpu"))
|
||||
print(f"[INFO] device={device}")
|
||||
U, d_s = args.users, args.d_bert * args.mux_factor
|
||||
emb = load_emb(args.embed_file, args.d_bert, args.pool_size, device)
|
||||
|
||||
results = {} # method -> list of rows
|
||||
|
||||
print("\n=== Proposed DRL (from checkpoint) ===")
|
||||
trx, mfn = load_drl(args.drl_ckpt, args, device)
|
||||
results["proposed_drl"] = sweep_metrics(trx, emb, mfn, device, U,
|
||||
args.eval_trials)
|
||||
|
||||
print("\n=== Fixed orthogonal masks (independent baseline) ===")
|
||||
Mfix = build_fixed_orthogonal_masks(U, d_s, seed=0)
|
||||
trx, mfn = train_transceiver(emb, args, device, "fixed", fixed_masks=Mfix)
|
||||
results["fixed_orth"] = sweep_metrics(trx, emb, mfn, device, U,
|
||||
args.eval_trials)
|
||||
|
||||
print("\n=== Static masking + semantic loss ===")
|
||||
trx, mfn = train_transceiver(emb, args, device, "sem")
|
||||
results["static_sem"] = sweep_metrics(trx, emb, mfn, device, U,
|
||||
args.eval_trials)
|
||||
|
||||
print("\n=== Static masking + CE loss ===")
|
||||
trx, mfn = train_transceiver(emb, args, device, "ce")
|
||||
results["static_ce"] = sweep_metrics(trx, emb, mfn, device, U,
|
||||
args.eval_trials)
|
||||
|
||||
os.makedirs(args.out_dir, exist_ok=True)
|
||||
out = os.path.join(args.out_dir, "taskmetric_sweep.csv")
|
||||
with open(out, "w", newline="") as f:
|
||||
w = csv.writer(f)
|
||||
w.writerow(["method", "snr_db", "cos_sim", "orthogonality", "top1_acc"])
|
||||
for meth, rows in results.items():
|
||||
for r in rows:
|
||||
w.writerow([meth, *r])
|
||||
print(f"\n[DONE] wrote {out}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user