MUSE-FM: Multi-Task Environment-Aware FM

• The paper introduces MUSE-FM, a transformer-based model that unifies heterogeneous wireless tasks through a prompt-guided encoder-decoder and scene-graph fusion.
• It addresses task and scenario uniformity challenges by incorporating multi-modal inputs and a shared GPT-2–style backbone for channel estimation, detection, and localization.
• Experimental results demonstrate improved performance and data efficiency across tasks, especially in low-data regimes and varied environmental scenarios.

Searching arXiv for the named paper and closely related multi-task/environment-aware works to ground the article. arxiv_search query: "MUSE-FM Multi-task Environment-aware Foundation Model for Wireless Communications" Multi-Task Environment-aware FM (MUSE-FM) is a transformer-based foundation model for wireless communications that is designed to unify multiple physical-layer tasks within a single framework while incorporating scenario information as a multi-modal input (Zheng et al., 2 Sep 2025). Its central objective is to address two limitations identified in prior wireless foundation-model designs: the lack of task uniformity when inputs and outputs differ substantially across tasks, and the lack of scenario uniformity when channel distributions shift across physical environments (Zheng et al., 2 Sep 2025). In the formulation reported for MUSE-FM, these issues are handled through a prompt-guided unified encoder-decoder pair, a scene-graph-based environment branch, and a shared GPT-2–style backbone that supports channel estimation, MIMO detection, multi-user precoding, channel decoding, and user localization in one model (Zheng et al., 2 Sep 2025).

1. Problem formulation and design rationale

MUSE-FM is motivated by the observation that existing wireless foundation models often cover only a narrow set of channel-related tasks, typically assume similar input-output formats, and commonly rely on task-specific encoders and decoders for each task or configuration (Zheng et al., 2 Sep 2025). This leads to large overhead and limited scalability when tasks vary in structure, such as regression versus classification, vector versus matrix input, or different pilot lengths, SNR values, user counts, and code lengths (Zheng et al., 2 Sep 2025). The model is therefore framed around a stronger notion of uniformity: one architecture should process heterogeneous data formats, heterogeneous objectives, and heterogeneous scenarios without redesigning the full network.

The scenario dimension is equally central. In the reported formulation, deployment in a new building layout or with a new base-station placement induces a channel-distribution shift that degrades models trained only on a single environment family (Zheng et al., 2 Sep 2025). MUSE-FM therefore treats environmental context as prior knowledge of the environment and channel distributions. The environment is not an auxiliary label attached post hoc; it is a primary input modality that is fused with wireless measurements inside the backbone (Zheng et al., 2 Sep 2025).

This multi-task formulation is also presented as a data-efficiency mechanism. When per-task data are limited, single-task models are described as prone to overfitting, whereas joint training across tasks can exploit inter-task relations and diverse datasets (Zheng et al., 2 Sep 2025). That claim is reflected in the paper’s reduced-data experiments, where the multi-task model improves over task-specific FM variants at low data fractions (Zheng et al., 2 Sep 2025).

2. Unified architecture and prompt-conditioned parameterization

MUSE-FM is organized around three components: a scene encoder, a prompt-guided unified data encoder-decoder pair, and an FM backbone implemented as a GPT-2–style transformer (Zheng et al., 2 Sep 2025). The environment branch processes a discretized scene graph; the data branch processes task-specific numerical inputs after preprocessing; and a textual task instruction controls the encoder and decoder weights through a hypernetwork (Zheng et al., 2 Sep 2025).

The prompt is a natural-language-like instruction containing the task identifier and key task parameters. Examples reported for the model include "Channel estimation", "MIMO detection", "Channel decoding", "Multi-user precoding", and "User localization", augmented with parameters such as pilot length, SNR, transmitting antenna number, data length, user number, encoded bit length, and information bit length (Zheng et al., 2 Sep 2025). Let the embedded instruction tokens be $X_t^{\mathrm{emb}}$. The hypernetwork maps them to the encoder and decoder parameters: $$\Theta = [\Theta_{en}, \Theta_{de}] = \mathrm{MLP}(X_t^{\mathrm{emb}}).$$ This is the core mechanism for task uniformity: the architecture of the encoder and decoder is fixed, but their weights are generated from the instruction (Zheng et al., 2 Sep 2025).

