Fast Wireless Foundation Models with Early-Exits

2026-06-28

AI Summary

This work addresses the high computational cost of wireless foundation models, which stems from the need to fully execute the backbone network, and their performance degradation on out-of-distribution tasks. To tackle these issues, the authors propose an efficient inference framework featuring early exits: the encoder is frozen, and lightweight task-specific heads are attached at multiple intermediate layers, enabling adaptive selection of inference depth based on input demands. A key insight is that features from intermediate layers exhibit superior transferability, and a fixed exit strategy consistently outperforms conventional dynamic routing methods that rely on per-sample difficulty estimation. Experiments demonstrate that the proposed approach achieves higher accuracy than the full model on unseen tasks while substantially reducing computational overhead—cutting FLOPs by up to 93%.

Abstract

While wireless foundation models (FMs) are demonstrating strong potential to enable AI-Native 6G networks, their high computational cost remains a critical barrier to deployment. The large computational cost stems from the rigid, full-depth execution of the FM backbone for every task, a process we show is not only inefficient but can also degrade performance on unseen out-of-distribution (OOD) tasks. In this paper, we propose a novel early-exit FM framework that attaches lightweight, per-task heads, at the most appropriate exit-stage of a frozen wireless FM encoder, enabling variable-depth inference tailored to each task's preferred representation depth. Our results demonstrate that these intermediate-layer features not only speed-up inference significantly (up to 93% fewer FLOPs), but also provide more transferable representations that exceed the full encoder accuracy on unseen tasks. We further demonstrate that a simple fixed-exit strategy per task is more effective than traditional early-exiting policies that route different samples to different exits based on their perceived difficulty levels.

Problem

• wireless foundation models
• computational cost
• out-of-distribution tasks
• AI-Native 6G
• model deployment

Innovation

• early-exit
• wireless foundation models
• variable-depth inference
• out-of-distribution generalization
• AI-Native 6G