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

TL;DR

WirelessJEPA introduces a multi-antenna foundation model for wireless signals using a novel predictive architecture, enabling versatile downstream tasks without extensive data augmentation.

Contribution

It adapts JEPA to wireless signals with a new 2D antenna time representation and spatio-temporal masking, advancing general-purpose wireless foundation models.

Findings

Robust performance across six downstream tasks

Strong generalization capabilities

Effective learning from real-world multi-antenna IQ data

Abstract

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.