Implementación de Mamba en un archivo de PyTorch

Artículo original ↗ Discusión en Hacker News ↗

Se está elogiando una implementación compacta de PyTorch en un solo archivo de la nueva arquitectura neuronal Mamba como una alternativa clara y educativa al código oficial, altamente optimizado. Los participantes explican cómo los modelos de espacio de estados estructurados de Mamba buscan rivalizar con los Transformers ofreciendo modelado de secuencias en tiempo lineal, menor uso de VRAM y una inferencia más rápida, al tiempo que debaten compromisos en legibilidad, trucos de rendimiento y herramientas (desde Fortran hasta Julia y Mojo). La conversación destaca lo rápido que evoluciona la jerga y las familias de modelos en ML, y comparte recursos para quienes quieren entender dónde encaja Mamba en el panorama más amplio de los grandes modelos de lenguaje.

Implementación educativa en un solo archivo

  • Muchos valoran la Mamba de PyTorch en un solo archivo como una herramienta de aprendizaje.
  • En comparación con el código oficial CUDA, muy optimizado, esta versión se considera mucho más fácil de entender.
  • El número de líneas y de archivos se consideran importantes aquí porque el objetivo es la claridad conceptual, no la velocidad.

Implementaciones y lenguajes alternativos

  • Se menciona una implementación de inferencia de Mamba/SSM basada en Fortran; se elogia Fortran por su sintaxis de arrays similar a NumPy, su velocidad compilada, su fácil paralelización y su buen soporte para álgebra lineal, pero se critica por lo torpe que resulta para tareas no numéricas.
  • Algunos sugieren considerar lenguajes más nuevos como Julia o Mojo, pero quien defiende Fortran acepta la “batalla cuesta arriba”.
  • Se destaca otra biblioteca que abstrae el código compartido de transformer para que modelos como BERT, LLaMA y MPT quepan en unas ~100 líneas.

Rendimiento y paralelización

  • Algunos comentaristas señalan que el bucle secuencial selective_scan puede paralelizarse usando un par de llamadas de PyTorch y que todas las salidas pueden calcularse con un solo einsum.
  • Otros observan que esto podría perjudicar la legibilidad, que es el propósito principal del repositorio; se sugieren rutas paralelas opcionales o mantener la versión simple como comentarios.

Qué son Mamba / SSM y por qué interesan

  • Mamba se presenta como una arquitectura basada en modelos de espacio de estados (SSM) que busca competir con los transformers: modelado de secuencias en tiempo lineal en lugar de atención cuadrática.
  • La idea clave: hacer que algunos parámetros del SSM dependan de la entrada (un mecanismo de “selección”) mientras se mantiene lineal la transición latente, lo que permite tanto expresividad como algoritmos paralelos eficientes.
  • Se describe como una continuación de trabajos previos de SSM/long-convolution (HIPPO, S4, Hyena, etc.), añadiendo ingredientes que faltaban como el gating dependiente de la entrada.

Jerga, nombres y velocidad del campo

  • Varios participantes tienen dificultades con la jerga acelerada y los nuevos nombres de modelos (Mamba, RetNet, RWKV, etc.), y califican el espacio de impulsado por modas y muy marcado por branding.
  • Otros sostienen que los nombres cortos son útiles y que la comprensión llega al “escuchar por encima” las conversaciones de la comunidad (HN, subreddits, boletines).
  • Se recomiendan recursos como glosarios, referencias en el README, introducciones en blogs y videos explicativos.

Capacidades, eficiencia y preguntas abiertas

  • Entre las afirmaciones se incluyen:
    • Mamba entrena más rápido e infiere mucho más rápido que transformers de tamaño similar, con menor uso de VRAM por token.
    • Puede ser más eficiente en cómputo/muestras y estar bien adaptado para contextos muy largos y otros dominios secuenciales.
  • Algunos informan que Mamba usa alrededor del 60% de la VRAM de RetNet (en una prueba concreta) y que es más barato por token que los transformers.
  • Otros señalan que:
    • Las aplicaciones aguas abajo siguen siendo inmaduras en comparación con los modelos basados en atención.
    • No está claro cómo se compara Mamba con RetNet y RWKV a gran escala; las pruebas con longitudes de contexto y las implementaciones eficientes siguen evolucionando.
    • Una persona afirma que los modelos Mamba actuales son “menos coherentes que GPT-2”, otra discrepa rotundamente y dice que son mejores; no se resuelve la cuestión.

Relación con RNNs / LSTMs

  • Una pregunta recurrente es en qué se diferencia Mamba de las RNNs/LSTMs.
  • La explicación ofrecida: es esencialmente una RNN lineal en el espacio latente (por paralelización y estabilidad), con transiciones dependientes de la entrada que aportan la capacidad expresiva tradicionalmente dada por las no linealidades.
  • Se menciona una limitación práctica: a diferencia de los transformers, empaquetar eficientemente secuencias de distinta longitud para entrenamiento puede ser más difícil.

Estilo de código, herramientas y proceso de investigación

  • Las opiniones sobre einops/einsum están divididas: algunos los encuentran elegantes y eficientes; otros dicen que perjudican la legibilidad.
  • Hay entusiasmo por los “esquemas de código” mínimos y autocontenidos como una forma de acelerar la investigación y reducir la “complejidad de Kolmogórov” de la experimentación.
  • Las implementaciones en un solo archivo son elogiadas como contrapeso a las bases de código de ML extensas e interdependientes.