TurboMamba combina Mamba y Transformer para modelos de lenguaje más rápidos e inteligentes
Una nueva arquitectura híbrida mejora la velocidad y la precisión combinando dos tipos de diseños de inteligencia artificial.
Una nueva forma de acelerar los modelos de lenguaje
Los investigadores han creado TurboMamba, un nuevo diseño para modelos de lenguaje que funciona más rápido tanto durante el entrenamiento como al generar texto. El trabajo fue publicado en arXiv por Hanxin Zhang 1.
La idea central es combinar dos tipos conocidos de arquitecturas de inteligencia artificial: Mamba y Transformer. Mamba (también llamada Mamba-3 en este trabajo) es un modelo muy eficiente para procesar largas secuencias de texto. Transformer es la arquitectura detrás de la mayoría de los modelos de lenguaje más populares hoy en día, conocido por su fuerte capacidad de predicción. TurboMamba intenta obtener lo mejor de ambos mundos.
Cómo funciona
Esta es la versión sencilla: un modelo global de Mamba se ejecuta una vez para todo un bloque de tokens —un grupo de palabras o partes de palabras que el modelo procesa junto—. Un Transformer más pequeño luego observa el contexto global que Mamba produjo, más los tokens anteriores del bloque actual, para predecir el siguiente token.
Tres pequeños ajustes mejoran aún más las predicciones locales. Primero, el modelo puede acceder directamente al último token del bloque anterior. Segundo, aplica una pequeña corrección a las puntuaciones para el siguiente token. Tercero, aprende a mezclar las predicciones del Transformer con la frecuencia con la que aparecieron ciertos tokens recientes.
Esto significa que el trabajo pesado lo realiza el rápido modelo Mamba, mientras que el Transformer más pequeño se enfoca únicamente en ajustar las predicciones dentro de cada bloque.
Mejor y más rápido
Los investigadores probaron doce versiones a pequeña escala de diferentes arquitecturas en un conjunto de datos público llamado FineWeb-Edu, manteniendo similar la cantidad de parámetros y el tiempo de entrenamiento entre todas. Los modelos de referencia incluían Transformer, Mamba-3, Jamba y Zamba. Para Jamba y Zamba, los investigadores reemplazaron los módulos originales por Mamba-3 1.
TurboMamba quedó primero. Logró la menor pérdida promedio de predicción en las pruebas entre los doce modelos. Su error de predicción fue un 9,4 % menor que el de Mamba-3, un 7,3 % menor que la versión modificada de Jamba y un 18,3 % menor que la versión modificada de Zamba 1.
Las ganancias de velocidad también fueron grandes. El entrenamiento fue 2,60 veces más rápido que Mamba-3, 2,81 veces más rápido que la versión Jamba y 4,82 veces más rápido que la versión Zamba 1. Para generar texto, TurboMamba fue 3,12 veces más rápido que Mamba-3, 3,53 veces más rápido que la versión Jamba y 6,21 veces más rápido que la versión Zamba en un lote de 32 secuencias con consignas de 256 tokens y 128 tokens generados .
Lo que importa aquí
El resultado muestra que combinar la velocidad de Mamba con la precisión de Transformer es una dirección prometedora. Las ganancias se midieron en modelos pequeños, por lo que el enfoque deberá probarse a escalas mayores. Aun así, la idea de que un modelo global rápido haga el trabajo pesado mientras un modelo más pequeño se encarga de las decisiones locales podría influir en cómo se construyan los modelos de lenguaje del futuro.
Los datos clave
- TurboMamba combina un modelo global Mamba-3 con un Transformer más pequeño para la predicción de tokens
- Logró un 9,4 % menos de pérdida de predicción que Mamba-3 en FineWeb-Edu
- El rendimiento durante el entrenamiento fue 2,60 veces más rápido que Mamba-3 y 4,82 veces más rápido que Zamba
- La aceleración en inferencia alcanzó ser 6,21 veces más rápido que la versión Zamba
- Las pruebas usaron doce modelos a pequeña escala con parámetros y presupuestos de entrenamiento comparables