# 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

```bash
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

```bibtex
@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}
}
```