La ampliación de la computación en tiempo de inferencia se ha convertido en el principal impulsor del rendimiento del modelo de lenguaje grande (LLM), cambiando el enfoque arquitectónico hacia la eficiencia de la inferencia junto con la calidad del modelo. Si bien las arquitecturas basadas en Transformer siguen siendo el estándar, su complejidad computacional cuadrática y sus requisitos de memoria lineal crean importantes cuellos de botella en la implementación. Un equipo de investigadores de la Universidad Carnegie Mellon (CMU), la Universidad de Princeton, Together AI y Cartesia AI han presentado Mamba-3, un modelo que aborda estas limitaciones mediante un diseño de “inferencia primero”.
Mamba-3 se basa en el marco del Modelo de espacio de estados (SSM), introduciendo tres actualizaciones metodológicas principales: discretización trapezoidal exponencial, actualizaciones de estado de valores complejos y una formulación de múltiples entradas y múltiples salidas (MIMO).
1. Discretización exponencial-trapezoidal
Los modelos de espacio de estados son sistemas de tiempo continuo que deben discretizarse para procesar secuencias discretas. Iteraciones anteriores como Mamba-1 y Mamba-2 utilizaron una heurística de primer orden conocida como discretización ‘exponencial-Euler’. Mamba-3 reemplaza esto con discretización trapezoidal exponencial, que proporciona una aproximación precisa de segundo orden de la integral estado-entrada.
Técnicamente, esta actualización cambia la recurrencia discreta de una actualización de dos períodos a una actualización de tres períodos:
$$h_{t}=e^{\Delta_{t}A_{t}}h_{t-1}+(1-\lambda_{t})\Delta_{t}e^{\Delta_{t}A_{t}}B_{t-1}x_{t-1}+\lambda_{t}\Delta_{t}B_{t}x_{t}$$
Esta fórmula equivale a aplicar una convolución de ancho 2 dependiente de los datos en el Btxt de entrada de estado dentro de la recurrencia central. En las pruebas empíricas, esta convolución implícita, combinada con sesgos B y C que se pueden aprender, permite que Mamba-3 funcione de manera efectiva sin las convoluciones causales cortas externas que normalmente requieren los modelos recurrentes.
2. Modelos de espacio de estados de valores complejos y el ‘truco RoPE’
Una limitación de los modelos lineales de valor real es su incapacidad para resolver tareas de “seguimiento de estado”, como determinar la paridad de secuencias de bits. Este fracaso surge de restringir los valores propios de la matriz de transición a números reales, que no pueden representar la dinámica “rotacional” requerida para tales tareas.
Mamba-3 incorpora SSM de valor complejo para resolver esto. El equipo de investigación estableció una equivalencia teórica entre los SSM complejos discretizados y los SSM de valor real que utilizan incrustaciones posicionales rotativas (RoPE) dependientes de datos en las proyecciones B y C.
Al utilizar el ‘truco RoPE’, el modelo aplica rotaciones agregadas dependientes de datos a lo largo de intervalos de tiempo. Esto permite a Mamba-3 resolver tareas sintéticas como paridad y aritmética modular, donde Mamba-2 y sus variantes de valor real no funcionan mejor que las adivinanzas aleatorias.
3. Formulación de múltiples entradas y múltiples salidas (MIMO)
Para abordar la ineficiencia del hardware de la decodificación vinculada a la memoria, Mamba-3 pasa de una recurrencia de entrada única y salida única (SISO) a una estructura de múltiples entradas y múltiples salidas (MIMO).
En la decodificación SSM estándar, la intensidad aritmética es de aproximadamente 2,5 operaciones por byte, muy por debajo del régimen de computación de las GPU modernas como la H100. MIMO aumenta el rango R de las proyecciones de entrada y salida (Bt E RNR y xt E RPR), transformando la actualización de estado de un producto externo a una multiplicación matriz-matriz.
Este cambio aumenta la decodificación de FLOP hasta 4 veces en relación con Mamba-2 en un tamaño de estado fijo. Debido a que el cálculo adicional se superpone con la E/S de memoria existente necesaria para la actualización del estado, MIMO mejora la calidad del modelado y la perplejidad al tiempo que mantiene una latencia de decodificación similar a la del reloj de pared.
Arquitectura y Normalización
El bloque Mamba-3 sigue el diseño estilo Llama, alternándose con bloques SwiGLU. Las mejoras clave incluyen:
Normalización BC/QK: la normalización RMS se aplica a las proyecciones B y C, reflejando QKNorm en Transformers. Esto estabiliza el entrenamiento y permite la eliminación del RMSNorm posterior a la puerta utilizado en versiones anteriores. Sesgos específicos de la cabeza: se agregan sesgos aprendibles por canal a los componentes B y C después de la normalización para inducir un comportamiento similar a la convolución. Integración híbrida: cuando se utiliza en arquitecturas híbridas (entrelazando capas lineales con autoatención), se descubrió que la adición de una RMSNorm agrupada previa a la puerta mejora la generalización de la longitud en las tareas de recuperación.
Resultados y Eficiencia
Las evaluaciones se realizaron en el conjunto de datos FineWeb-Edu en cuatro escalas de modelo (180 millones a 1,5 mil millones).
Rendimiento posterior: en la escala de 1.500 millones, la variante SISO Mamba-3 supera a Mamba-2 y Gated DeltaNet (GDN). La variante MIMO (R=4) mejora aún más la precisión descendente promedio en 1,2 puntos con respecto a la línea base SISO. Frontera de Pareto: Mamba-3 logra una perplejidad previa al entrenamiento comparable a Mamba-2 mientras usa solo la mitad del tamaño del estado (por ejemplo, Mamba-3 con un tamaño de estado de 64 coincide con Mamba-2 con 128). Rendimiento del kernel: Los kernels Triton optimizados (para precarga) y CuTe DSL (para decodificación) garantizan que los componentes matemáticos adicionales sigan siendo livianos. Los kernels SISO Mamba-3 demuestran una latencia más baja que los kernels Mamba-2 y GDN lanzados en la configuración estándar BF16.
Mamba-3 demuestra que los ajustes fundamentales al punto de vista del modelo de espacio de estados pueden cerrar la brecha entre la eficiencia subcuadrática teórica y la capacidad práctica de modelado.
Consulte el documento, la página de GitHub y los detalles técnicos. Además, no dude en seguirnos en Twitter y no olvide unirse a nuestro SubReddit de más de 120.000 ML y suscribirse a nuestro boletín. ¡Esperar! estas en telegrama? Ahora también puedes unirte a nosotros en Telegram.