PaperlayerView exact revision on arXiv

Source-linked AI summary

WirelessJEPA: A Multi-Antenna Foundation Model using Spatio-temporal Wireless Latent Predictions

Viet Chu, Omar Mashaal, Hatem Abou-Zeid

arXiv:2601.20190v1eess.SP

TL;DR

WirelessJEPA addresses limitations of contrastive wireless representation learning by predicting latent representations of masked regions in raw multi-antenna IQ data. Its antenna-time representation and spatio-temporal masking support transferable representations, with gains over contrastive baselines on most evaluated out-of-distribution datasets.

  • Problem

    Contrastive wireless representation learning relies on carefully crafted augmentations that can encourage shortcut features and provide weak inductive bias for raw-signal dynamics.

  • Method

    WirelessJEPA uses JEPA masked latent prediction on multi-antenna IQ data, reshaped into 2D antenna-time grids with structured spatio-temporal mask geometries.

  • Results

    WirelessJEPA outperforms the contrastive baseline on four of five OOD datasets, including absolute gains of 8.9% on RML16 and 8.4% on 5G Interference.

  • Takeaways & Limitations

    The results establish JEPA-style pretraining as a practical direction for wireless foundation models and motivate masking strategies that capture spatial and temporal structure.

Abstract

from arXiv · show

We propose WirelessJEPA, a novel wireless foundation model (WFM) that uses the Joint Embedding Predictive Architecture (JEPA). WirelessJEPA learns general-purpose representations directly from real-world multi-antenna IQ data by predicting latent representations of masked signal regions. This enables multiple diverse downstream tasks without reliance on carefully engineered contrastive augmentations. To adapt JEPA to wireless signals, we introduce a 2D antenna time representation that reshapes multi-antenna IQ streams into structured grids, allowing convolutional processing with block masking and efficient sparse computation over unmasked patches. Building on this representation, we propose novel spatio temporal mask geometries that encode inductive biases across antennas and time. We evaluate WirelessJEPA across six downstream tasks and demonstrate it's robust performance and strong task generalization. Our results establish that JEPA-based learning as a promising direction for building generalizable WFMs.

I. INTRODUCTION

WirelessJEPA addresses limited label efficiency and cross-task generalization in wireless learning by applying JEPA to real-world multi-antenna IQ data. It combines antenna-time representations and structured masking, then evaluates transfer across six downstream tasks.

  • Supervised wireless models require costly labeled data and typically do not generalize across tasks.
  • WirelessJEPA learns latent representations from real-world multi-antenna IQ data without carefully crafted contrastive augmentations.
  • A 2D antenna-time grid enables convolutional processing, block masking, and sparse computation over unmasked patches.
  • Structured spatio-temporal mask geometries introduce inductive biases across antennas and time without handcrafted augmentations.
  • WirelessJEPA is evaluated on six diverse downstream tasks to assess its performance across wireless applications.
  • The reported results support JEPA-based learning as a direction for generalizable wireless foundation models.

II. PROBLEM FORMULATION

The formulation represents raw multi-antenna complex baseband signals as real-valued tensors and trains a shared encoder to produce transferable representations. JEPA feature prediction over masked regions provides the self-supervised learning objective.

  • A. Signal Model: The receiver observes raw complex baseband signals from a uniform rectangular array.
  • A. Signal Model: Each input is a real-valued tensor with two I/Q channels, H receiver antennas, and W time samples.
  • A. Signal Model: Unitmax normalization scales each signal by its maximum magnitude, while residual CFO and timing drift remain uncompensated in software.
  • B. Learning Objective: Given unlabeled raw IQ data, the objective is to learn an encoder whose representations transfer across downstream wireless tasks with low-shot and lightweight adaptation.
  • B. Learning Objective: The encoder and predictor are optimized by minimizing JEPA feature-prediction loss over masked regions.
  • III. WIRELESSJEPA: The methodology covers architecture, masking strategies, training algorithms, and pretraining and downstream datasets.

A. Proposed Architecture: Learning from the Latent Space

