Source-linked AI summary
Scaling of Wireless Foundation Models via Representation Diversity and Multi-Branch Architectures
Ahmed Mohamed, Ahmed Aboulfotouh, Hatem Abou-Zeid
TL;DR
Wireless self-supervised learning has less consistent scaling gains when pretraining data are limited. This paper combines complementary reconstructive, predictive, and contrastive representations through fusion and a shared-trunk, multi-branch architecture. Fusion outperforms the evaluated single-objective encoders on all six tasks at equal output dimension, using roughly one-third of their parameters.
Problem
Different downstream tasks depend on different signal properties, motivating objective diversity as an alternative to relying on a single self-supervised objective.
Source: §I · p. 1Method
The paper fuses frozen representations from independently pretrained reconstructive, predictive, and contrastive encoders and introduces a jointly trained multi-branch model with a shared trunk and objective-specific branches.
Source: §I · pp. 1–2, §V · p. 12Results
At matched output embedding dimension, three-objective fusion outperforms single-objective encoders on all six tasks while using roughly one-third of their parameters.
Source: A. Contributions · p. 2Takeaways & Limitations
The findings support objective diversity as an effective strategy for scaling wireless foundation models, with much of the multi-branch model’s benefit retained within a single-encoder parameter budget.
Source: §V · p. 12
Abstract
from arXiv · showhide
Wireless foundation models learn representations from unlabeled radio signals for reuse across downstream tasks. Scaling model capacity is a common strategy for learning richer representations and improving downstream performance. However, its gains are less consistent in wireless self-supervised learning when pretraining data are limited. We investigate objective diversity as an alternative scaling axis: different self-supervised objectives emphasize different signal properties, and combining their representations can preserve more information to enable diverse tasks. We develop a fusion framework that combines frozen representations from independent encoders trained through reconstructive, predictive, and contrastive learning. This provides a reference for the benefits of diversity, but requires the maintenance of multiple encoders. To retain these benefits within the parameter budget of a standard single-objective encoder, we introduce a jointly trained multi-branch architecture with a shared trunk and objective-specific branches. We pretrain on a heterogeneous corpus of spectrogram and channel state information data and evaluate six downstream tasks that span communication, sensing, and positioning. At an equal output embedding dimension, fusion outperforms the evaluated single-objective encoders on all six tasks while using roughly one-third of their parameters. The multi-branch architecture retains much of the fusion benefit within a single-encoder parameter budget. Representation analyses indicate complementary contributions across objectives, with much of the added benefit retained in components orthogonal to the reconstruction representation subspace. These findings support objective diversity as an effective strategy for scaling wireless foundation models.
I. INTRODUCTION
Wireless pretraining data are limited, and increasing model capacity alone may not improve downstream performance. The paper investigates objective diversity as an alternative scaling strategy, using fusion and a multi-branch architecture to combine complementary representations.
I. INTRODUCTION
- Wireless corpora are small and fragmented, and increasing capacity alone can improve pretraining-task performance without broadening downstream-useful information.Prior work found that larger wireless self-supervised models do not necessarily improve downstream performance, a trend this work examines and confirms.Source: §I · p. 1
- Reconstruction, prediction, and contrastive objectives emphasize different signal properties, motivating combinations that may preserve features useful across downstream tasks.Temporal and spectral characteristics support modulation classification, while inter-antenna phase relationships support angle-of-arrival estimation.Source: §I · p. 1
A. Contributions
- The fusion framework combines representations learned independently under multiple objectives, while the multi-branch design uses a shared trunk and objective-specific branches within a standard encoder’s size budget.Fusion provides a reference for objective diversity but requires maintaining separate encoders.Source: §I · pp. 1–2
- With matched output embedding dimensions, three-objective fusion outperforms single-objective encoders on all six tasks while using roughly one-third of their parameters.The comparisons are conducted under limited pretraining data.Source: A. Contributions · p. 2
- The multi-branch architecture recovers most fusion gains within a single standard encoder’s size budget and remains stable across evaluated shared-depth configurations.It uses shared early transformer blocks and objective-specific branches in deeper layers.Source: A. Contributions · p. 2
- Representation analyses find structural differences among encoders, complementary contributions from all three objectives, and added benefits largely retained outside the MAE representation subspace.Pairwise comparisons also show that the strongest objective combination varies across downstream tasks.Source: A. Contributions · p. 2
B. Notation
Wireless foundation-model work has largely used single-objective pretraining, while combining objectives remains relatively unexplored. The paper frames objective diversity as an alternative to capacity scaling and seeks to retain fusion benefits within a single-encoder budget.
II. RELATED WORK
- Wireless foundation-model research has largely committed to one objective family, with reconstruction the most developed line and latent prediction and contrastive learning forming other approaches.Prior shared-encoder branching designs targeted supervised multi-task learning or multiple augmented views of one objective.Source: §II · pp. 2–3
- The work combines reconstructive MAE, predictive JEPA, and contrastive SimCLR, compares diversity with width scaling at matched parameter budgets, and proposes a shared-trunk, objective-specific-branch design.Prior wireless work combined at most two objectives; the proposed branching design aims to recover most fusion gains without parameter overhead over a single-objective encoder.Source: §II · p. 3
A. Problem Formulation
- For a fixed dataset size, the paper proposes diversity across sub-representations as a more effective scaling axis than increasing capacity under one objective.The formulation posits that a D-dimensional composition of diverse representations can outperform a D-dimensional representation trained with any single objective.Source: A. Problem Formulation · p. 5
- The formulation identifies the practical challenge as retaining fusion benefits without maintaining M independent encoders.The paper attributes diminishing returns from capacity increases to the limited size and diversity of available wireless pretraining corpora.Source: A. Problem Formulation · pp. 4–5
B. Fusion Framework
The fusion framework independently pretrains three encoders with complementary self-supervised objectives, freezes them, and concatenates their pooled outputs. Its input representation processes antenna channels separately while allowing varying channel counts.
B. Fusion Framework
- Three independently pretrained, frozen ViT encoders use MAE reconstruction, JEPA latent prediction, and SimCLR temporal-view invariance; their pooled outputs form a 768-dimensional fused representation.Each encoder has 8 blocks and hidden dimension 256; MAE and JEPA use mean pooling, while SimCLR uses its learned [CLS] token.Source: B. Fusion Framework · p. 5
1) Input Representation:
- Each antenna channel is patchified independently, and a shared linear projection embeds the resulting single-channel patches into a common token sequence.The sequence allows the encoder to model interactions across tokens from different channels.Source: 1) Input Representation · p. 5
- Uniformly sampling an arbitrary subset of available antenna channels during pretraining exposes the model to varying channel counts without structural modification.The effective sequence length varies across training iterations, preventing reliance on a specific channel count.Source: 1) Input Representation · p. 5
2) Masked Autoencoder (MAE):
The three pretraining objectives impose distinct learning signals: MAE reconstructs masked input, JEPA predicts masked latent representations from context, and SimCLR learns invariance across temporal crops.
2) Masked Autoencoder (MAE):
- MAE masks 80% of input tokens and uses a lightweight decoder to reconstruct masked patches, minimizing mean squared error over masked positions.The sparse visible tokens encourage context-dependent representations that capture signal structure.Source: 2) Masked Autoencoder (MAE) · p. 5
3) Joint Embedding Predictive Architecture (JEPA):
- JEPA uses six contiguous region masks, encodes context tokens, and predicts the masked regions’ latent representations using a predictor and an EMA target encoder.Its coherent spatial blocks encourage higher-level predictive dependencies while latent-space prediction avoids committing to low-level reconstruction details.Source: 3) Joint Embedding Predictive… · pp. 5–6
4) Simple Contrastive Learning of Representations (SimCLR):
- SimCLR contrasts two random temporal crops of each signal, learning features that are stable across temporal shifts and globally discriminative across views.The crops retain the same frequency extent and antenna channels while spanning 30% to 100% of the signal’s time axis.Source: 4) Simple Contrastive Learning of… · p. 6
C. Multi-Branch Architecture
The multi-branch architecture combines a shared trunk with objective-specific branches to retain representation diversity within a single-encoder size budget. It trains MAE, JEPA, and SimCLR objectives jointly and concatenates their branch representations at inference.
C. Multi-Branch Architecture
- A shared trunk followed by objective-specific branches realizes three-objective representation diversity within a single dedicated encoder.Early blocks are shared, while deeper blocks specialize for each objective.Source: C. Multi-Branch Architecture · pp. 6–7
- 8 transformer blocks match the single-method encoder budget: 2 shared blocks plus 3 branches of 2 blocks each.The stated configuration adds zero parameter overhead relative to an 8-block single-method encoder.Source: C. Multi-Branch Architecture · p. 6
- All three objectives process the same input in parallel, and their losses jointly update branch blocks and the shared trunk.Each objective masks or augments the input, processes it through shared and dedicated branch blocks, then computes its loss.Source: C. Multi-Branch Architecture · p. 7
- The JEPA target encoder is an EMA of the composed shared-trunk and JEPA-branch mapping, with its momentum scheduled from τ_min to τ_max.This keeps prediction targets aligned with the evolving shared and JEPA-specific parameters.Source: C. Multi-Branch Architecture · p. 7
- At inference, the input passes through the shared trunk once, then all three branch outputs are pooled and concatenated into a 3d-dimensional representation.The resulting dimension matches the fusion framework's output.Source: C. Multi-Branch Architecture · p. 7
D. Pretraining Data
Pretraining uses four unlabeled datasets spanning spectrogram and CSI modalities, totaling approximately 25,000 samples. The corpus combines heterogeneous signal sources and antenna configurations.
D. Pretraining Data
- Approximately 25,000 samples across four unlabeled datasets form a corpus spanning spectrogram and CSI modalities.The datasets include roughly 3,200 proprietary over-the-air spectrogram samples, 6,000 RF-fingerprinting spectrogram samples, 15,000 five-channel indoor 5G NR CSI samples, and 840 three-channel WiFi CSI samples.Source: D. Pretraining Data · pp. 7–8
- The corpus includes heterogeneous antenna configurations ranging from one to five channels.The proprietary spectrogram dataset contains WiFi, LTE, Bluetooth, 5G NR, and ISM signals captured across sub-6 GHz center frequencies at 10–60 MHz sampling rates.Source: D. Pretraining Data · p. 8
- All inputs are resized to 224 × 224 and normalized using per-channel means and standard deviations computed over the respective pretraining dataset.Spectrograms are first converted to decibel scale using 20 log10(·).Source: D. Pretraining Data · p. 8
E. Downstream Tasks
The evaluation covers six downstream tasks across communication, sensing, and positioning. Three tasks use CSI and three use spectrogram inputs.
E. Downstream Tasks
- Six downstream tasks span communication, sensing, and positioning domains.The evaluation is divided evenly between three CSI tasks and three spectrogram tasks.Source: E. Downstream Tasks · p. 8
- Three tasks operate on CSI inputs and three on spectrogram inputs.The paper describes this as an assessment across both modalities.Source: E. Downstream Tasks · p. 8
- The task suite is intended to provide a comprehensive assessment across both input modalities.Source: E. Downstream Tasks · p. 8
1) Spectrogram Tasks:
The spectrogram task suite evaluates signal-standard recognition, GNSS jamming-type classification, and interference-modulation classification.
1) Spectrogram Tasks:
- RF Signal Classification identifies the wireless standard from a spectrogram using over-the-air captures of multiple radio technologies.The task uses the CommRad RF dataset.Source: 1) Spectrogram Tasks · p. 8
- GNSS Jamming Detection classifies jamming type from spectrograms converted from synthesized raw IQ recordings.The simulated scenarios include narrowband, AM, FM, chirp, and DME jamming.Source: 1) Spectrogram Tasks · p. 8
- Interference Classification identifies interference modulation type from spectrograms of LTE IQ recordings captured under clean and DSSS-interference conditions.The recordings are over the air.Source: 1) Spectrogram Tasks · p. 8
2) CSI Tasks:
The CSI-task suite covers positioning and two DeepMIMO classification problems, with task-specific metrics for classification and location error.
2) CSI Tasks:
- 5G NR Positioning regresses user-equipment location from five-channel CSI measurements in the outdoor scenario of the 5G CSI dataset.Source: 2) CSI Tasks · p. 8
- Beam Prediction uses DeepMIMO CSI to classify the optimal beam index.Source: 2) CSI Tasks · p. 8
- LoS/NLoS Classification uses DeepMIMO data to classify whether the transmitter–receiver channel contains a line-of-sight path.Source: 2) CSI Tasks · p. 8
- Classification tasks use per-class accuracy, the mean of each class recall, while positioning uses mean Euclidean error between predicted and true 3D coordinates.Source: 2) CSI Tasks · p. 8
IV. RESULTS AND DISCUSSION
The results compare objective diversity with single-objective scaling and test a shared-trunk, multi-branch alternative. Fusion performs strongly across the six tasks, while the multi-branch design retains much of its benefit within a single-encoder parameter budget.
A. Representation Diversity versus Single-Objective Scaling
- Increasing single-objective width from 256 to 768 dimensions gives inconsistent returns: MAE RFS accuracy falls from 85.9% to 84.2%.The same width increase improves MAE on beam prediction and interference classification.Source: A. Representation Diversity versus… · p. 8
- At 512-d output, the best pairwise fusion per task matches or exceeds the best single-objective encoder on five of six tasks, using roughly half the parameters.Single-objective JEPA retains the best interference accuracy, and the strongest pair varies by task.Source: A. Representation Diversity versus… · p. 8, Table IV · p. 9
- At matched 768-d output, fusion outperforms all three single-objective encoders on all six tasks while using roughly one-third of their parameters.Fusion also beats the strongest 256-d baseline, MAE, on all six tasks; Table III reports test performance across downstream tasks.Source: A. Representation Diversity versus… · p. 8, Table III · p. 9
B. Multi-Branch Performance across Shared Depths
- Across shared-depth configurations, RFS and beam accuracy vary by only 0.5 percentage points, while positioning error ranges from 1.63 to 1.73 m.The tested configurations use different numbers of shared and per-branch blocks.Source: B. Multi-Branch Performance across… · p. 9, Table V · p. 10
- Within a single 256-d, eight-block encoder’s parameter budget, the multi-branch model produces a 768-d representation and exceeds the best single-method 768-d encoder’s mean performance on four of six tasks.Its largest accuracy gain is +2.0 percentage points on interference classification; performance is lower on LoS/NLoS detection and jamming classification.Source: B. Multi-Branch Performance across… · p. 9
- Table V reports a JEPA + MAE + SimCLR multi-branch shared-depth ablation (768-d; mean ± std over 20 seeds); overhead versus 256-d, 8-block single-objective encoders is (S + 3 × B − 8)/8.Source: Table V · p. 10
IV. RESULTS AND DISCUSSION
- Pairwise linear CKA compares encoder representations per downstream dataset; lower values indicate greater structural differences, and 1 indicates perfect similarity.Source: §IV · p. 8, Figure 3 · p. 10
C. Pairwise Complementarity
Pairwise results show that the strongest objective combination varies by task, while representation analyses find distinct structures and useful contributions across objectives. Much of the added benefit lies outside MAE’s representation subspace.
C. Pairwise Complementarity
- MAE + SimCLR is the strongest pair overall, with the highest RFS (88.6%) and Jamming accuracy (97.0%) among 512-d configurations; JEPA + MAE leads Positioning (1.86 m) and Beam accuracy (62.9%).JEPA + SimCLR is consistently the weakest pair; one possible explanation is that MAE’s reconstruction preserves signal details complementary to prediction or contrastive features.Source: C. Pairwise Complementarity · p. 10
D. Representation Similarity across Objectives
- Across five datasets, cross-method CKA ranges from 0.52 for MAE–JEPA on RFS to 0.84 for JEPA–SimCLR on Jamming, indicating distinct representation geometries.JEPA–SimCLR CKA is 0.61 on Positioning, showing that even the more aligned pair retains task-dependent divergence.Source: D. Representation Similarity across… · p. 10
E. Contributions of Individual Objectives
- Probe weight mass remains at least 23% for every encoder slice on every task, with JEPA largest on Positioning (43.1%) and MAE largest on Jamming (39.0%).The task-dependent distribution indicates the probes use features from all three objectives.Source: E. Contributions of Individual… · p. 10, Figure 4 · p. 11
- Removing MAE causes the largest leave-one-out degradation on every task, including a 26.2% relative increase in positioning error and a 10.5% relative drop in beam prediction accuracy.Removing any encoder degrades performance on at least three of six tasks, although JEPA removal has negligible or slightly positive effects on RFS, LoS/NLoS, and Jamming.Source: E. Contributions of Individual… · p. 10, Figure 4 · p. 11
F. Subspace Analysis of Complementary Information
- For both SimCLR and JEPA, the component orthogonal to MAE’s subspace retains nearly all or most of the full-embedding gain on several tasks, whereas the parallel component adds effectively nothing.The parallel component’s gains are within noise of zero on every task and sometimes fall slightly below the MAE-only baseline.Source: F. Subspace Analysis of Complementary… · p. 11, Figure 5 · p. 11
V. CONCLUSION AND FUTURE WORK
- The conclusion supports objective diversity as a scaling strategy when wireless pretraining data are limited, and identifies larger corpora, adaptive loss weighting, branch partitioning, more objectives, and more modalities as future directions.Source: §V · p. 12