Beam: Un Modelo MoE Escaso de 501B Parámetros para Carga de Trabajo de Codificación y Agentes

El desarrollo de modelos de lenguaje grandes (LLMs) ha demostrado que la escala de parámetros y datos de entrenamiento es crucial para el rendimiento. Sin embargo, el costo computacional de la inferencia de modelos densos de miles de millones de parámetros es prohibitivo para muchas aplicaciones. Beam aborda este problema fundamental de la computación al emplear una arquitectura Mixture-of-Experts (MoE) escasa, que permite un gran número total de parámetros (501B) mientras mantiene un número significativamente menor de parámetros activos (23B) durante la inferencia. Esto se traduce en una mayor eficiencia computacional por token, haciendo que modelos de gran capacidad sean más accesibles y rentables para cargas de trabajo específicas como la codificación y las tareas de agentes.

La tesis central es que la combinación de una arquitectura MoE escasa con un entrenamiento de Reinforcement Learning (RL) a gran escala y una curación de datos de alta calidad puede producir modelos de lenguaje de frontera con una eficiencia de inferencia superior. Esto es particularmente relevante en un contexto donde la demanda de inteligencia artificial en el borde y en aplicaciones empresariales requiere un equilibrio entre capacidad y costo operativo. La inversión en infraestructura de RL de alto cómputo y algoritmos para manejar la inestabilidad de la política asíncrona a escala es un pilar clave de esta estrategia.

Arquitectura del Sistema

Beam se basa en una arquitectura Mixture-of-Experts (MoE) escasa, donde solo un subconjunto de expertos se activa para cada token de entrada, lo que reduce el costo computacional de la inferencia. La arquitectura incorpora atención local y global entrelazada, expertos enrutados de grano fino, un flujo residual controlado y múltiples formas de balanceo de carga para asegurar una utilización estable y equilibrada de los expertos. Se utiliza un balanceo de carga sin pérdidas auxiliares, con decaimiento coseno de las actualizaciones de sesgo de expertos para reducir las perturbaciones de enrutamiento, y balanceo a nivel de secuencia para datos fuera de la distribución de preentrenamiento.

El sistema de entrenamiento se divide en preentrenamiento, midtraining y Reinforcement Learning (RL). El preentrenamiento utiliza 23.8 billones de tokens de alta calidad, con una curación de datos centrada en código, explicaciones técnicas y conocimiento científico/matemático. La infraestructura de preentrenamiento incluye un scheduler consciente de la topología basado en Kubernetes, un sistema de ciclo de vida de nodos con monitoreo continuo y un sistema de detección de corrupción de datos silenciosa (SDC) con rebobinados y reinicios semiautónomos. El RL a gran escala se ejecuta en una plataforma asíncrona que coordina el flujo de experiencia y las actualizaciones del modelo, permitiendo la ejecución independiente de la generación de rollouts, la ejecución de herramientas, la evaluación de resultados y la actualización del modelo. Se utilizan gradientes de política asíncronos, con algoritmos para mitigar la inestabilidad causada por la obsolescencia de la política y el desajuste numérico entre motores de entrenamiento e inferencia. La distribución de pesos del modelo se realiza jerárquicamente sobre RoCE y NVLink para actualizaciones rápidas.

Fundamentos Teóricos

La arquitectura Mixture-of-Experts (MoE) tiene sus raíces en trabajos académicos que exploran cómo combinar múltiples modelos especializados para mejorar el rendimiento general, como se ve en el paper "Adaptive Mixture of Local Experts" de Jacobs et al. (1991). La idea de la escasez en redes neuronales para mejorar la eficiencia computacional también ha sido un tema de investigación recurrente, buscando reducir la complejidad sin sacrificar la capacidad. La aplicación de Reinforcement Learning (RL) a gran escala para el entrenamiento de LLMs se basa en principios de aprendizaje por refuerzo, donde un agente aprende a tomar decisiones en un entorno para maximizar una señal de recompensa, un campo con fundamentos en trabajos como los de Sutton y Barto (1998) en "Reinforcement Learning: An Introduction".

Los desafíos de estabilidad en el entrenamiento de RL asíncrono a gran escala, particularmente la "policy staleness" (obsolescencia de la política), son bien conocidos en la literatura de RL distribuido. La gestión de este problema y el desajuste entre entrenamiento e inferencia son áreas activas de investigación que buscan hacer que los algoritmos de RL sean más robustos y escalables. La curación de datos de alta calidad y la deduplicación, así como la gestión de la repetición de datos, se conectan con la investigación en muestreo óptimo y la prevención del sobreajuste en el aprendizaje automático, temas explorados en numerosos papers sobre generalización y robustez de modelos.