WirelessJEPA predicts latent representations of masked multi-antenna IQ regions rather than reconstructing raw patches. Its CNN-JEPA design uses structured antenna-time inputs and mask geometries that impose spatial, temporal, or joint inductive biases.

  • Latent-space prediction: WirelessJEPA masks spatial-temporal regions of IQ samples and predicts their latent representations instead of reconstructing raw patches.
  • Latent-space prediction: An encoder represents unmasked IQ regions, while a predictor estimates the latent representations of masked regions.
  • Latent-space prediction: After pretraining, the learned encoder serves as the wireless foundation model for multiple downstream tasks.
  • CNN-JEPA components: CNN-JEPA combines a sparse convolution context encoder, lightweight predictor, and momentum teacher encoder.
  • CNN-JEPA components: ShuffleNetV2-x0.5 provides a compact encoder and matches the IQFM baseline for controlled JEPA-versus-contrastive comparison.
  • Antenna-time representation: Antenna-dimension upsampling converts the nonsquare four-antenna by 256-sample input into a square 256×256 antenna-time plane for 2D masking.
  • Mask geometry: Antenna, time, and multi-block masks respectively emphasize cross-antenna, temporal, and joint spatio-temporal dependencies.

C. Training Algorithm

WirelessJEPA trains a context encoder and predictor to infer masked latent representations from multi-antenna IQ data, using a dense EMA teacher for stable targets. Sparse masking controls information flow through the encoder while the loss is computed only on masked locations.

  • Training pipeline: The training procedure combines context encoding, latent prediction, and teacher-generated targets in two main stages.The context encoder processes masked input, while the prediction stage estimates masked-region IQ embeddings.
  • Sparse computation: A sparse convolution process reapplies the upsampled latent mask after convolution and pooling layers to suppress activations originating from masked regions.Normalization layers operate only on unmasked positions, while masked locations remain zeroed through other layers.
  • Context encoder and predictor: The context encoder processes masked IQ input, inserts a learnable mask token at masked latent indices, and predicts their latent representations.The masked input and latent mask are prepared before prediction.
  • Teacher encoder: The teacher encoder processes unmasked input densely to produce stable target embeddings, with parameters updated by exponential moving average rather than gradient descent.The momentum coefficient τ controls the teacher update rate.
  • Optimization: The model minimizes an L2 regression loss over masked indices, updating the context encoder and predictor while the teacher follows an EMA update.This trains the model to infer missing spatial–temporal latent representations rather than reconstruct raw waveforms directly.

D. Downstream Tasks

WirelessJEPA is evaluated on in-distribution and out-of-distribution wireless tasks spanning modulation, spatial estimation, device identification, protocol recognition, jamming, and interference classification. The evaluation covers the original testbed and five additional datasets.

  • Evaluation scope: The evaluation measures both in-distribution and out-of-distribution generalization across two testbed tasks and five additional downstream datasets.The protocol spans seven downstream datasets or task settings in total.
  • In-distribution tasks: The in-distribution testbed tasks are modulation classification and angle-of-arrival estimation using the pre-training dataset.These tasks assess modulation recognition and spatial estimation under the original data setting.
  • Out-of-distribution datasets: POWDER RF Fingerprinting evaluates four-device classification from WiFi signals collected across distinct hardware.The task tests device-level identification under hardware variation.
  • Out-of-distribution datasets: RML2016.10a evaluates 11 modulation types under SNRs from −20 dB to +18 dB using single-channel I/Q inputs.The inputs have shape (2, 1, 128).
  • Out-of-distribution datasets: The remaining datasets cover six-class GNSS jamming, four-class WiFi protocol classification, and four-class 5G NR interference recognition.The interference classes are LFM, collision, bleedover, and wideband.

IV. EXPERIMENTS

