WiMamba: Linear-Scale Wireless Foundation Model

1. Problema y Motivación

Los modelos fundacionales (foundation models) han demostrado un gran éxito en procesamiento de lenguaje natural y visión por computadora, permitiendo la transferencia de representaciones aprendidas a múltiples tareas. Sin embargo, su aplicación a los sistemas inalámbricos, específicamente para la representación del estado del canal (CSI), enfrenta desafíos críticos:

• Complejidad Computacional Cuadrática: Los modelos existentes basados en Transformers utilizan mecanismos de atención que escalan cuadráticamente ($O(N^2)$) con la longitud de la secuencia de entrada (número de tokens). Esto genera una alta latencia de inferencia y un consumo masivo de memoria, lo que dificulta su despliegue en canales de gran escala (muchas antenas y subportadoras).
• Limitaciones de los Enfoques Actuales:
  • Adaptación de modelos genéricos: Usar modelos de lenguaje/visión preentrenados requiere pipelines de pre/post-procesamiento complejos y mantiene una alta carga computacional.
  • Modelos dedicados actuales: Aunque se han desarrollado modelos específicos para canales inalámbricos, casi todos dependen de arquitecturas Transformer, heredando sus limitaciones de escalabilidad.
• Requisitos No Satisfechos: Se necesitan modelos que sean agnósticos a la tarea, escalables a diferentes dimensiones de CSI y capaces de ofrecer inferencia de baja latencia con un footprint de memoria reducido.

2. Metodología: WiMamba

El artículo propone WiMamba, un modelo fundacional inalámbrico diseñado para superar las limitaciones de los Transformers mediante el uso de la arquitectura Mamba (Modelos de Espacio de Estado Selectivos o SSMs).

A. Arquitectura Base

• Mamba Bidireccional: A diferencia de los datos temporales o de lenguaje que tienen una dirección causal, las matrices de canales inalámbricos (antenas x subportadoras) no tienen una dirección temporal inherente. WiMamba utiliza un backbone bidireccional compuesto por capas de Mamba que procesan la secuencia en direcciones hacia adelante y hacia atrás simultáneamente.
• Modelos de Espacio de Estado Selectivos (SSMs): Reemplaza la atención por SSMs selectivos. Estos permiten un procesamiento en tiempo lineal ($O(N)$) respecto a la longitud de la secuencia y eliminan la necesidad de almacenar cachés de claves-valores (key-value caches) durante la inferencia, reduciendo drásticamente el uso de memoria.

B. Mecanismo de Tokenización Adaptativa

Para manejar CSI de dimensiones variables sin reentrenar o modificar la arquitectura:

• Se introduce un mecanismo de granularidad adaptativa. El canal se tokeniza en bloques de diferentes tamaños ($L_e$), permitiendo controlar el compromiso entre resolución (detalle espacial-frecuencial) y costo computacional.
• Se soportan múltiples resoluciones de tokenización dinámicamente, ajustando la longitud de la secuencia de entrada según las restricciones de latencia y memoria.

C. Entrenamiento y Pre-entrenamiento

• Pre-entrenamiento Auto-supervisado: Se utiliza un marco agnóstico a la tarea basado en el Modelado de Lenguaje enmascarado (Masked Language Modeling).
• Máscara Inteligente: Se enmascaran tokens del canal (componentes reales e imaginarios) y el modelo aprende a reconstruirlos. Esto fuerza al modelo a aprender representaciones transferibles de la estructura subyacente del canal sin necesidad de etiquetas específicas de tarea.
• Pipeline de Inferencia: Para tareas específicas (clasificación, regresión), se utiliza el backbone pre-entrenado para extraer embeddings latentes, que luego son procesados por cabezas de red neuronal (MLP) específicas para la tarea, entrenadas solo con datos limitados.

3. Contribuciones Clave

1. Primera Aplicación de Mamba a Modelos Fundacionales Inalámbricos: WiMamba es, según los autores, el primer trabajo que demuestra la eficacia de la arquitectura Mamba como base para modelos fundacionales en comunicaciones, logrando una complejidad sub-cuadrática.
2. Escalabilidad Lineal: Logra una complejidad de inferencia y memoria que escala linealmente con las dimensiones del canal ($N$ y $M$), a diferencia de la escala cuadrática de los Transformers.
3. Mecanismo de Granularidad Adaptativa: Permite operar sobre canales de diferentes resoluciones (tamaños de parche) sin reentrenamiento, ofreciendo flexibilidad operativa.
4. Marco de Pre-entrenamiento Agnóstico: Un enfoque de auto-supervisión diseñado específicamente para la naturaleza compleja y no temporal de los datos CSI.

4. Resultados Experimentales

Los autores evaluaron WiMamba utilizando el dataset DeepMIMO en cuatro tareas de descarga:

1. Clasificación LOS vs. NLOS.
2. Predicción de haces (beam prediction).
3. Interpolación de canales.
4. Localización de usuarios.

Hallazgos Principales:

• Rendimiento de Tarea: WiMamba logra un rendimiento comparable o superior a los modelos basados en Transformers (como el modelo LWM de referencia) en todas las tareas, tanto en escenarios de clasificación como de regresión.
• Eficiencia Computacional:
  • Latencia: WiMamba muestra reducciones dramáticas en la latencia de inferencia. Por ejemplo, en parches de $4\times4$, la latencia de WiMamba fue de 23.70 ms frente a 363.59 ms del Transformer.
  • Memoria: El consumo de memoria pico en GPU se redujo en casi un orden de magnitud (de ~4836 MB a ~114 MB en el caso de parches pequeños).
• Estabilidad: Mientras que el rendimiento del Transformer se degrada significativamente al aumentar la resolución (longitud de secuencia) debido a la atención cuadrática, WiMamba mantiene un rendimiento estable y eficiente.

5. Significado e Impacto

El trabajo de WiMamba es significativo porque resuelve el cuello de botella de escalabilidad que ha impedido el despliegue práctico de modelos fundacionales en sistemas inalámbricos de gran escala.

• Viabilidad en Tiempo Real: Al reducir la complejidad de cuadrática a lineal, WiMamba hace factible la implementación de modelos fundacionales en entornos con restricciones estrictas de latencia y memoria, como estaciones base masivas MIMO.
• Paradigma de Diseño: Establece un nuevo estándar para el diseño de modelos fundacionales en comunicaciones, demostrando que los SSMs selectivos pueden igualar la expresividad de los Transformers con una fracción del costo computacional.
• Generalización: La capacidad de aprender representaciones transferibles sin etiquetas específicas y adaptarse a diferentes resoluciones de hardware o escenarios de canal posiciona a WiMamba como una solución robusta para la próxima generación de sistemas de comunicación inteligentes.

En resumen, WiMamba demuestra que es posible construir modelos fundacionales inalámbricos que no solo son precisos, sino también computacionalmente eficientes y escalables, superando las limitaciones inherentes de la arquitectura Transformer en el dominio de las comunicaciones inalámbricas.