Context-Aware Embedding Masking for Shared-Embedding Semantic Multiplexing

Source code for the IEEE Wireless Communications Letters paper

Context-Aware Embedding Masking Based on Reinforcement Learning for Semantic Multiplexing (WCL2026-1544) K.-H. Lee, H.-H. Choi, and J.-R. Lee.

A proximal policy optimization (PPO) actor, conditioned on the wireless context (SNR, U), emits the user mask matrix of the shared-embedding (SE) transceiver through a low-rank generator, under a reward that maximizes per-user cosine similarity and penalizes mask non-orthogonality.

Requirements

pip install -r requirements.txt

Python 3.8+ with torch, transformers, numpy, matplotlib.

Files

File	Role
drl_mask_policy.py	Main trainer. Modes: drl (proposed), joint (static masking, MSE+CosSim), joint_ce (cross-entropy), fixed_orth (fixed-orthogonal scheme). Implements the PPO actor/critic, low-rank generator, and SE transceiver.
extract_bert_embeddings.py	Produces the frozen bert-base-uncased embeddings of 8,000 AG News headlines (bert_agnews_8000.pt).
eval_task_oriented.py	Top-1 semantic retrieval accuracy / semantic SER (single run) — Table II.
eval_multiseed.py	Six-seed aggregation (mean ± std) of the retrieval metric — Table II.
fixed_orth_byU.py	Fixed-orthogonal-mask scheme swept over the user count U.
plot_drl_wcl.py	Regenerates all figures (Figs. 2–3) from the result CSVs.
run_*.sh	Experiment drivers (multi-seed training, SNR sweeps, ablations).

Mapping to the paper

Paper artifact	How to reproduce
Fig. 2(a) training-time O(M)	run_multiseed_all.sh then plot_drl_wcl.py
Fig. 2(b) ablation (beta, rank r)	run_beta02_U4.sh, run_ablation_U26.sh
Fig. 3(a) per-user CosSim vs. SNR	run_multiseed_100ep.sh
Fig. 3(b) throughput vs. U	run_ablation_U26.sh
Table II top-1 retrieval (6 seeds)	eval_multiseed.py
Fixed-orthogonal reference	fixed_orth_byU.py

Quick start

pip install -r requirements.txt
python extract_bert_embeddings.py     # -> bert_agnews_8000.pt
bash   run_multiseed_all.sh           # train all methods over six seeds
python eval_multiseed.py              # -> Table II (retrieval accuracy)
python plot_drl_wcl.py                # -> figures in fig/

Hyperparameters match Table I of the paper (PPO clip 0.2, four optimizer epochs per buffer, buffer size 64, rank r = 64, beta = 0.5, 100 epochs of 200 iterations, six seeds {0, 42, 123, 7, 2025, 2026}).

Citation

@article{lee2026contextaware,
  author  = {Lee, Ki-Ho and Choi, Hyun-Ho and Lee, Jung-Ryun},
  title   = {Context-Aware Embedding Masking Based on Reinforcement
             Learning for Semantic Multiplexing},
  journal = {IEEE Wireless Communications Letters},
  year    = {2026},
  note    = {WCL2026-1544}
}