The experiments use low-shot transfer learning with a frozen pretrained encoder and lightweight classifiers, comparing WirelessJEPA against IQFM under a matched encoder setup. Training uses a fixed optimization and teacher-update configuration.

  • Evaluation protocol: The evaluation freezes the pretrained encoder and uses linear probing and non-parametric k-nearest neighbor classification for low-shot transfer.The comparison benchmark is IQFM, whose I/Q encoder uses spatial–temporal augmentations.
  • Baseline comparison: WirelessJEPA is benchmarked against IQFM using the same ShuffleNetV2 (0.5×) encoder.This keeps the encoder architecture consistent across the comparison.
  • Training configuration: Training runs for 100 epochs with AdamW, cosine learning-rate decay, and EMA teacher updates whose momentum τ increases from 0.996 to 1.0.The predictor uses three depthwise separable convolutional layers with Batch Normalization and ReLU.

A. Mask Geometry and Inductive Bias

Mask geometry determines which wireless dependencies WirelessJEPA emphasizes: temporal masking favors temporal tasks, antenna masking favors spatial estimation, and multi-block masking balances both. Structured continuity generally outperforms unstructured random masking for these signal dependencies.

  • Comparison: Table I reports modulation and AoA linear-probe accuracy for different WirelessJEPA mask geometries.The comparison directly contrasts temporal and spatial task behavior across masking strategies.
  • Random masking: Random masking is generally suboptimal because scattered isolated patches do not effectively exploit the physical continuity of wireless signals.The resulting encoder representations are described as less expressive.
  • Time masking: Time masking achieves the highest accuracy on 3 OOD tasks and 80.75% 1-shot modulation accuracy, but its 1-shot AoA accuracy is lowest at 2.70%.Its temporal inductive bias helps modulation classification while reducing spatial discrimination.
  • Antenna masking: Antenna masking achieves the highest 1-shot AoA accuracy at 40.39% but is least effective on modulation classification tasks.It occludes entire antenna inputs and emphasizes inter-antenna phase relationships.
  • Multi-block masking: Multi-block masking produces robust performance across both temporal and spatial tasks by grouping occlusions into contiguous spatiotemporal regions.The results support local continuity as a useful inductive bias for masked prediction.

B. Out-of-Distribution Performance: WirelessJEPA vs. IQFM

WirelessJEPA outperforms IQFM on four of five unseen OOD datasets, with substantial gains on RML16 and 5G Interference despite a slight GNSS Jamming deficit.

  • WirelessJEPA outperforms IQFM on four of five unseen OOD datasets using the time-masked model.
  • 8.9% absolute improvement on RML16 brings WirelessJEPA to 74.8% versus IQFM’s 65.9%.
  • 8.4% absolute improvement on 5G Interference brings WirelessJEPA to 76.3% versus IQFM’s 67.9%.
  • WirelessJEPA reaches 90.5% versus IQFM’s 83.4% on RF Fingerprinting, but trails 64.7% to 63.1% on GNSS Jamming.

C. In-Distribution Performance: WirelessJEPA vs. IQFM

On the two in-distribution tasks, performance depends on the task and masking geometry: AoA benefits from spatial inductive biases, while temporal and multi-block masking improve modulation accuracy.

  • AoA is generally more challenging than modulation classification and requires spatial inductive biases for strong performance.
  • WirelessJEPA with antenna masking performs comparably to IQFM on the in-distribution tasks.
  • WirelessJEPA with time and multi-block masking exceeds IQFM’s modulation accuracy.

D. Representation Space Analysis: Linear vs. k-NN

Frozen-encoder k-NN evaluation shows highly discriminative representation clusters, while masking-induced inductive biases shape which signal properties are captured most effectively.

  • k-NN analysis: 99.9% k-NN accuracy on modulation and 87.8% on RF Fingerprinting demonstrate discriminative clusters without gradient updates.
  • k-NN analysis: Lower k-NN accuracy on AoA reflects the time-masked encoder’s emphasis on waveform structure rather than spatial phase differences.
  • Evaluation: WirelessJEPA representations are evaluated with both linear probing and k-NN classification on frozen encoder outputs.
  • Masking inductive biases: Temporal masks favor waveform structure, whereas antenna masks emphasize spatial cues.
Loading 2601.20190v1…