⚡ electrical engineering

WiMamba: Linear-Scale Wireless Foundation Model

WiMamba is a novel wireless foundation model that leverages the linear-time Mamba architecture and a self-supervised pre-training framework to achieve scalable, low-latency inference while matching or outperforming existing transformer-based models in representational expressivity and downstream task performance.

Original authors: Tomer Raviv, Nir Shlezinger

Published 2026-03-30

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Tomer Raviv, Nir Shlezinger

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to understand the "weather" of a wireless signal. In the world of Wi-Fi and 5G, this "weather" is called the Channel State Information (CSI). It's a massive, complex map of how radio waves bounce off walls, furniture, and people to get from your router to your phone.

For a long time, the best way to teach robots to read this map was using a type of AI called a Transformer (the same tech behind chatbots like me). But Transformers have a big problem: they are like a student who tries to read every single word of a book and compare it to every other word to understand the story. If the book is short, that's fine. But if the book is a massive encyclopedia (which wireless channels often are), the student gets overwhelmed, takes forever to read, and runs out of memory.

Enter WiMamba.

The Problem: The "Library" Bottleneck

Think of the old Transformer models as a librarian who insists on walking to every single bookshelf in a giant library to check if a book is related to the one you're holding.

The Issue: If the library has 1,000 books, the librarian makes 1,000 trips. If the library has 1,000,000 books, they make a million trips. This is called "quadratic complexity." It's slow, expensive, and impossible to scale for huge wireless networks.

The Solution: The "Smart Reader" (WiMamba)

The authors, Tomer Raviv and Nir Shlezinger, built a new kind of AI called WiMamba. Instead of a librarian who checks everything against everything else, WiMamba uses a "Smart Reader" based on a new architecture called Mamba.

Here is how it works, using simple analogies:

1. The "Selective Memory" (The SSM)

Imagine you are listening to a long podcast.

The Old Way (Transformer): To understand the ending, you pause and mentally replay the entire podcast from the beginning, comparing every sentence to the current one.
The WiMamba Way (Mamba): You listen to the podcast and keep a smart, selective memory. You remember the important parts, forget the boring parts, and update your understanding as you go. You don't need to re-read the whole thing; you just carry the "gist" forward.

In technical terms, this is called a Selective State-Space Model (SSM). It allows the AI to process the signal in linear time. If the signal is 10 times bigger, it takes 10 times longer to process, not 100 times longer. It's like switching from a snail that has to visit every house on a street to a high-speed train that just stops at the stations.

2. The "Adaptive Lens" (Granularity)

Wireless signals come in all shapes and sizes. Sometimes you have a small room with few devices; other times, you have a stadium with thousands of antennas.

The Old Way: You had to build a different camera for every room size. If you wanted to see more detail, you had to rebuild the whole camera.
The WiMamba Way: WiMamba has an adaptive lens. It can zoom in or out. If the signal is huge, it looks at it in "chunks" (like looking at a map of a country by looking at states). If the signal is small, it looks at the fine details. It can handle any size without needing to be retrained or rebuilt.

3. The "Universal Translator" (Foundation Model)

Usually, AI models are trained for one specific job (like "predicting where a user is standing"). If you want it to do something else (like "predict which direction the signal is coming from"), you have to retrain it from scratch.

WiMamba is a Foundation Model. Think of it like a universal translator who learns the "language" of radio waves first.

Pre-training: The model reads millions of different radio maps (without needing labels) to learn the "grammar" of how signals behave.
Downstream Tasks: Once it knows the language, you can ask it to do specific tasks (like "find the user" or "classify the signal") with just a tiny bit of extra training. It's like teaching a polyglot a new dialect; they pick it up much faster than someone starting from zero.

The Results: Fast, Light, and Smart