The raw task input $X_n^{\mathrm{input}}$ is first converted by a preprocessor into a common token feature size $2N_t$: $$X_n^{\mathrm{pre}} = \mathrm{Preprocessor}(X_n^{\mathrm{input}}).$$ For pilots, the feature dimension is zero-padded to $2N_t$; for detection and precoding inputs, batch normalization is applied directly; for decoding, the input vector is converted into a diagonal matrix (Zheng et al., 2 Sep 2025). The unified encoder then projects the preprocessed tokens into the transformer embedding space: $$X_n^{\mathrm{emb}} = \mathrm{Linear}(X_n^{\mathrm{pre}}; W_{en}, b_{en}),$$ where $W_{en} \in \mathbb{R}^{D \times 2N_t}$ and $b_{en} \in \mathbb{R}^{D \times 1}$ are generated from the prompt (Zheng et al., 2 Sep 2025).

Environmental context is encoded by a Vision Transformer operating on a discretized top-view scene graph. After scene and data embeddings are formed, they are concatenated with a learnable $[CLS]$ token: $$\Theta = [\Theta_{en}, \Theta_{de}] = \mathrm{MLP}(X_t^{\mathrm{emb}}).$$0 and processed by $$\Theta = [\Theta_{en}, \Theta_{de}] = \mathrm{MLP}(X_t^{\mathrm{emb}}).$$1 transformer blocks with LayerNorm, multi-head self-attention, and MLP sublayers with residual connections (Zheng et al., 2 Sep 2025). After the backbone, the decoder maps the retained data tokens back to a common output space: $$\Theta = [\Theta_{en}, \Theta_{de}] = \mathrm{MLP}(X_t^{\mathrm{emb}}).$$2 and a task-specific postprocessor extracts the final task output by slicing the decoded tensor or reading the $$\Theta = [\Theta_{en}, \Theta_{de}] = \mathrm{MLP}(X_t^{\mathrm{emb}}).$$3 token (Zheng et al., 2 Sep 2025).

This design replaces five separate encoder-decoder pairs with one prompt-conditioned encoder-decoder pair. In the reported comparison against a multi-task FM with task-specific encoders and decoders, the unified design slightly improves channel estimation, precoding, detection, and decoding, while localization is comparable (Zheng et al., 2 Sep 2025).

3. Task set, input-output interfaces, and optimization objectives

The model is evaluated on five physical-layer tasks in a multi-user MISO-OFDM mmWave setting with a $$\Theta = [\Theta_{en}, \Theta_{de}] = \mathrm{MLP}(X_t^{\mathrm{emb}}).$$4-antenna UPA at the base station, $$\Theta = [\Theta_{en}, \Theta_{de}] = \mathrm{MLP}(X_t^{\mathrm{emb}}).$$5 subcarriers, center frequency $$\Theta = [\Theta_{en}, \Theta_{de}] = \mathrm{MLP}(X_t^{\mathrm{emb}}).$$6 GHz, subcarrier spacing $$\Theta = [\Theta_{en}, \Theta_{de}] = \mathrm{MLP}(X_t^{\mathrm{emb}}).$$7 kHz, and $$\Theta = [\Theta_{en}, \Theta_{de}] = \mathrm{MLP}(X_t^{\mathrm{emb}}).$$8 users by default (Zheng et al., 2 Sep 2025). These tasks differ substantially in both their signal interfaces and optimization targets.

Task	Canonical input	Canonical output
Channel estimation	Received pilot OFDM matrix $$\Theta = [\Theta_{en}, \Theta_{de}] = \mathrm{MLP}(X_t^{\mathrm{emb}}).$$9	Estimated channel $X_n^{\mathrm{input}}$0
MIMO detection	$X_n^{\mathrm{input}}$1	Estimated symbols $X_n^{\mathrm{input}}$2
Multi-user precoding	Estimated channel $X_n^{\mathrm{input}}$3	Precoding matrix $X_n^{\mathrm{input}}$4
Channel decoding	$X_n^{\mathrm{input}}$5	Recovered bits $X_n^{\mathrm{input}}$6
User localization	Pilot signal $X_n^{\mathrm{input}}$7	Estimated position $X_n^{\mathrm{input}}$8

