Megalodon: Pretrenamiento e Inferencia Eficiente de Modelos de Lenguaje con Longitud de Contexto Ilimitada

Resumen

La complejidad cuadrática y la débil extrapolación de longitud de los Transformers limitan su capacidad para escalar a secuencias largas, y aunque existen soluciones subcuadráticas como la atención lineal y los modelos de espacio de estados, empíricamente tienen un rendimiento inferior al de los Transformers en eficiencia de preentrenamiento y precisión en tareas posteriores. Presentamos Megalodon, una arquitectura neuronal para el modelado eficiente de secuencias con longitud de contexto ilimitada. Megalodon hereda la arquitectura de Mega (media móvil exponencial con atención gated), e introduce además múltiples componentes técnicos para mejorar su capacidad y estabilidad, incluyendo la media móvil exponencial compleja (CEMA), la capa de normalización por paso de tiempo, el mecanismo de atención normalizada y el pre-norm con configuración residual de dos saltos. En una comparación controlada cara a cara con Llama2, Megalodon logra una mejor eficiencia que Transformer en la escala de 7 mil millones de parámetros y 2 billones de tokens de entrenamiento. Megalodon alcanza una pérdida de entrenamiento de 1.70, situándose a medio camino entre Llama2-7B (1.75) y 13B (1.67). Código: https://github.com/XuezheMax/megalodon

English

The quadratic complexity and weak length extrapolation of Transformers limits their ability to scale to long sequences, and while sub-quadratic solutions like linear attention and state space models exist, they empirically underperform Transformers in pretraining efficiency and downstream task accuracy. We introduce Megalodon, a neural architecture for efficient sequence modeling with unlimited context length. Megalodon inherits the architecture of Mega (exponential moving average with gated attention), and further introduces multiple technical components to improve its capability and stability, including complex exponential moving average (CEMA), timestep normalization layer, normalized attention mechanism and pre-norm with two-hop residual configuration. In a controlled head-to-head comparison with Llama2, Megalodon achieves better efficiency than Transformer in the scale of 7 billion parameters and 2 trillion training tokens. Megalodon reaches a training loss of 1.70, landing mid-way between Llama2-7B (1.75) and 13B (1.67). Code: https://github.com/XuezheMax/megalodon

Megalodon: Pretrenamiento e Inferencia Eficiente de Modelos de Lenguaje con Longitud de Contexto Ilimitada

Megalodon: Efficient LLM Pretraining and Inference with Unlimited Context Length

Resumen

Support