LWM-CDE: dataset distancing for wireless datasets

A new base station often has only a few labeled samples, while a library holds the datasets of many other sites. Deciding which of those datasets to use usually means training one model per candidate and testing each one. This repository measures instead how close two wireless datasets are, in the space of a pretrained wireless foundation model, the Large Wireless Model (LWM 1.1), and uses that distance to make the decision without training a model per candidate. It is the code of the paper LWM-CDE: A Representation Space for Wireless Data Reasoning and Transferability.

What the distance decides

The same distance serves four decisions about data:

1. Transfer learning. A new site learns from the library dataset closest to its few labeled samples.
2. Data augmentation. A site with few labels adds samples from the closest datasets, because data that is too different can lower its accuracy.
3. Pretraining data selection. When only k datasets fit the budget, k-medoids on the distances picks the k that cover the library best, because diverse training data generalizes best.
4. Model monitoring. The distance between new data and a model's training data needs no labels, so a site can track it and retrain, switch models or fall back when it passes a threshold.

A live demo, the Dataset Distancing Lab, computes these distances in the browser. The LWM family is presented at lwm-wireless.net.

Installation

git clone https://huggingface.co/wi-lab/lwm-cde
cd lwm-cde
pip install -r requirements.txt
python scripts/download_lwm.py

The code needs Python 3.9 or newer and PyTorch 2. A GPU speeds up the experiments but is not required for the quick start. scripts/download_lwm.py downloads the LWM 1.1 weights from wi-lab/lwm-v1.1 and checks them. DeepMIMO 4.0.1 generates the channels and labels, and it downloads each scenario the first time it is used (into ~/.cache/lwm/deepmimo_scenarios).

Quick start

from dataset_distancing import DatasetCollection, DistanceConfig, compute_distances
from dataset_distancing.applications import augmentation_ranking, budgeted_pretraining, source_selection
from dataset_distancing.embeddings import load_runner

library = DatasetCollection.from_scenarios(
    [{"name": "city_0_newyork_3p5", "tx_set": 1}, {"name": "city_1_losangeles_3p5", "tx_set": 1},
     {"name": "city_2_chicago_3p5", "tx_set": 1}, {"name": "city_3_houston_3p5", "tx_set": 1}],
    task="los-nlos", n_ant=32, n_sub=32, max_samples=600)
runner = load_runner()                                # LWM 1.1, used as is
D = compute_distances(library, runner, cfg=DistanceConfig(distance="swd", label_aware=True))

source_selection(D, target_indices=[0], top_k=1)      # the dataset site 0 should learn from
augmentation_ranking(D, 0)                            # the datasets to add to site 0, closest first
budgeted_pretraining(D, k=2)                          # the two datasets that cover the library best

python examples/quickstart.py runs these steps on eight DeepMIMO base stations and adds a monitoring check, which compares the distance of new data to the training data with the distance between two random parts of the training data.

How the distance is computed

LWM 1.1 splits each channel (32 antennas by 32 subcarriers) into patches, and its [CLS] token gives one 128-value embedding per channel. The distance between two datasets is the label-aware sliced Wasserstein distance. For each class that both datasets hold (LOS and NLOS, or the beam index), it projects the two sets of embeddings on random directions, compares the one-dimensional distributions, and averages over directions and classes. DistanceConfig selects the distance ("swd", "centroid" or "cosine"), the number of samples and directions, and a directed variant (asymmetric=True) that scores source-to-target transfer. The tutorial also computes the distance on raw channels and in a UMAP space, the baselines of the paper.

LWM-CDE fine-tuning

dataset_distancing.cde.CDETrainer fine-tunes the top eight layers of LWM together with a metric head, using a measured transfer matrix of the library and the four-term objective of the paper: a listwise ranking term, a view-consistency term, a correlation term and a distillation term. After training, the distances are computed on the fine-tuned [CLS] embeddings. The fine-tuned checkpoints of the paper are not part of this release; the tutorial's --cde option trains one.

Reproducing the paper

• tutorial.py generates datasets from DeepMIMO, measures the transfer matrix by training one model per source and testing it on every target, computes the distances, and reports how well they rank the measured transfer (Pearson and Spearman correlation over trials). Its last lines list example commands for LOS/NLOS identification and beam prediction.
• experiments/ holds the experiments behind Table II (transfer learning), Fig. 2 (pretraining data selection) and Table VIII (from 3.5 GHz to 28 GHz). experiments/README.md gives their commands.

Repository layout

• dataset_distancing/ is the package: datasets from DeepMIMO (data.py), LWM embeddings (embeddings.py), distances (distances.py), the transfer matrix (performance.py), correlations (correlation.py), dataset selection (selection.py, applications.py) and LWM-CDE fine-tuning (cde.py).
• third_party/lwm-1.1/ holds the code of LWM 1.1; scripts/download_lwm.py adds its weights.
• examples/quickstart.py, tutorial.py and tutorial.ipynb are the entry points; experiments/ reproduces the paper's results.

Citation

@article{alikhani2026lwmcde,
  title   = {{LWM-CDE}: A Representation Space for Wireless Data Reasoning and Transferability},
  author  = {Alikhani, Sadjad and Malhotra, Akshay and Hamidi-Rad, Shahab and Alkhateeb, Ahmed},
  journal = {arXiv preprint arXiv:2605.24077},
  year    = {2026}
}

License

The code is released under the Apache License 2.0 (LICENSE). It includes the LWM 1.1 code under the same license (NOTICE). DeepMIMO is installed as a dependency and is licensed under the GNU General Public License, version 2 or later.