For channel estimation, the objective is

$X_n^{\mathrm{input}}$9

with MSE training loss and NMSE evaluation (Zheng et al., 2 Sep 2025).

For MIMO detection, the system model is

$2N_t$0

and the estimator solves

$2N_t$1

again using MSE training and NMSE evaluation (Zheng et al., 2 Sep 2025).

For multi-user precoding, the model optimizes the sum rate under a power constraint. With

$2N_t$2

the optimization is

$2N_t$3

with the negative sum rate used as the training loss (Zheng et al., 2 Sep 2025).

For channel decoding, MUSE-FM follows the ECCT preprocessing pipeline for polar codes. The model predicts a binary multiplicative noise representation and is trained with cross-entropy: $2N_t$4 with BER used for evaluation, reported as $2N_t$5 in the results (Zheng et al., 2 Sep 2025).

For user localization, the objective is

$2N_t$6

with MSE training and normalized error distance evaluation (Zheng et al., 2 Sep 2025).

Across tasks, the overall multi-task training loss is

$2N_t$7

where $2N_t$8 balances the distinct objective scales. Validation is tracked through

$2N_t$9

with task-specific validation metrics such as NMSE, negative sum rate, BER, and normalized localization error (Zheng et al., 2 Sep 2025).

4. Environment-aware representation and scene-conditioned inference

The environment branch is based on a discretized scene graph derived from a simulated 3D indoor scene constructed in Blender (Zheng et al., 2 Sep 2025). The reported environment consists of a $$X_n^{\mathrm{pre}} = \mathrm{Preprocessor}(X_n^{\mathrm{input}}).$$0 room with concrete floor, brick walls, $$X_n^{\mathrm{pre}} = \mathrm{Preprocessor}(X_n^{\mathrm{input}}).$$1–$$X_n^{\mathrm{pre}} = \mathrm{Preprocessor}(X_n^{\mathrm{input}}).$$2 internal walls, $$X_n^{\mathrm{pre}} = \mathrm{Preprocessor}(X_n^{\mathrm{input}}).$$3–$$X_n^{\mathrm{pre}} = \mathrm{Preprocessor}(X_n^{\mathrm{input}}).$$4 cylindrical obstacles, a fixed BS at $$X_n^{\mathrm{pre}} = \mathrm{Preprocessor}(X_n^{\mathrm{input}}).$$5, and randomly placed user equipments (Zheng et al., 2 Sep 2025). The top-view scene is discretized into a binary matrix $$X_n^{\mathrm{pre}} = \mathrm{Preprocessor}(X_n^{\mathrm{input}}).$$6, with $$X_n^{\mathrm{pre}} = \mathrm{Preprocessor}(X_n^{\mathrm{input}}).$$7 representing obstacles and $$X_n^{\mathrm{pre}} = \mathrm{Preprocessor}(X_n^{\mathrm{input}}).$$8 free space (Zheng et al., 2 Sep 2025).

The scene encoder splits $$X_n^{\mathrm{pre}} = \mathrm{Preprocessor}(X_n^{\mathrm{input}}).$$9 into $2N_t$0 patches, linearly projects them, adds positional embeddings, and processes them with a transformer encoder to produce scene tokens $2N_t$1 (Zheng et al., 2 Sep 2025). These tokens are concatenated with task-data tokens and enter the shared backbone, allowing scene tokens and wireless-signal tokens to attend to one another. This joint self-attention is the mechanism through which environment structure acts as prior information on channel and localization behavior (Zheng et al., 2 Sep 2025).

The reported quantitative comparison between task-specific FM(s) without scene input and environment-aware FM(s) with scene input is especially informative because it isolates the value of the environmental branch.

