El problema fundamental que ERNIE 5.5 Flash aborda es la creciente demanda de modelos de lenguaje grandes (LLMs) que sean tanto potentes como computacionalmente eficientes. A medida que los LLMs escalan en tamaño y contexto, los costos de entrenamiento y, crucialmente, de inferencia se vuelven prohibitivos. Esto se manifiesta en la necesidad de gestionar el consumo de memoria del KV cache y la carga computacional de la atención en secuencias largas, así como la eficiencia del entrenamiento distribuido a escala de miles de GPUs.

La solución propuesta es una reevaluación holística de la arquitectura Transformer, integrando técnicas de dispersión y optimización en cada componente, desde la atención hasta las capas MoE y la cabeza de salida. El objetivo es proporcionar un modelo base que sea un punto de partida asequible para la adaptación a dominios específicos, el ajuste de instrucciones y la investigación de agentes, sin comprometer significativamente la capacidad. Esto es particularmente relevante ahora, dada la explosión de aplicaciones de IA que requieren LLMs personalizados y la necesidad de democratizar el acceso a modelos de alto rendimiento.

Arquitectura del Sistema

ERNIE 5.5 Flash se basa en 43 capas Transformer, con un tamaño oculto de 4096 y 64 cabezas de atención. La arquitectura de atención es híbrida: la mayoría de las capas (37 de 43) utilizan Heavily Compressed Attention (HCA), que comprime cada 128 tokens en una entrada de caché única, reduciendo drásticamente el tamaño del KV cache. Las capas restantes (6) utilizan Multi-head Latent Attention (MLA), que mantiene acceso a nivel de token al contexto completo. Las capas MLA emplean Sparse Attention (DSA) con un indexador ligero para seleccionar los top-k tokens más relevantes, manteniendo el costo computacional constante con la longitud del contexto. Además, se utiliza Low-Rank Gated Attention y Virtual Head Attention (VHA) para mejorar la calidad y la estabilidad.

Las capas Feed-Forward, excepto la primera, implementan Multi-Router Latent MoE. En lugar de un único router, utilizan varios routers sigmoidales inicializados independientemente que puntúan a los expertos conjuntamente, mejorando la selección y la robustez frente a ruido. Cada capa MoE proyecta tokens a un espacio latente de menor dimensión (4096 a 2048) antes de que los expertos operen, reduciendo el tráfico de datos. El diseño Stable LatentMoE con escalares aprendibles asegura la estabilidad de la optimización. Las conexiones residuales estándar son reemplazadas por Manifold-Constrained Hyper-Connections (mHC), que permiten a las capas intercambiar información a través de 4 flujos paralelos. La cabeza de salida utiliza MultiMax, una temperatura de softmax por piezas aprendible, para suprimir la cola larga de tokens implausibles manteniendo la diversidad de opciones plausibles. Finalmente, ECHO reemplaza el módulo de predicción multi-token (MTP) reutilizando la capa final del backbone, entrenando directamente los mismos parámetros para predecir tokens futuros y mejorando las capacidades de razonamiento.

Flujo de Atención Híbrida

  1. 1 Entrada de Token Secuencia de tokens de entrada al modelo.
  2. 2 Capas HCA (37/43) Comprime 128 tokens en 1 entrada de caché. Atiende a entradas comprimidas + v...
  3. 3 Capas MLA (6/43) Mantiene acceso a nivel de token. Utiliza Sparse Attention (DSA) con indexado...
  4. 4 Low-Rank Gated Attention Multiplica salida de atención por un gate aprendido de la latente de consulta...
  5. 5 Virtual Head Attention (VHA) Paso de mezcla de bajo rango entre cabezas con conexión de salto de identidad.
  6. 6 Salida de Capa de Atención Representación de tokens procesada por la capa de atención.

