Fast Wireless Foundation Models with Early-Exits
This addresses the computational bottleneck of deploying large wireless foundation models in 6G networks, offering a practical solution for efficient inference.
The paper proposes an early-exit framework for wireless foundation models that attaches lightweight heads at intermediate encoder stages, achieving up to 93% reduction in FLOPs while improving accuracy on unseen out-of-distribution tasks compared to full-depth execution.
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.