Task metric	FM(s)	Env-aware FM(s)
CE NMSE	0.1387	0.0982
Precoding rate	19.44	19.71
Detection NMSE	0.001829	0.001632
Decoding $2N_t$2	$2N_t$3	$2N_t$4
Localization error	0.1516	0.1233

The corresponding improvements are reported as 1.50 dB for channel estimation, 1.38% for sum rate, 0.49 dB for MIMO detection, 0.20 dB for decoding, and 18.67% for localization error (Zheng et al., 2 Sep 2025). In this formulation, environment information is therefore most consequential for tasks that depend directly on channel geometry or user position, especially channel estimation and localization.

5. Experimental results, scaling behavior, and configuration transfer

The training set comprises 2600 indoor scenarios, with 50 samples per scenario for a total of 130,000 samples; the split is 2000 training scenarios, 300 validation scenarios, and 300 test scenarios (Zheng et al., 2 Sep 2025). The FM backbone is the smallest GPT-2 variant with embedding dimension $2N_t$5, trained with Adam using $2N_t$6, initial learning rate $2N_t$7, cosine decay, batch size $2N_t$8, and 500 epochs on 4×RTX 4090 (Zheng et al., 2 Sep 2025).

On multi-user precoding, MUSE-FM is reported to achieve the highest sum rate across SNRs, with clear gains over ZF, WMMSE, CNN, shallow transformer, and the environment-aware task-specific FM in low-SNR regimes (Zheng et al., 2 Sep 2025). On MIMO detection, environment-aware FM(s) is the strongest method, while MUSE-FM remains close behind and is especially strong at low SNR, with the prompt’s explicit SNR prior cited as beneficial (Zheng et al., 2 Sep 2025). On user localization at $2N_t$9 dB, the fraction of samples with normalized error smaller than $$X_n^{\mathrm{emb}} = \mathrm{Linear}(X_n^{\mathrm{pre}}; W_{en}, b_{en}),$$0 is 64.13% for environment-aware FM(s), 60.38% for MUSE-FM, 49.27% for the transformer baseline, and 45.37% for the MLP baseline (Zheng et al., 2 Sep 2025).

For channel estimation, the reported NMSE values in dB show consistent superiority over LS, ReEsNet, and CE-ViT. At SNR $$X_n^{\mathrm{emb}} = \mathrm{Linear}(X_n^{\mathrm{pre}}; W_{en}, b_{en}),$$1 dB, the table reports $$X_n^{\mathrm{emb}} = \mathrm{Linear}(X_n^{\mathrm{pre}}; W_{en}, b_{en}),$$2 dB for MUSE-FM versus $$X_n^{\mathrm{emb}} = \mathrm{Linear}(X_n^{\mathrm{pre}}; W_{en}, b_{en}),$$3 dB for CE-ViT, $$X_n^{\mathrm{emb}} = \mathrm{Linear}(X_n^{\mathrm{pre}}; W_{en}, b_{en}),$$4 dB for ReEsNet, and $$X_n^{\mathrm{emb}} = \mathrm{Linear}(X_n^{\mathrm{pre}}; W_{en}, b_{en}),$$5 dB for LS (Zheng et al., 2 Sep 2025). For channel decoding, the reported $$X_n^{\mathrm{emb}} = \mathrm{Linear}(X_n^{\mathrm{pre}}; W_{en}, b_{en}),$$6 at $$X_n^{\mathrm{emb}} = \mathrm{Linear}(X_n^{\mathrm{pre}}; W_{en}, b_{en}),$$7 is 13.72 for MUSE-FM versus 13.82 for environment-aware FM(s), 12.32 for ECCT, and 7.69 for Hyper BP (Zheng et al., 2 Sep 2025). The decoding result is therefore competitive with the strongest FM baseline and above the classical and transformer baselines.