Flujo de Procesamiento MoE

  1. 1 Entrada de Token (Capa MoE) Representación de token de la capa anterior.
  2. 2 Proyección a Espacio Latente Tokens proyectados de 4096 a 2048 dimensiones.
  3. 3 Multi-Router Varios routers sigmoidales puntúan expertos conjuntamente.
  4. 4 Selección de Expertos (Top-10 + 1 Com... Los tokens son dirigidos a los expertos seleccionados.
  5. 5 Ejecución de Expertos Latentes Expertos (tamaño intermedio 2048) operan en el espacio latente.
  6. 6 Up-proyección Compartida Resultado mapeado de vuelta a 4096 dimensiones.
  7. 7 Salida de Capa MoE Representación de token procesada por la capa MoE.
CapaTecnologíaJustificación
compute Transformer Architecture Marco fundamental para el procesamiento de secuencias, modificado con atención dispersa y MoE para eficiencia. 43 capas, hidden size 4096, 64 attention heads.
compute Sparse Mixture-of-Experts (MoE) Escalar la capacidad del modelo a 280B parámetros mientras se activan solo 14B por token, reduciendo el costo computacional de inferencia. vs Dense models, Smaller MoE configurations 512 expertos ruteados (top-10) + 1 experto compartido, Latent Experts (2048-d).
compute Hybrid Sparse Attention (MLA + HCA + DSA) Optimizar el uso del KV cache y la computación de atención para contextos largos, equilibrando capacidad y eficiencia. vs Dense Attention, Full Sparse Attention (CSA) HCA (128x compresión) en 37 capas, MLA (token-level) con DSA en 6 capas.
compute FlashMask-4 Kernel de atención optimizado para GPUs SM90/SM100, diseñado para manejar grandes dimensiones de cabeza de atención de manera eficiente. vs FlashAttention-4 MaskMod Soporta qk_head_dim=576, v_head_dim=512. Utiliza 2-CTA tiled parallel computation y double-buffer rotation.
data-processing BPE Tokenizer Convertir texto en secuencias de tokens para el modelo, con una tokenización especializada para números (Right-to-Left) para mejorar el rendimiento aritmético. Vocabulario de ~200K entradas, tokenización R2L de números en grupos de hasta 3 dígitos.
orchestration Distributed Muon Optimizer Optimizar el entrenamiento a escala de miles de GPUs, gestionando la comunicación y la distribución de memoria de manera eficiente para parámetros densos y MoE. vs Adam, Otros optimizadores element-wise Actualizaciones ortogonales de matrices de parámetros, sincronización de parámetros densos y MoE combinada, sharding de memoria balanceado entre nodos.

Trade-offs

Ganancias
  • ▲ Eficiencia de inferencia
  • ▲ Eficiencia de entrenamiento
  • ▲ Capacidad del modelo (parámetros totales)
  • ▲ Rendimiento en benchmarks (razonamiento, código, exámenes)
  • ▲ Soporte de contexto largo
Costes
  • △ Complejidad arquitectónica
  • △ Posible dificultad en la depuración de MoE con múltiples routers 
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_name = "baidu/ERNIE-5.5-Flash-Base-BF16"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto")

model_inputs = tokenizer(["The hybrid attention architecture of ERNIE 5.5 Flash"], add_special_tokens=True, return_tensors="pt").to(model.device)
print(tokenizer.decode(model.generate(**model_inputs, max_new_tokens=256)[0], skip_special_tokens=True))
Ejemplo básico de cómo cargar el modelo ERNIE 5.5 Flash Base y generar texto utilizando la biblioteca Transformers de Hugging Face.

Fundamentos Teóricos

La arquitectura de ERNIE 5.5 Flash se fundamenta en varios principios teóricos y avances académicos. La idea de la atención dispersa (Sparse Attention) se remonta a trabajos como 'Longformer: The Long-Document Transformer' (Beltagy et al., 2020) o 'Reformer: The Efficient Transformer' (Kitaev et al., 2020), que buscan reducir la complejidad cuadrática de la atención estándar a lineal o log-lineal. La implementación de Mixture-of-Experts (MoE) tiene sus raíces en 'Adaptive Computation Time for Recurrent Neural Networks' (Graves et al., 2016) y 'Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer' (Shazeer et al., 2017), que introdujeron la idea de activar solo un subconjunto de expertos por token para escalar la capacidad del modelo sin aumentar proporcionalmente el costo computacional. La optimización del KV cache es un área activa de investigación, con trabajos como 'FlashAttention' (Dao et al., 2022) que buscan optimizar el uso de la memoria on-chip de las GPUs para mejorar el rendimiento de la atención. La tokenización Right-to-Left para números se inspira en técnicas de procesamiento de lenguaje natural para mejorar el manejo de información numérica, un desafío conocido en los LLMs.