Cola2 supera a Llama-3-8B en HellaSwag con 43 veces menos tokens

Un equipo de investigadores ha presentado Cola2, un framework de modelado jerárquico del lenguaje en el que el significado vive en un espacio latente continuo y el texto lo ensambla un decodificador discreto. Con 346,03 mil millones de tokens de origen, un checkpoint que aún no había convergido obtuvo una media de 60,10 en ocho tareas y superó a Llama-3-8B en HellaSwag y StoryCloze. El presupuesto de preentrenamiento declarado de Llama-3 es, además, más de 43 veces mayor.

Qué comprime exactamente el espacio latente

La idea de la generación jerárquica lleva tiempo en el aire: primero el modelo planifica el significado a grandes rasgos y después lo materializa en tokens. El punto débil siempre es el mismo. Si los latentes comprimen lo que no deben o lo hacen mal, el decodificador o bien los ignora o bien se convierte en una simple fotocopiadora.

Cola2 aborda el problema con varias medidas. El texto se divide dinámicamente en spans semánticos, no en fragmentos de longitud fija. La tasa de compresión se regula mediante retroalimentación negativa, sin pérdidas auxiliares: el propio sistema ajusta el umbral sin añadir términos de penalización extra al entrenamiento. Y para que los latentes sean compactos y, a la vez, predictivos, los autores combinan dos señales: la reconstrucción del texto original y la generación convencional.

Decodificador de difusión y equilibrio entre copia y autonomía

El decodificador es bloque-causal, de difusión enmascarada. Se entrena con tres tipos de condiciones: con latentes limpios, con latentes ruidosos y con latentes descartados por completo. Esto permite aplicar classifier-free guidance. Al variar la intensidad del guidance, se puede desplazar el comportamiento del modelo entre la reconstrucción fiel del significado a partir del latente y la generación autónoma.

Un detalle de ingeniería concierne al flujo de gradientes. Los gradientes de la predicción de tokens pasan por los latentes hasta el codificador, de modo que las representaciones se ajustan a las necesidades del decodificador. En cambio, el entrenamiento del prior con Flow Matching está desacoplado del entrenamiento de las representaciones mediante stop-gradient. El prior aprende a predecir los siguientes latentes, pero no arrastra al codificador ni destruye el espacio que este construye.

Hay también un matiz en la inferencia. El texto generado se vuelve a codificar con una segmentación de spans fija, para que el historial de latentes coincida con el contexto textual real y no con lo que el modelo «pensó» antes. Sin este paso, los errores entre la parte continua y la discreta se acumularían de un bloque a otro.

Los números y qué hay que matizar en ellos

• Presupuesto de entrenamiento en los experimentos principales: unos 210 mil millones de tokens, cinco tareas downstream.
• Escalado: los autores lo describen como competitivo y sin saturar.
• Entrenamiento a gran escala: 346,03 mil millones de tokens de origen, con el checkpoint aún sin converger.
• Media en ocho tareas: 60,10.
• Supera a Llama-3-8B en HellaSwag y StoryCloze con un presupuesto de preentrenamiento más de 43 veces menor.

La comparación con Llama-3-8B conviene leerla con cautela. HellaSwag y StoryCloze miden la plausibilidad de la continuación y el sentido común, es decir, justo lo que capta bien un modelo con una capa explícita de estructura semántica. En benchmarks de conocimiento y razonamiento el panorama puede ser distinto, y según el texto del trabajo la victoria se declara solo en dos de las ocho tareas. En cambio, la curva de escalado no llega a una meseta, lo que implica que el potencial con presupuestos mayores no está agotado.

Dónde puede resultar útil

Para quienes construyen modelos, la conclusión práctica principal es que la eficiencia de datos se puede buscar en la arquitectura del transformer o en la curación del corpus, pero también en cómo se diseña la capa semántica intermedia. Si los latentes comprimen el texto en spans y el decodificador genera por bloques, el coste de inferencia podría bajar: menos pasos por unidad de significado.

Para los investigadores de modelos de lenguaje de difusión, Cola2 ofrece una receta lista para unir la parte continua y la discreta: retroalimentación negativa para la compresión, tres regímenes de condiciones para el decodificador, stop-gradient para el prior y recálculo de latentes durante la generación. Cada uno de estos elementos se puede probar y trasladar por separado.

Para la industria, es una señal de que las alternativas al esquema puramente autorregresivo ganan peso. Por ahora se trata de un modelo de clase 8B con un entrenamiento inacabado, pero si el escalado se confirma con decenas y cientos de miles de millones de parámetros, los modelos latentes jerárquicos serán un candidato serio como siguiente paso tras el next-token prediction clásico.