STAPLE añade miles de millones de parámetros a OLMo-2 y decodifica el doble de rápido que MoE

Un equipo de investigadores ha publicado en arXiv STAPLE, un método que permite aumentar el número de parámetros de un modelo de lenguaje ya entrenado sin recurrir a un costoso preentrenamiento disperso desde cero. En OLMo-2 1B aporta +4,27 puntos en HumanEval, +2,80 en MBPP, +2,27 en GSM8K y +1,65 en DROP respecto a la simple continuación del entrenamiento dense. Según los autores, la generación se ralentiza no más de un 1,5% y el entrenamiento, no más de un 3%, a pesar de que se añaden al modelo miles de millones de parámetros en una fase tardía.

Por qué no MoE

La computación dispersa se considera desde hace tiempo una vía para aumentar la capacidad del modelo sin inflar el cómputo activo por token. El problema es el coste: el preentrenamiento de arquitecturas dispersas exige muchos datos y mucho dinero. Los métodos de upcycling existentes, es decir, la conversión de un checkpoint dense ya entrenado en un modelo disperso, apuntan sobre todo a Mixture-of-Experts. Los autores identifican ahí tres cuellos de botella: la homogeneidad de los expertos, la baja sample efficiency y la sobrecarga de sistema por el enrutamiento y la memoria.

La memoria estática dispersa funciona de otra manera. Su acceso es predecible: los parámetros se eligen por el identificador del token, no por la decisión de un router. Las actualizaciones también están ligadas a cada token concreto. Pero, como señalan los autores, insertar de forma ingenua esta memoria en un modelo preentrenado ofrece un rendimiento débil.

Cómo funciona STAPLE

El método se basa en STEM, donde las proyecciones up de los bloques FFN se sustituyen por embeddings indexados por token. Para entrenar desde cero funciona, pero en un checkpoint ya entrenado la sustitución rompe el cómputo acumulado. STAPLE hace otra cosa: las proyecciones up se quedan en su sitio y los embeddings se añaden como camino residual. La parte preentrenada sigue funcionando como antes, mientras que la nueva memoria aprende aquello que el modelo conoce peor sobre los tokens: información adicional específica de cada uno.

Hay un detalle de ingeniería importante: la implementación descarga los embeddings al host. En la GPU se mantiene solo alrededor del 5% de los embeddings más frecuentes y el resto vive en la memoria del host. Por la distribución de los tokens, eso basta para que el rendimiento apenas caiga.

Las cifras

Principales resultados del preprint:

• OLMo-2 1B, mejora frente al continued dense training: HumanEval +4,27, MBPP +2,80, GSM8K +2,27, DROP +1,65.
• El resultado medio crece junto con la capacidad estática y las mejoras se mantienen en OLMo-2 7B.
• STAPLE supera a las líneas base estáticas dispersas: STEM y un modelo al estilo Engram.
• Frente al upcycling a MoE: mayor puntuación media, decodificación 2 veces más rápida, menos de la mitad de memoria de GPU para el modelo y menos energía por token.
• Sobrecarga: generación dentro del 1,5% respecto a la base dense y entrenamiento dentro del 3%.

La salvedad es importante: la comparación con MoE en velocidad y energía se hizo en una evaluación orientada al rendimiento (throughput) de decodificación. En escenarios con requisitos estrictos de latencia por petición individual, el resultado puede ser distinto. Además, por ahora se trata de un preprint con modelos de 1B y 7B, y no se ha demostrado el escalado a decenas y cientos de miles de millones de parámetros.

Qué aporta en la práctica

Para los equipos que ajustan modelos abiertos, STAPLE ofrece una vía barata para subir en capacidad: se toma un checkpoint existente, se le añade memoria estática y el modelo sigue entrenándose con una pequeña pérdida de velocidad. La mejora es especialmente visible en código y matemáticas, donde el conocimiento exacto de tokens y patrones resulta particularmente útil. En inferencia, lo interesante es el perfil de memoria: la mayor parte de los parámetros reside en el host y en la GPU solo hace falta una caché pequeña. Eso reduce los requisitos de memoria de vídeo frente a un MoE de calidad comparable.

Antes de llevar el enfoque a producción conviene comprobar tres cosas: el comportamiento en modelos mayores de 7B, la robustez ante un cambio de dominio en el que la distribución de los tokens frecuentes sea distinta, y la latencia con batches pequeños. Si estas pruebas salen bien, el upcycling mediante embeddings estáticos se convertirá en una alternativa seria a MoE para ampliar modelos ya entrenados.