The data-scaling study clarifies the role of multi-task learning. With full data, task-specific environment-aware FM(s) is usually slightly stronger; with reduced data, MUSE-FM frequently surpasses it, especially for channel estimation, precoding, and localization (Zheng et al., 2 Sep 2025). At 10% training data for precoding, the paper reports 17.80 bps/Hz for MUSE-FM versus 17.28 bps/Hz for the task-specific FM; at 2% data, MUSE-FM improves performance across all tasks relative to task-specific FM, which is presented as evidence that multi-task learning reduces overfitting in data-limited regimes (Zheng et al., 2 Sep 2025).

A central scalability experiment changes the pilot length from 4 during training to 6 at test time. In task-specific encoder-decoder designs, this change breaks the input interface and requires a new encoder. In MUSE-FM, the updated textual instruction and the same preprocessor-hypernetwork mechanism enable zero-shot evaluation, where performance is reported to be better than LS and ReEsNet and comparable to CE-ViT, and few-shot adaptation with only 1% of data for 5 epochs, where the adapted model surpasses all baselines (Zheng et al., 2 Sep 2025). This experiment is the clearest demonstration of configuration-level transfer through prompt-conditioned encoder-decoder generation.

6. Position within the broader research landscape and reported limitations

MUSE-FM belongs to a broader family of systems that combine multi-task learning with explicit or implicit environment modeling, but its particular combination of prompt-generated signal interfaces and scene-graph conditioning is specific to wireless physical-layer processing (Zheng et al., 2 Sep 2025). In Wi-Fi sensing, Uni-Fi formalizes multi-task sensing through shared event sets, shared feature sets, and a forward-model network, arguing that the multi-task inverse model is not simply the union of single-task inverse models (Li et al., 16 Jan 2026). Wimuse addresses multi-task CSI sensing through a shared encoder, task-specific residual adaptors, and knowledge distillation from single-task experts for gesture recognition, indoor localization, and user identification (Zhang et al., 2021). AM-FM approaches ambient intelligence through WiFi as a foundation-model problem trained on 9.2 million unlabeled CSI samples and evaluated on nine downstream tasks, emphasizing cross-task and cross-environment reuse of a shared backbone (Zhu et al., 4 Feb 2026). These works share the premise that environment variability and task coupling should be modeled jointly rather than absorbed into isolated per-task systems.

Related environment-aware formulations appear in other domains as well. FM-EAC uses feature-based models and an enhanced actor-critic design for multi-task UAV control in dynamic environments, with environment features supplied to critics and task-specific reward streams handled by multiple critics (Zhou et al., 17 Dec 2025). MUSE-VAE decomposes long-term trajectory prediction into macro and micro stages, where environment maps, goal prediction, and fine-grained trajectory generation form a coarse-to-fine environment-aware hierarchy (Lee et al., 2022). By contrast, the vision paper titled MUSE uses timestep embeddings as native task steering for one-step dense prediction; despite the shared acronym, that mechanism is distinct from the prompt-guided encoder-decoder and scene-graph fusion used in wireless MUSE-FM (Zhou et al., 29 Jun 2026).

The limitations reported for MUSE-FM are substantial and define its current scope. The experiments are simulation-based, with the wireless channel and scene data generated from 2600 indoor scenarios rather than field measurements (Zheng et al., 2 Sep 2025). Environmental input is restricted to a simple 2D discretized scene graph, whereas richer modalities such as point clouds, RGB images, or semantic maps are left for future work (Zheng et al., 2 Sep 2025). The backbone is a moderate-size GPT-2 small model, and the authors explicitly note that larger models and parameter-efficient adaptation methods such as partial tuning, LoRA, or layer-norm tuning remain open directions (Zheng et al., 2 Sep 2025). They also identify end-to-end transceiver optimization, broader task coverage, and real-world validation as future steps (Zheng et al., 2 Sep 2025).

Taken as a whole, MUSE-FM is best understood as a wireless foundation-model architecture that treats task specification, signal-interface adaptation, and environmental context as coequal components of the inference process. Its distinctive contribution is not only multi-task training, but a concrete method for unifying heterogeneous wireless tasks through prompt-conditioned encoder-decoder generation and scene-aware fusion inside a shared transformer backbone (Zheng et al., 2 Sep 2025).