The authors tested WiMamba against the best Transformer models on four different tasks (like finding a user's location or predicting signal beams).

Performance: WiMamba was just as good as the Transformers. It understood the signals perfectly.
Speed & Memory: This is where it shines.
- Latency (Speed): WiMamba was 10 to 20 times faster.
- Memory: It used 40 to 80 times less memory.

The Bottom Line

WiMamba is like upgrading from a heavy, slow-moving tank to a sleek, high-speed drone. It does the same job (understanding complex wireless signals) but does it so efficiently that it can run on smaller devices and handle massive networks without breaking a sweat.

By swapping the "compare-everything" approach of Transformers for the "smart-memory" approach of Mamba, the authors have created a wireless AI that is finally ready for the real world, where speed and efficiency are everything.

1. Problem Statement

The paper addresses the computational bottlenecks associated with applying Foundation Models (FMs) to wireless communication systems, specifically for processing Channel State Information (CSI).

Current Limitations: Existing wireless foundation models predominantly rely on Transformer architectures. While powerful, Transformers utilize self-attention mechanisms that scale quadratically ( $O(N^2)$ ) with respect to the input sequence length. This results in prohibitive inference latency and memory consumption when processing large-scale MIMO channels (high numbers of antennas and subcarriers).
The Gap: There is a need for a wireless foundation model that maintains the representational expressivity of Transformers but offers linear-time complexity and lower memory footprints to enable practical, low-latency deployment in real-world wireless systems.

2. Methodology: WiMamba Architecture

The authors propose WiMamba, a wireless foundation model built upon the Mamba architecture, which utilizes Selective State-Space Models (SSMs) instead of self-attention.

A. Core Architecture

Selective SSMs: WiMamba replaces the attention mechanism with SSMs. Unlike Transformers that store Key-Value caches, SSMs process sequences via recurrent-like evolution equations with linear complexity relative to sequence length.
Bidirectional Backbone: Since wireless channel matrices (antenna $\times$ subcarrier) lack an inherent temporal causal order (unlike text or time-series), a standard causal Mamba is unsuitable. WiMamba employs a bidirectional Mamba backbone consisting of parallel forward and backward Mamba layers to capture non-causal spatial-frequency dependencies.
Adaptive Tokenization: To handle varying CSI dimensions ( $N$ $N$ antennas, $M$ $M$ subcarriers) without retraining, the model uses an adaptive granularity mechanism.
- The input CSI matrix is vectorized and split into tokens based on a configurable embedding dimension ( $L_e$ ).
- This allows the model to operate at different resolutions (e.g., $4\times4$ or $8\times8$ patches) seamlessly, balancing the trade-off between representational detail and computational cost.

B. Pre-training Framework

Task-Agnostic & Self-Supervised: WiMamba is pre-trained on a large, unlabeled dataset of CSI realizations ( $D_{pre}$ ) without task-specific labels.
Masked Modeling: The model uses a Masked Language Modeling (MLM) objective adapted for wireless data.
- A subset of tokens is masked.
- Crucially, the masking strategy ensures that if a real component token is masked, its corresponding imaginary component token is also masked to preserve signal structure.
- The model reconstructs the masked tokens using a linear decoder, learning rich, transferable channel embeddings.

C. Inference Pipeline

Downstream Adaptation: For specific tasks (e.g., beam prediction, localization), the pre-trained backbone is frozen.
Representation Extraction: The model outputs embeddings that can be aggregated in three ways:
1. Class Token: Aggregates global channel information.
2. Flattened Patches: Preserves localized spatial-frequency details.
3. Mean-Pooled: Captures overall channel structure.
Lightweight Heads: A small, task-specific Multi-Layer Perceptron (MLP) is trained on top of these embeddings using a small labeled dataset ( $D_{task}$ ), requiring minimal fine-tuning.

3. Key Contributions

First Mamba-based Wireless FM: WiMamba is the first work to demonstrate the effectiveness of Mamba (SSM) architectures as a foundation model for wireless channels, moving beyond the Transformer dominance.
Linear Scalability: By leveraging SSMs, WiMamba achieves linear complexity ( $O(N \cdot M)$ ) with respect to channel dimensions, eliminating the quadratic bottleneck of attention mechanisms.
Adaptive Granularity: The introduction of a resolution-adaptive tokenization mechanism allows a single model to handle varying MIMO configurations without architectural changes or retraining.
Dedicated Wireless Pre-training: A novel self-supervised pre-training framework tailored to the complex-valued, non-causal nature of CSI data.

4. Experimental Results

The model was evaluated on the DeepMIMO dataset across four downstream tasks: LOS/NLOS classification, beam prediction, channel interpolation, and user localization.

Performance Accuracy: WiMamba achieves performance comparable to or exceeding Transformer-based baselines (specifically the Large Wireless Model, LWM). In scenarios with long sequence lengths (small patch sizes), the performance gap is negligible, proving SSMs can match Transformer expressivity.
Efficiency Gains:
- Latency: WiMamba demonstrates dramatic latency reductions. For a $4\times4$ patch size, latency dropped from ~363ms (Transformer) to ~23ms (WiMamba).
- Memory: Peak GPU memory usage was reduced by nearly an order of magnitude (e.g., from ~4836 MB to ~114 MB for $4\times4$ patches).
- Scalability: While Transformer latency and memory spiked as patch sizes decreased (sequence length increased), WiMamba maintained stable, low resource consumption due to its linear scaling.

5. Significance

This work represents a paradigm shift in wireless AI infrastructure. By proving that Selective State-Space Models can serve as effective foundation models for communications, the authors provide a pathway to deploy large-scale, generalizable AI models on resource-constrained wireless devices (e.g., base stations, user equipment). WiMamba solves the critical trade-off between model capability and computational feasibility, enabling real-time, high-fidelity channel understanding for next-generation 6G networks.