Source-linked AI summary
JEPA-Inspired Latent Field Prediction for Sparse-Pilot MIMO-OFDM Channel Reconstruction
Yuzhi Yang, Brahim Mefgouda, Hang Zou, Lina Bariah, Anis Bara, Yuhuan Lu, Hao Zhang, Merouane Debbah
TL;DR
Sparse-pilot CSI acquisition matters because dense pilots consume spectral resources, while direct raw-CSI prediction is fragile despite strong channel correlations. The paper predicts latent propagation-field states and fuses them with sparse noisy current pilots; the best JEPA-based method outperforms simple estimation and direct-training baselines.
Problem
Dense pilot transmission can consume considerable spectral resources, while direct raw-CSI prediction is fragile and existing prediction approaches leave sparse-pilot reconstruction with predicted latent priors less explored.
Source: §I · p. 1Method
A JEPA-inspired field world model predicts masked or future latent propagation-field states from historical multi-band CSI and fuses the predicted prior with sparse noisy current pilots for reconstruction.
Source: §I · p. 2Results
The best JEPA-based method outperforms both simple estimation and direct-training baselines; among prediction-aided objectives, cov and cov + diff perform best across the noise range, similarly to each other.
Source: F. Channel Prediction and… · p. 6Takeaways & Limitations
Abstract
from arXiv · showhide
Accurate acquisition of full channel state information (CSI) is essential for multiple-input multiple-output orthogonal frequency-division multiplexing (MIMO-OFDM) wireless systems, while dense pilot-based estimation incurs substantial overhead. Directly minimizing the prediction error over raw CSI coefficients is fragile in realistic propagation environments, since small-scale fading, mobility, local scattering, and small spatial perturbations can induce large phase-sensitive variations. Nevertheless, wireless channels contain strong spatial, temporal, and frequency-domain correlations that can be more predictably represented in a latent propagation-field domain. Motivated by this observation, we propose a joint-embedding predictive architecture (JEPA)-inspired field world model, defined as a latent predictive model of the wireless propagation field, for sparse-pilot channel reconstruction. The model maps CSI observations into latent time-frequency tokens, predicts masked or future latent field states from historical multi-band context, and fuses the predicted latent prior with sparse noisy current pilots to reconstruct the full MIMO-OFDM channel. Ray-tracing experiments show that the proposed method improves normalized mean square error (NMSE) over pilot-only and direct reconstruction baselines, supporting latent field prediction as an effective prior for sparse-pilot channel acquisition.
I. INTRODUCTION
Dense pilots can consume spectral resources, while direct prediction of raw CSI is fragile to phase-sensitive variation. The paper proposes predicting latent propagation-field structure and combining it with sparse current pilots for full-channel reconstruction.
I. INTRODUCTION
- Sparse-pilot acquisition is motivated by the substantial spectral resources that repeated pilots can consume as channels vary across time slots and frequency bands.Channel-observation correlations can support full-CSI reconstruction from sparse current pilots.Source: §I · p. 1
- Raw CSI prediction is fragile to phase-sensitive variations, whereas propagation fields retain spatial, temporal, and frequency-domain correlations that can serve as a structured prior.The prior complements rather than replaces pilot-based estimation.Source: §I · p. 1
- Prior prediction-oriented methods often operate on raw CSI, and sparse-pilot full-channel reconstruction using predicted latent states as priors remains less explored.The paper identifies this as a gap alongside the separation of predictable propagation structure from instantaneous fluctuations.Source: §I · pp. 1–2
- The proposed JEPA-inspired field world model predicts masked or future latent field states from historical multi-band context instead of extrapolating raw CSI coefficients.It is defined as a latent predictive model of the wireless propagation field, with the predicted state used as a prior alongside sparse noisy pilots.Source: §I · p. 2
- The pilot-assisted framework fuses its predicted latent prior with sparse noisy current pilots, while overly detail-preserving auxiliary objectives can reduce downstream reconstruction usefulness.The contribution frames latent prediction as a prior for full-channel recovery, not a replacement for pilots.Source: §I · p. 2
II. PROPOSED METHOD
The method represents CSI as time-frequency latent tokens, predicts masked or future cells from visible historical context, and uses the prediction as a prior for pilot-assisted reconstruction. Its pipeline includes CSI tokenization, JEPA-style prediction, and a downstream reconstruction head.
A. System Overview
- The system maps each time-subband CSI observation to latent tokens, predicts latent representations for masked or future cells, and combines the predicted prior with sparse pilots in a downstream head.The downstream head reconstructs the full CSI from the prior and pilot observations.Source: A. System Overview · p. 2, Figure 2 · p. 3
B. JEPA-Based Time-Frequency Latent Prediction
- Each visible CSI cell is tokenized into K latent tokens of dimension D, and the cells are organized as a time-frequency field.Training masks all cells at the final time slot for next-step prediction and randomly masks historical cells to encourage completion from incomplete context.Source: B. JEPA-Based Time-Frequency Latent… · p. 2, Figure 1 · p. 3
- A context encoder processes visible cells while learnable query tokens represent masked cells; the predictor maps its outputs back into the tokenizer latent space.Masked cells are not provided to the context encoder.Source: B. JEPA-Based Time-Frequency Latent… · p. 2
- The prediction task infers latent states of missing or future CSI cells from observed time-frequency context rather than extrapolating raw channel coefficients directly.The predicted latent field serves as a structured prior for reconstruction from sparse current pilots.Source: A. System Overview · p. 2, B. JEPA-Based Time-Frequency Latent… · p. 3
C. Self-Supervised Training Objective
The self-supervised objective predicts masked-cell latent targets using a momentum tokenizer and adds covariance regularization to discourage redundant latent dimensions. CSI-domain grounding losses are examined as ablations, with different degrees of detail preservation.
C. Self-Supervised Training Objective
- The main JEPA loss predicts masked-cell targets using stable latent targets from a momentum tokenizer with the same architecture as the online tokenizer.The final model uses the JEPA loss plus weighted covariance loss; the grounding loss is enabled only in ablations.Source: C. Self-Supervised Training Objective · pp. 3–4
- Covariance regularization reduces redundancy among latent dimensions, addressing the risk that pure latent prediction maps CSI observations to nearly constant codes.The covariance matrix is computed over flattened latent features in the current batch.Source: C. Self-Supervised Training Objective · p. 3
- Grounding-loss ablations compare direct CSI recovery, amplitude recovery, and conditional diffusion grounding to test whether preserving CSI details affects reconstruction.Direct CSI recovery imposes the strongest detail-preservation constraint, while diffusion grounding is the weakest.Source: C. Self-Supervised Training Objective · p. 3
III. PILOT-ASSISTED CHANNEL RECONSTRUCTION WITH LATENT PRIORS
The reconstruction task combines a predicted latent prior from historical observations with sparse, noisy current pilot estimates to recover full CSI. At inference, predicted rather than teacher-forced latent states condition the reconstruction network.
III. PILOT-ASSISTED CHANNEL RECONSTRUCTION WITH LATENT PRIORS
- The reconstruction task uses a predicted latent representation based on historical channel observations and a sparse, noisy initial estimate from current pilots.The pilot estimate is modeled using a binary pilot mask and estimation noise.Source: §III · p. 4
- The final reconstruction regressor combines the predicted latent prior with the sparse, noisy pilot-based channel estimate to recover current full CSI.The latent prior captures predictable historical propagation information, while the pilot estimate preserves instantaneous measurement evidence.Source: §III · p. 4
- During evaluation, predicted latent states from historical visible cells replace teacher-forced latent conditions, and no current full CSI is used at inference.The tokenizer and latent predictor are frozen during downstream training, when only the reconstruction network is optimized.Source: §III · p. 4
A. Dataset Generation
The dataset uses ray-traced channels from a Chicago city-center scenario with three base stations and mobile-user trajectories. Training and testing are separated by complete trajectories.
A. Dataset Generation
- The dataset contains 200 randomized trajectories for each of three base stations, with 16 points per trajectory and adjacent points separated by 50 ms.UE positions receive Gaussian perturbations with 5 cm standard deviation; channels span 16 time samples, 8 subbands, 32 subcarriers, and 16 antennas.Source: A. Dataset Generation · p. 4
- Complete UE trajectories are split into disjoint training and testing sets before temporal slices are generated, preventing trajectory overlap across splits.Trajectory endpoints are within 100 m horizontally of their corresponding base station, and the endpoint distance is 1–10 meters.Source: A. Dataset Generation · pp. 4–5
B. Visibility Mask Generation
JEPA training always masks the final time step and randomly masks historical steps. The historical masking rate shifts over training from a fixed starting rate to a stochastic target distribution.
B. Visibility Mask Generation
- The final time step is always fully masked, while historical time steps are randomly masked to train prediction from incomplete context.The initial masking ratio is fixed because the JEPA tokenizer is not yet powerful enough early in training.Source: B. Visibility Mask Generation · p. 5
- The historical masking rate interpolates from 0.3 toward a sample from a Beta(4.0, 1.2)-based distribution over [0.1, 0.9].The interpolation weight increases linearly from 0 to 1, progressively increasing task difficulty while the final time step remains fully masked.Source: B. Visibility Mask Generation · p. 5
C. Network Topology and Training Settings
The network tokenizes complex CSI into latent embeddings, predicts latent states with a ViT-style predictor, and conditions downstream reconstruction on the predicted latent prior. Training also includes auxiliary recovery heads, sparse noisy pilots, and an exponential-moving-average teacher.
Tokenizer
- The CNN tokenizer maps two-channel complex CSI through four residual blocks into 192-dimensional tokens, while layer normalization is used to avoid collapse.The tokenizer begins with a 3 × 3 convolutional stem; the second residual module downsamples by a 2× stride.Source: C. Network Topology and Training… · p. 5
Predictor and auxiliary heads
- The JEPA predictor uses a 4-layer, 8-head transformer module followed by a shallow 2-layer, 4-head ViT-style stack that predicts in the 192-dimensional latent space.Auxiliary heads include a symmetric CNN masked-recovery decoder, an amplitude-recovery branch, and a token-conditional DiT-style diffusion head.Source: C. Network Topology and Training… · p. 5
Downstream reconstruction
- The downstream regressor embeds the rough channel estimate on a 32×16 grid and uses the predicted latent prior as conditioning for six transformer-style reconstruction blocks.Each block combines self-attention, cross-attention to the latent prior, and a feed-forward layer.Source: C. Network Topology and Training… · p. 5
Pilot and noise setup
- Pilots are sampled every eight positions per subband with a random offset, and Gaussian noise is applied only to observed pilot entries.Noise power is sampled uniformly in NMSE range [0, 1].Source: C. Network Topology and Training… · p. 5
Training settings
- Training uses batch size 32, weight decay 10^-4, a warmup-and-cosine learning-rate schedule, and teacher-student EMA coefficient μ = 0.95.The learning rate warms from 0 to 10^-4 over the first 10% of training.Source: C. Network Topology and Training… · p. 5
D. Baseline Methods
The comparison includes prediction-only, current-pilot-only, direct end-to-end training, and an idealized current-latent reconstruction upper bound. These baselines distinguish the value of historical prediction, the JEPA two-stage training scheme, and predicted versus perfectly available latents.
Simple Prediction
- Simple Prediction uses past channel matrices to recover next-step channels, without current pilot-based estimation.The paper notes that this prediction experiment fails to converge.Source: D. Baseline Methods · p. 5
Simple Estimation
- Simple Estimation reconstructs the target channel using only the current sparse, noisy pilot-based estimate, without past information..Source: D. Baseline Methods · p. 5
Direct Training
- Direct Training ablates the JEPA two-stage training scheme by training the overall network directly on the final reconstruction target for 400 epochs.Source: D. Baseline Methods · p. 5
Current-Latent Reconstruction
- Current-Latent Reconstruction is an unrealistic idealized upper bound that maps the current channel directly to the latent domain before downstream reconstruction.It uses the same downstream network and represents the investigated scheme with perfect prediction.Source: D. Baseline Methods · p. 5
E. Effect of Different Regularization Items in Pretraining
The pretraining comparisons examine covariance and recovery-oriented regularization, while downstream reconstruction is used to assess the usefulness of the learned latents. Moderate covariance regularization performs best among the reported pretraining settings, but latent prediction NMSE alone does not establish downstream gains.
Downstream evaluation
- Fig. 5 compares downstream reconstruction NMSE across noise levels and reports that the simple prediction experiment fails to converge.Source: Figure 5 · p. 6
Evaluation interpretation
- Latent-prediction NMSE is reported because it aligns more directly with downstream wireless reconstruction quality, but good prediction NMSE is necessary rather than sufficient for downstream gains.The paper also notes that phase-level prediction is difficult in this dataset because of position shift.Source: E. Effect of Different Regularization… · p. 6
Pretraining results
- λcov = 0.1 gives the best result, and the results indicate that moderate covariance regularization is necessary.Source: E. Effect of Different Regularization… · p. 6
Recovery-oriented regularization
- Pretraining performance degrades from diffusion to amplitude recovery to direct recovery, consistent with finer-detail restoration making JEPA prediction harder.Diffusion grounding is least harmful among the grounding variants because its denoising target can be partly solved from the corrupted sample without encoding all fine-grained CSI details.Source: E. Effect of Different Regularization… · p. 6
F. Channel Prediction and Reconstruction Results
The best JEPA-based method outperforms simple-estimation and direct-training baselines, while the ablations favor predictive objectives without recovery-oriented regularization. The conclusion attributes the approach's utility to latent field prediction combined with sparse real-time pilots.
F. Channel Prediction and Reconstruction Results
- Across the noise range, “cov” and “cov + diff” yield the best prediction-aided reconstruction performance and perform similarly.Adding amplitude recovery or direct recovery consistently increases reconstruction NMSE; a moderate covariance term without recovery-oriented regularization is most effective.Source: F. Channel Prediction and… · p. 6
- The best JEPA-based method clearly outperforms simple-estimation and direct-training baselines in channel reconstruction.The simple prediction baseline fails to converge under random disturbance in the locations.Source: F. Channel Prediction and… · p. 6
V. CONCLUSION
- The framework predicts a compact latent propagation field and combines the predicted state with sparse real-time pilots as a reconstruction prior.The conclusion says the latent representation emphasizes stable, predictable propagation structure while reducing sensitivity to unstable small-scale fading and phase-related details.Source: §V · p. 6
- CSI-domain grounding objectives do not necessarily improve reconstruction; preserving excessive fine-grained CSI details can reduce latent predictability and usefulness as a prior.Source: §V · p. 6