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.

    Source: §I · p. 1

  • Method

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

    Source: §I · p. 1

  • 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.

    Source: B. Out-of-Distribution Performance:… · p. 5

  • 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.

    Source: §V · p. 5

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.Source: §I · p. 1
  • WirelessJEPA learns latent representations from real-world multi-antenna IQ data without carefully crafted contrastive augmentations.Source: §I · p. 1
  • A 2D antenna-time grid enables convolutional processing, block masking, and sparse computation over unmasked patches.Source: §I · p. 1
  • Structured spatio-temporal mask geometries introduce inductive biases across antennas and time without handcrafted augmentations.Source: §I · p. 1
  • WirelessJEPA is evaluated on six diverse downstream tasks to assess its performance across wireless applications.Source: §I · p. 1
  • The reported results support JEPA-based learning as a direction for generalizable wireless foundation models.Source: §I · p. 2

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.Source: A. Signal Model · p. 2
  • Each input is a real-valued tensor with two I/Q channels, H receiver antennas, and W time samples.Source: A. Signal Model · p. 2
  • Unitmax normalization scales each signal by its maximum magnitude, while residual CFO and timing drift remain uncompensated in software.Source: A. Signal Model · p. 2

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.Source: B. Learning Objective · p. 2
  • The encoder and predictor are optimized by minimizing JEPA feature-prediction loss over masked regions.Source: B. Learning Objective · p. 2

III. WIRELESSJEPA

  • The methodology covers architecture, masking strategies, training algorithms, and pretraining and downstream datasets.Source: §III · p. 2

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

CNN-JEPA components

Antenna-time representation

Mask geometry

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.Source: C. Training Algorithm · p. 3

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.Source: C. Training Algorithm · p. 3

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.Source: C. Training Algorithm · p. 3

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.Source: C. Training Algorithm · p. 3

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.Source: C. Training Algorithm · pp. 3–4

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.Source: D. Downstream Tasks · p. 4

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.Source: D. Downstream Tasks · p. 4

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.Source: D. Downstream Tasks · p. 4
  • 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).Source: D. Downstream Tasks · p. 4
  • 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.Source: D. Downstream Tasks · p. 4

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.Source: §IV · p. 4

Baseline comparison

  • WirelessJEPA is benchmarked against IQFM using the same ShuffleNetV2 (0.5×) encoder.This keeps the encoder architecture consistent across the comparison.Source: §IV · p. 4

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.Source: §IV · p. 4

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

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.Source: A. Mask Geometry and Inductive Bias · p. 4

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.Source: A. Mask Geometry and Inductive Bias · p. 4

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.Source: A. Mask Geometry and Inductive Bias · p. 4

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.Source: A. Mask Geometry and Inductive Bias · p. 4

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.

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.

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

Evaluation

Masking inductive biases

  • Temporal masks favor waveform structure, whereas antenna masks emphasize spatial cues.Source: §V · p. 5

Related papers

Loading 2601.20190v1…