WiMamba: Linear-Scale Wireless Foundation Model

Probleemstelling

Foundation modellen (grote, taak-onafhankelijke modellen) hebben revolutionaire resultaten geboekt in natuurlijke taalverwerking en computer vision, en er is groeiende interesse in het toepassen hiervan op draadloze communicatiesystemen. Bestaande draadloze foundation modellen zijn echter grotendeels gebaseerd op transformer-architecturen. Deze architecturen hebben een kwadratische computatie- en geheugencomplexiteit ($O(N^2)$) ten opzichte van de lengte van de invoersequentie.

Dit vormt een groot obstakel voor de praktische implementatie in draadloze systemen, vooral bij het verwerken van grote kanalen (bijv. MIMO met veel antennes en subcarriers). De noodzaak om "key-value caches" te onderhouden tijdens inferentie en de kwadratische schaling van de attention-mechanismen leiden tot hoge latentie en een groot geheugengebruik, wat de schaalbaarheid beperkt.

Methodologie: WiMamba

De auteurs stellen WiMamba voor, een nieuw draadloos foundation model dat de recente Mamba-architectuur (gebaseerd op Selective State-Space Models of SSMs) gebruikt in plaats van self-attention. Dit biedt lineaire complexiteit ($O(N)$) voor sequentiemodellering.

De kerncomponenten van de methodologie zijn:

1. Bidirectionele Mamba Backbone:

   • Omdat draadloze kanaalmatrices (CSI) geen inherente causale tijdsrichting hebben (in tegenstelling tot taal), is een standaard causale Mamba niet ideaal. WiMamba gebruikt daarom een bidirectionele opzet met een voorwaartse en een achterwaartse Mamba-laag die parallel werken.
   • De output van beide stromen wordt geconcateneerd om een rijke representatie te vormen.

2. Adaptieve Tokenisatie en Granulariteit:

   • Om te kunnen omgaan met variërende CSI-dimensies (verschillende aantallen antennes en subcarriers) zonder de architectuur aan te passen, introduceert WiMamba een adaptief granulariteitsmechanisme.
   • De invoer (CSI-matrix) wordt opgesplitst in tokens van een vaste grootte ($L_e$). Door de grootte van deze tokens aan te passen, kan de gebruiker een afweging maken tussen resolutie (detail) en rekentijd.
   • Dit maakt het model schaalbaar voor verschillende resoluties zonder hertraining.

3. Taak-onafhankelijke Zelftoezicht Pre-training:

   • Het model wordt vooraf getraind op een groot, ongelabeld dataset van kanaalrealisaties ($D_{pre}$).
   • Er wordt gebruik gemaakt van een Masked Language Modeling (MLM) aanpak: een deel van de tokens wordt gemaskeerd en het model moet deze reconstrueren.
   • Een speciaal maskeringsbeleid zorgt ervoor dat zowel het reële als het imaginaire deel van de CSI-data consistent worden gemaskeerd.
   • Dit resulteert in een generieke, overdraagbare representatie van draadloze kanalen.

4. Inferentiepijplijn:

   • Voor downstream taken (zoals beam-predicatie of localisatie) wordt de voorgeprogrammeerde WiMamba gebruikt om de CSI om te zetten in een latentere presentatie (embeddings).
   • Alleen een kleine, taak-specifieke "head" (bijv. een MLP) hoeft te worden getraind met een beperkt gelabeld dataset.

Belangrijkste Bijdragen

• Introductie van WiMamba: Het eerste draadloze foundation model dat de Mamba-architectuur (SSM) gebruikt, waardoor lineaire schaling wordt bereikt in plaats van kwadratische schaling.
• Architecturale Innovatie: Een bidirectionele opzet die specifiek is ontworpen voor de niet-causale structuur van draadloze kanalen, gecombineerd met een adaptief tokenisatiemechanisme voor variërende invoergroottes.
• Efficiëntie: Het elimineren van de noodzaak voor key-value caches en het bereiken van lineaire complexiteit, wat leidt tot drastische verbeteringen in latentie en geheugengebruik.
• Pre-training Framework: Een robuust, zelftoezichtend pre-training protocol dat transferleerbaarheid garandeert voor diverse downstream taken.

Resultaten

Het model is geëvalueerd op het DeepMIMO-dataset voor vier downstream taken:

1. Line-of-Sight (LOS) vs. Non-LOS (NLOS) classificatie.
2. Beam-predicatie.
3. Kanaal-interpolatie.
4. Gebruikerslocalisatie.

Kernbevindingen:

• Prestaties: WiMamba presteert gelijkwaardig aan of beter dan transformer-based foundation modellen (zoals de Large Wireless Model - LWM) op alle vier de taken, ongeacht de patch-grootte (resolutie).
• Schaalbaarheid en Latentie:
  • Bij kleinere patches (hogere resolutie, langere sequenties) explodeert de latentie en het geheugengebruik van de transformer-architectuur door de kwadratische complexiteit.
  • WiMamba behoudt een constante, lage latentie en een laag geheugengebruik, ongeacht de invoergrootte.
  • In tests toonde WiMamba een drastische reductie in inferentie-tijd (bijv. van ~363ms naar ~23ms bij 4x4 patches) en geheugengebruik (van ~4800MB naar ~114MB) ten opzichte van de transformer-basislijn.

Betekenis en Conclusie

WiMamba bewijst dat Selective State-Space Models (SSMs) een krachtig alternatief zijn voor transformers in de context van draadloze communicatie. Het paper toont aan dat het mogelijk is om foundation modellen te bouwen die:

1. Schaalbaar zijn voor grote MIMO-systemen zonder kwadratische kosten.
2. Efficiënt zijn in termen van energie en hardware-gebruik (cruciaal voor randapparatuur in 6G).
3. Flexibel zijn in het verwerken van verschillende kanaalresoluties zonder herarchitecturering.

Dit opent de deur voor de praktische implementatie van geavanceerde, op AI gebaseerde draadloze systemen die real-time kunnen opereren op grote schaal, iets wat met huidige transformer-architecturen vaak onhaalbaar is.