Code and stored results for the IEEE Communications Letters submission

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.

bert_semcom.py

@@ -0,0 +1,785 @@
+# =========================================================
+# bert_semcom.py
+# BERT-based Multi-User Semantic Communication
+# with Expanded Shared Embedding + User-Wise Attention
+#
+# Key Idea:
+#   Each user's BERT embedding (768-dim) is EXPANDED to a
+#   larger shared embedding space (e.g., 768*K where K=4),
+#   enabling statistical multiplexing gain through masking
+#   in the over-provisioned shared space.
+#
+# Architecture:
+#   Tx:
+#     b_u = BERT(text_u)                    # (d_bert,)  e.g., 768
+#     e_u = W_tx(b_u)                       # (d_shared,) e.g., 3072
+#     x_u = e_u ⊙ m_u                      # user-masked in expanded space
+#     y   = Σ_u x_u                         # over-the-air superposition
+#     y_rx = channel(y)
+#
+#   Rx (User-wise attention demux):
+#     R_i = y_rx ⊙ m_i                     # candidate projections
+#     a_{u,i} = softmax( <norm(R_i), norm(q_u)> / sqrt(d) )
+#     z_u = Σ_i a_{u,i} R_i
+#     b̂_u = W_rx(z_u)                      # (d_bert,) recovered
+#
+#   Loss:
+#     L = (1-λ) * MSE(b_u, b̂_u) + λ * (1 - CosSim(b_u, b̂_u))
+#
+#   Metrics: Cosine similarity, MSE
+#
+# Statistical Multiplexing Gain:
+#   Even with U=2 users, projecting to d_shared = d_bert*4
+#   provides over-provisioned subspace for each user's mask,
+#   reducing inter-user interference and improving separation.
+#
+# Train modes: joint / maml_decoder / maml_full
+# =========================================================
+
+import argparse
+import math
+import random
+import csv
+import os
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+try:
+    from torch.func import functional_call
+except Exception:
+    from torch.nn.utils.stateless import functional_call
+
+
+# =========================================================
+# BERT Embedding Extractor (frozen)
+# =========================================================
+class BertEmbeddingExtractor:
+    """
+    Extract sentence-level embeddings from pre-trained BERT.
+    Uses mean pooling over non-padding tokens to mitigate
+    BERT [CLS] anisotropy. Returns un-normalized embeddings;
+    centering and L2-norm are applied at the cache level.
+    """
+    def __init__(self, model_name="bert-base-uncased", device="cpu"):
+        from transformers import BertModel, BertTokenizer
+        self.tokenizer = BertTokenizer.from_pretrained(model_name)
+        self.model = BertModel.from_pretrained(model_name).to(device)
+        self.model.eval()
+        self.device = device
+        self.embed_dim = self.model.config.hidden_size  # 768
+
+    @torch.no_grad()
+    def encode(self, texts):
+        """Mean-pooled token embedding (excluding padding)."""
+        inputs = self.tokenizer(
+            texts, padding=True, truncation=True,
+            max_length=64, return_tensors="pt"
+        ).to(self.device)
+        outputs = self.model(**inputs)
+        last_hidden = outputs.last_hidden_state          # (B, T, d)
+        mask = inputs["attention_mask"].unsqueeze(-1).float()  # (B, T, 1)
+        summed = (last_hidden * mask).sum(dim=1)         # (B, d)
+        count = mask.sum(dim=1).clamp(min=1.0)           # (B, 1)
+        mean_emb = summed / count                        # (B, d)
+        return mean_emb
+
+
+# =========================================================
+# Text Dataset
+# =========================================================
+def load_sentences(source="ag_news", max_sentences=50000,
+                   min_len=5, max_len=30):
+    """Load sentences from various sources."""
+    sentences = []
+
+    if source == "europarl":
+        try:
+            from datasets import load_dataset
+            ds = load_dataset("wmt14", "de-en", split="train", streaming=True)
+            for example in ds:
+                sent = example["translation"]["en"]
+                words = sent.split()
+                if min_len <= len(words) <= max_len:
+                    sentences.append(sent)
+                if len(sentences) >= max_sentences:
+                    break
+        except Exception:
+            print("[INFO] WMT14 not available, trying ag_news...")
+            source = "ag_news"
+
+    if source == "ag_news":
+        try:
+            from datasets import load_dataset
+            ds = load_dataset("ag_news", split="train")
+            for example in ds:
+                sent = example["text"]
+                first_sent = sent.split(".")[0].strip()
+                words = first_sent.split()
+                if min_len <= len(words) <= max_len:
+                    sentences.append(first_sent)
+                if len(sentences) >= max_sentences:
+                    break
+        except Exception:
+            print("[INFO] ag_news not available, using synthetic...")
+            source = "synthetic"
+
+    if source == "synthetic" or len(sentences) < 1000:
+        print("[INFO] Generating synthetic diverse sentences...")
+        templates = [
+            "The {} {} the {} in the {}.",
+            "A {} {} quickly {} the {}.",
+            "Several {} {} near the {} {}.",
+            "The {} and {} {} {} together.",
+            "Every {} must {} its own {}.",
+            "Under the {}, a {} {} {} softly.",
+            "The {} of {} {} a {} signal.",
+            "Without {}, the {} cannot {} properly.",
+            "A new {} {} from the {} {}.",
+            "The {} {} through the {} channel.",
+        ]
+        nouns = ["system", "signal", "network", "channel", "user", "device",
+                 "antenna", "receiver", "transmitter", "waveform", "protocol",
+                 "data", "packet", "frequency", "power", "noise", "beam",
+                 "satellite", "tower", "base station", "terminal", "sensor",
+                 "message", "code", "sequence", "spectrum", "bandwidth"]
+        verbs = ["processes", "transmits", "receives", "analyzes", "encodes",
+                 "decodes", "modulates", "filters", "amplifies", "detects",
+                 "estimates", "optimizes", "adapts", "allocates", "schedules"]
+        adjs = ["wireless", "digital", "analog", "robust", "adaptive",
+                "cognitive", "massive", "distributed", "cooperative", "mobile",
+                "reliable", "efficient", "dynamic", "intelligent", "semantic"]
+        for _ in range(max(max_sentences, 50000)):
+            tmpl = random.choice(templates)
+            n_slots = tmpl.count("{}")
+            fillers = [random.choice(random.choice([nouns, verbs, adjs]))
+                       for _ in range(n_slots)]
+            sentences.append(tmpl.format(*fillers))
+
+    random.shuffle(sentences)
+    return sentences[:max_sentences]
+
+
+# =========================================================
+# Channel
+# =========================================================
+def apply_channel(y, snr_db, channel="rayleigh"):
+    """
+    y: (d_shared,) real-valued, power-normalized embedding vector.
+    Rayleigh: magnitude fading |h| where h ~ CN(0,1).
+    """
+    snr_lin = 10 ** (snr_db / 10.0)
+    noise_var = 1.0 / snr_lin
+
+    if channel == "rayleigh":
+        # Complex Rayleigh fading -> magnitude |h|
+        h_real = torch.randn((), device=y.device)
+        h_imag = torch.randn((), device=y.device)
+        h_mag = torch.sqrt(h_real**2 + h_imag**2) / math.sqrt(2.0)
+        y = h_mag * y
+
+    noise = torch.randn_like(y) * math.sqrt(noise_var)
+    return y + noise
+
+
+# =========================================================
+# Model: Expanded Shared Embedding + User-Wise Attention
+# =========================================================
+class BertSemComMux(nn.Module):
+    """
+    BERT-based Multi-User Semantic Communication.
+
+    Key design: BERT embedding (d_bert=768) is EXPANDED to
+    d_shared = d_bert * mux_factor (e.g., 768*4=3072) to
+    provide statistical multiplexing gain.
+
+    Even with fewer users than mux_factor, the expanded space
+    gives each user more "room" for orthogonal mask allocation,
+    reducing inter-user interference.
+    """
+    def __init__(self, U, d_bert, d_shared, hidden=512):
+        super().__init__()
+        self.U = U
+        self.d_bert = d_bert          # 768
+        self.d_shared = d_shared      # e.g., 768*4 = 3072
+
+        # Tx: single-layer linear projection to the shared space
+        # (no sub-d_bert bottleneck, no costly d_shared x d_shared layer).
+        # Followed by LayerNorm for numerical stability.
+        self.tx_proj = nn.Sequential(
+            nn.Linear(d_bert, d_shared),
+            nn.LayerNorm(d_shared),
+        )
+
+        # User-specific masks in expanded shared space
+        self.user_mask = nn.Embedding(U, d_shared)
+
+        # Rx: user-wise attention queries in shared space
+        self.user_query = nn.Embedding(U, d_shared)
+
+        # Rx: single-layer linear projection back to BERT space.
+        self.rx_proj = nn.Sequential(
+            nn.LayerNorm(d_shared),
+            nn.Linear(d_shared, d_bert),
+        )
+
+        nn.init.normal_(self.user_mask.weight, std=0.5)
+        nn.init.normal_(self.user_query.weight, std=0.5)
+
+    def forward(self, bert_embs, snr_db, channel, params=None,
+                return_intermediate=False):
+        """
+        bert_embs: (U, d_bert) - BERT sentence embeddings
+        Returns: (U, d_bert) - recovered embeddings
+        """
+        dev = bert_embs.device
+
+        # ----- Tx: Project to expanded shared space -----
+        if params is None:
+            e = self.tx_proj(bert_embs)            # (U, d_shared)
+            m = self.user_mask.weight              # (U, d_shared)
+        else:
+            e = functional_call(self.tx_proj, params["tx_proj"],
+                                (bert_embs,))
+            m = functional_call(
+                self.user_mask, params["user_mask"],
+                (torch.arange(self.U, device=dev),))
+
+        # Masking in expanded space
+        x = e * m                                   # (U, d_shared)
+        y_tx = x.sum(dim=0)                         # (d_shared,)
+
+        # ----- Power normalization -----
+        # OFDM model: each subcarrier (dimension) has unit power.
+        # Larger d_shared = more subcarriers = more bandwidth.
+        # Per-dimension SNR is the same regardless of d_shared.
+        # This models the bandwidth-quality trade-off:
+        #   d_shared=128: compressed (low bandwidth, lossy)
+        #   d_shared=768: matched (1:1)
+        #   d_shared=3072: expanded (high bandwidth, room for separation)
+        y_tx = y_tx / (torch.sqrt((y_tx ** 2).mean() + 1e-8))
+
+        # ----- Channel -----
+        y_rx = apply_channel(y_tx, snr_db, channel)
+
+        # ----- Rx: user-wise attention -----
+        R = y_rx.unsqueeze(0) * m                   # (U, d_shared)
+        Rn = F.normalize(R, p=2, dim=-1)
+
+        recovered = []
+        attn_weights = []
+        for u in range(self.U):
+            if params is None:
+                q = self.user_query(torch.tensor(u, device=dev))
+            else:
+                q = functional_call(
+                    self.user_query, params["user_query"],
+                    (torch.tensor(u, device=dev),))
+
+            qn = F.normalize(q, p=2, dim=-1)
+            # Native cosine-similarity scores (no extra 1/sqrt(d_s)
+            # damping). Since Rn, qn are unit-norm, the raw dot product
+            # is already in [-1, 1] and softmax-friendly. The standard
+            # 1/sqrt(d) attention temperature would collapse softmax to
+            # near-uniform when d_s is large (e.g., 3072).
+            scores = (Rn @ qn) * 8.0  # mild temperature sharpening
+            attn = F.softmax(scores, dim=0)
+            attn_weights.append(attn.detach())
+
+            z_u = (attn.unsqueeze(-1) * R).sum(dim=0)  # (d_shared,)
+
+            # Reverse projection: shared space -> BERT space
+            if params is None:
+                b_hat_u = self.rx_proj(z_u)
+            else:
+                b_hat_u = functional_call(self.rx_proj, params["rx_proj"],
+                                          (z_u,))
+
+            recovered.append(b_hat_u)
+
+        result = torch.stack(recovered, dim=0)  # (U, d_bert)
+
+        if return_intermediate:
+            return result, {
+                "y_tx": y_tx.detach(),
+                "y_rx": y_rx.detach(),
+                "attn": torch.stack(attn_weights, dim=0),  # (U, U)
+                "masks": m.detach(),
+                "projected": e.detach(),
+            }
+
+        return result
+
+
+# =========================================================
+# Loss Function
+# =========================================================
+def semantic_loss(b_orig, b_hat, lam=0.5):
+    """
+    Combined MSE + Cosine Similarity loss.
+    b_orig, b_hat: (U, d_bert)
+    """
+    mse = F.mse_loss(b_hat, b_orig)
+    cos_sim = F.cosine_similarity(b_hat, b_orig, dim=-1).mean()
+    loss = (1.0 - lam) * mse + lam * (1.0 - cos_sim)
+    return loss, mse.item(), cos_sim.item()
+
+
+# =========================================================
+# MAML Helpers (same structure as original maml.py)
+# =========================================================
+def split_params(model: BertSemComMux):
+    return {
+        "tx_proj": dict(model.tx_proj.named_parameters()),
+        "user_mask": dict(model.user_mask.named_parameters()),
+        "user_query": dict(model.user_query.named_parameters()),
+        "rx_proj": dict(model.rx_proj.named_parameters()),
+    }
+
+
+def select_inner_keys(train_mode: str):
+    if train_mode == "maml_decoder":
+        return ["user_query", "rx_proj"]
+    if train_mode == "maml_full":
+        return ["tx_proj", "user_mask", "user_query", "rx_proj"]
+    return []
+
+
+def ordered_param_items(param_dict: dict):
+    return [(k, param_dict[k]) for k in sorted(param_dict.keys())]
+
+
+def gather_inner_params(fast_params: dict, inner_keys: list):
+    flat_list, meta_index = [], []
+    for key in inner_keys:
+        for name, p in ordered_param_items(fast_params[key]):
+            flat_list.append(p)
+            meta_index.append((key, name))
+    return flat_list, meta_index
+
+
+def apply_inner_update(fast_params, inner_keys, meta_index, grads,
+                       inner_lr, maml_order):
+    new_fast = {k: dict(v) for k, v in fast_params.items()}
+    for (key, name), g in zip(meta_index, grads):
+        if maml_order == "first":
+            g = g.detach()
+        new_fast[key][name] = new_fast[key][name] - inner_lr * g
+    return new_fast
+
+
+# =========================================================
+# Embedding Cache (pre-compute BERT embeddings)
+# =========================================================
+class EmbeddingCache:
+    """Pre-compute, center, and L2-normalize BERT embeddings.
+    Centering removes the dominant mean direction to reduce
+    anisotropy; L2 normalization gives unit-norm embeddings
+    suitable for cosine-similarity-based loss/metrics.
+    """
+    def __init__(self, bert_extractor, sentences, batch_size=64):
+        self.embeddings = []
+        self.sentences = sentences
+        print(f"[INFO] Pre-computing BERT embeddings for "
+              f"{len(sentences)} sentences...")
+        for i in range(0, len(sentences), batch_size):
+            batch = sentences[i:i+batch_size]
+            emb = bert_extractor.encode(batch)
+            self.embeddings.append(emb.cpu())
+        self.embeddings = torch.cat(self.embeddings, dim=0)
+
+        # Center (subtract mean direction) to reduce anisotropy
+        self.mean = self.embeddings.mean(dim=0, keepdim=True)
+        self.embeddings = self.embeddings - self.mean
+        # L2-normalize
+        self.embeddings = F.normalize(self.embeddings, p=2, dim=-1)
+
+        # Sanity check: random pairwise cos sim should be near 0
+        n_check = min(500, len(self.embeddings))
+        idx1 = torch.randperm(len(self.embeddings))[:n_check]
+        idx2 = torch.randperm(len(self.embeddings))[:n_check]
+        cs = F.cosine_similarity(self.embeddings[idx1],
+                                  self.embeddings[idx2], dim=-1)
+        print(f"[INFO] Cached {self.embeddings.shape[0]} embeddings, "
+              f"shape={self.embeddings.shape}")
+        print(f"[INFO] After centering: random pair cos sim "
+              f"mean={cs.mean():.4f}, std={cs.std():.4f}")
+
+    def sample(self, U, device):
+        """Sample U random embeddings."""
+        idx = torch.randint(0, len(self.embeddings), (U,))
+        return self.embeddings[idx].to(device), idx
+
+    def get_sentences(self, idx):
+        """Get sentences by index."""
+        return [self.sentences[i] for i in idx]
+
+
+# =========================================================
+# Evaluation
+# =========================================================
+@torch.no_grad()
+def evaluate(model, cache, args, device):
+    model.eval()
+    total_cos, total_mse, n_trials = 0.0, 0.0, 0
+
+    if not args.eval_adapt:
+        for _ in range(args.eval_trials):
+            b, _ = cache.sample(args.users, device)
+            b_hat = model(b, args.snr, args.channel, params=None)
+            cos = F.cosine_similarity(b_hat, b, dim=-1).mean().item()
+            mse = F.mse_loss(b_hat, b).item()
+            total_cos += cos
+            total_mse += mse
+            n_trials += 1
+        return total_cos / n_trials, total_mse / n_trials
+
+    # Few-shot adaptation
+    inner_keys = select_inner_keys(args.eval_adapt_mode)
+    base_params = split_params(model)
+
+    for _ in range(args.eval_trials):
+        b_sup, _ = cache.sample(args.users, device)
+        fast_params = {k: {n: p for n, p in v.items()}
+                       for k, v in base_params.items()}
+
+        with torch.enable_grad():
+            for _ in range(args.eval_inner_steps):
+                b_hat_sup = model(b_sup, args.snr, args.channel,
+                                  params=fast_params)
+                loss_sup, _, _ = semantic_loss(b_sup, b_hat_sup,
+                                               args.loss_lambda)
+                flat_list, meta_index = gather_inner_params(
+                    fast_params, inner_keys)
+                grads = torch.autograd.grad(loss_sup, flat_list,
+                                            create_graph=False)
+                fast_params = apply_inner_update(
+                    fast_params, inner_keys, meta_index, grads,
+                    inner_lr=args.eval_inner_lr, maml_order="first")
+
+        b_q, _ = cache.sample(args.users, device)
+        b_hat_q = model(b_q, args.snr, args.channel, params=fast_params)
+        cos = F.cosine_similarity(b_hat_q, b_q, dim=-1).mean().item()
+        mse = F.mse_loss(b_hat_q, b_q).item()
+        total_cos += cos
+        total_mse += mse
+        n_trials += 1
+
+    return total_cos / n_trials, total_mse / n_trials
+
+
+# =========================================================
+# Train & Eval
+# =========================================================
+def train_and_eval(args):
+    if args.cuda and torch.cuda.is_available():
+        device = torch.device("cuda")
+    elif torch.backends.mps.is_available():
+        device = torch.device("mps")
+    else:
+        device = torch.device("cpu")
+
+    torch.manual_seed(args.seed)
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+
+    # ---- Load BERT and pre-compute embeddings ----
+    bert = BertEmbeddingExtractor(args.bert_model, device)
+    d_bert = bert.embed_dim  # 768
+
+    # Compute d_shared = d_bert * mux_factor
+    d_shared = d_bert * args.mux_factor
+    print(f"[INFO] d_bert={d_bert}, mux_factor={args.mux_factor}, "
+          f"d_shared={d_shared}")
+
+    sentences = load_sentences(args.data_source, args.max_sentences)
+    print(f"[INFO] Loaded {len(sentences)} sentences")
+
+    cache = EmbeddingCache(bert, sentences,
+                           batch_size=args.bert_batch_size)
+
+    # Free BERT from GPU
+    del bert
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    # ---- Build model ----
+    model = BertSemComMux(args.users, d_bert, d_shared,
+                          args.hidden).to(device)
+    opt = torch.optim.Adam(model.parameters(), lr=args.lr)
+
+    train_snr_list = list(args.train_snr)
+    inner_keys = select_inner_keys(args.train_mode)
+
+    # ---- Config ----
+    print("=" * 60)
+    print(" BERT Semantic Communication - Expanded Shared Embedding")
+    print("=" * 60)
+    print(f" Device              : {device}")
+    print(f" Train mode          : {args.train_mode}")
+    print(f" MAML order          : {args.maml_order}")
+    print(f" Users (U)           : {args.users}")
+    print(f" BERT dim (d_bert)   : {d_bert}")
+    print(f" Mux factor (K)      : {args.mux_factor}")
+    print(f" Shared dim (d_shared): {d_shared}")
+    print(f" Hidden (MLP)        : {args.hidden}")
+    print(f" Channel             : {args.channel}")
+    print(f" Loss lambda         : {args.loss_lambda}")
+    print("-" * 60)
+    print(f" Train SNRs (dB)     : {train_snr_list}")
+    print(f" Eval SNR (dB)       : {args.snr}")
+    print("-" * 60)
+    print(f" Epochs              : {args.epochs}")
+    print(f" Steps/epoch         : {args.train_steps_per_epoch}")
+    print(f" Eval trials         : {args.eval_trials}")
+    if args.train_mode != "joint":
+        print("-" * 60)
+        print(f" Inner steps         : {args.inner_steps}")
+        print(f" Inner LR            : {args.inner_lr}")
+        print(f" Meta-batch          : {args.meta_batch}")
+        print(f" Inner keys          : {inner_keys}")
+    print("=" * 60)
+
+    # ---- CSV ----
+    os.makedirs(args.save_dir, exist_ok=True)
+    csv_path = os.path.join(
+        args.save_dir,
+        f"bert_{args.train_mode}_{args.users}U_K{args.mux_factor}.csv"
+    )
+
+    if not os.path.exists(csv_path):
+        with open(csv_path, "w", newline="") as f:
+            writer = csv.writer(f)
+            writer.writerow([
+                "epoch", "train_mode", "maml_order",
+                "users", "d_bert", "mux_factor", "d_shared",
+                "channel", "train_snrs", "eval_snr",
+                "inner_steps", "inner_lr", "meta_batch",
+                "avg_loss", "cos_sim", "mse"
+            ])
+
+    # ---- Training ----
+    for ep in range(1, args.epochs + 1):
+        model.train()
+        loss_meter = 0.0
+
+        for _ in range(args.train_steps_per_epoch):
+            if args.train_mode == "joint":
+                b, _ = cache.sample(args.users, device)
+                snr = float(random.choice(train_snr_list))
+                b_hat = model(b, snr, args.channel, params=None)
+                loss, _, _ = semantic_loss(b, b_hat, args.loss_lambda)
+                opt.zero_grad(set_to_none=True)
+                loss.backward()
+                opt.step()
+                loss_meter += loss.item()
+
+            else:
+                # MAML
+                if args.meta_batch > len(train_snr_list):
+                    snr_tasks = [float(random.choice(train_snr_list))
+                                 for _ in range(args.meta_batch)]
+                else:
+                    snr_tasks = [float(x) for x in
+                                 random.sample(train_snr_list,
+                                               args.meta_batch)]
+
+                base_params = split_params(model)
+                meta_loss = torch.tensor(0.0, device=device)
+
+                for snr in snr_tasks:
+                    fast_params = {
+                        k: {n: p for n, p in v.items()}
+                        for k, v in base_params.items()
+                    }
+
+                    for _ in range(args.inner_steps):
+                        b_sup, _ = cache.sample(args.users, device)
+                        b_hat_sup = model(b_sup, snr, args.channel,
+                                          params=fast_params)
+                        loss_sup, _, _ = semantic_loss(
+                            b_sup, b_hat_sup, args.loss_lambda)
+
+                        flat_list, meta_index = gather_inner_params(
+                            fast_params, inner_keys)
+                        create_graph = (args.maml_order == "second")
+                        grads = torch.autograd.grad(
+                            loss_sup, flat_list,
+                            create_graph=create_graph)
+                        fast_params = apply_inner_update(
+                            fast_params, inner_keys, meta_index, grads,
+                            inner_lr=args.inner_lr,
+                            maml_order=args.maml_order)
+
+                    b_q, _ = cache.sample(args.users, device)
+                    b_hat_q = model(b_q, snr, args.channel,
+                                    params=fast_params)
+                    qloss, _, _ = semantic_loss(
+                        b_q, b_hat_q, args.loss_lambda)
+                    meta_loss = meta_loss + qloss
+
+                meta_loss = meta_loss / float(args.meta_batch)
+                opt.zero_grad(set_to_none=True)
+                meta_loss.backward()
+                opt.step()
+                loss_meter += meta_loss.item()
+
+        avg_loss = loss_meter / float(args.train_steps_per_epoch)
+
+        # ---- Eval ----
+        cos_sim, mse = evaluate(model, cache, args, device)
+
+        print(f"[Epoch {ep:03d}/{args.epochs}] loss={avg_loss:.4f} | "
+              f"SNR={args.snr:.1f}dB | CosSim={cos_sim:.4f} | "
+              f"MSE={mse:.4e}")
+
+        with open(csv_path, "a", newline="") as f:
+            writer = csv.writer(f)
+            writer.writerow([
+                ep, args.train_mode, args.maml_order,
+                args.users, d_bert, args.mux_factor, d_shared,
+                args.channel, train_snr_list, args.snr,
+                (args.inner_steps if args.train_mode != "joint" else ""),
+                (args.inner_lr if args.train_mode != "joint" else ""),
+                (args.meta_batch if args.train_mode != "joint" else ""),
+                avg_loss, cos_sim, mse
+            ])
+
+    print(f"\n✅ Results saved to: {csv_path}")
+    return model
+
+
+# =========================================================
+# Main
+# =========================================================
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="BERT-based Multi-User Semantic Communication "
+                    "with Expanded Shared Embedding")
+
+    # training mode
+    parser.add_argument("--train-mode",
+                        choices=["joint", "maml_decoder", "maml_full"],
+                        default="joint")
+    parser.add_argument("--maml-order", choices=["first", "second"],
+                        default="first")
+
+    # model
+    parser.add_argument("--users", type=int, default=4)
+    parser.add_argument("--mux-factor", type=int, default=4,
+                        help="Shared dim = d_bert * mux_factor "
+                             "(e.g., 4 -> 768*4=3072)")
+    parser.add_argument("--hidden", type=int, default=512)
+
+    # BERT
+    parser.add_argument("--bert-model", type=str,
+                        default="bert-base-uncased")
+    parser.add_argument("--bert-batch-size", type=int, default=64)
+
+    # data
+    parser.add_argument("--data-source",
+                        choices=["europarl", "ag_news", "synthetic"],
+                        default="ag_news")
+    parser.add_argument("--max-sentences", type=int, default=50000)
+
+    # channel
+    parser.add_argument("--channel", choices=["awgn", "rayleigh"],
+                        default="rayleigh")
+
+    # loss
+    parser.add_argument("--loss-lambda", type=float, default=0.5)
+
+    # training
+    parser.add_argument("--epochs", type=int, default=100)
+    parser.add_argument("--train-steps-per-epoch", type=int, default=2000)
+    parser.add_argument("--train-snr", type=float, nargs="+",
+                        default=[0, 5, 10, 15, 20, 25])
+    parser.add_argument("--lr", type=float, default=1e-3)
+
+    # MAML
+    parser.add_argument("--inner-lr", type=float, default=1e-3)
+    parser.add_argument("--inner-steps", type=int, default=1)
+    parser.add_argument("--meta-batch", type=int, default=4)
+
+    # evaluation
+    parser.add_argument("--snr", type=float, default=10.0)
+    parser.add_argument("--eval-trials", type=int, default=5000)
+
+    # test-time adaptation
+    parser.add_argument("--eval-adapt", action="store_true")
+    parser.add_argument("--eval-adapt-mode",
+                        choices=["maml_decoder", "maml_full"],
+                        default="maml_decoder")
+    parser.add_argument("--eval-inner-steps", type=int, default=1)
+    parser.add_argument("--eval-inner-lr", type=float, default=1e-3)
+
+    # misc
+    parser.add_argument("--save-dir", type=str, default="results_bert")
+    parser.add_argument("--cuda", action="store_true")
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--debug-maml", action="store_true")
+
+    args = parser.parse_args()
+    train_and_eval(args)
+
+
+"""
+========================
+Example Commands
+========================
+
+[1] Joint training (U=4, K=4 -> d_shared=3072, Rayleigh)
+python3 bert_semcom.py \
+  --train-mode joint \
+  --users 4 --mux-factor 4 \
+  --snr 10 --channel rayleigh \
+  --data-source ag_news \
+  --epochs 50 --train-steps-per-epoch 2000 \
+  --eval-trials 5000 --cuda
+
+[2] Statistical mux gain: U=2, still K=4 (over-provisioned)
+python3 bert_semcom.py \
+  --train-mode joint \
+  --users 2 --mux-factor 4 \
+  --snr 10 --channel rayleigh \
+  --epochs 50 --cuda
+
+[3] Compare mux factors: K=1,2,4,8
+for K in 1 2 4 8; do
+  python3 bert_semcom.py \
+    --train-mode joint \
+    --users 4 --mux-factor $K \
+    --snr 10 --epochs 50 --cuda
+done
+
+[4] Decoder-only MAML (U=4, K=4)
+python3 bert_semcom.py \
+  --train-mode maml_decoder --maml-order first \
+  --users 4 --mux-factor 4 \
+  --meta-batch 4 --inner-lr 5e-4 --inner-steps 1 \
+  --snr 10 --epochs 50 --cuda
+
+[5] Full MAML
+python3 bert_semcom.py \
+  --train-mode maml_full --maml-order first \
+  --users 4 --mux-factor 4 \
+  --meta-batch 4 --inner-lr 5e-4 --inner-steps 1 \
+  --snr 10 --epochs 50 --cuda
+
+[6] SNR sweep
+for snr in 0 5 10 15 20 25 30; do
+  python3 bert_semcom.py \
+    --train-mode joint \
+    --users 4 --mux-factor 4 \
+    --snr $snr --epochs 50 --cuda
+done
+
+[7] User count comparison (all with K=4)
+for u in 2 4 8 16; do
+  python3 bert_semcom.py \
+    --train-mode joint \
+    --users $u --mux-factor 4 \
+    --snr 10 --epochs 50 --cuda
+done
+"""