Cola2 beats Llama-3-8B on HellaSwag with 43 times fewer tokens

A team of researchers has presented Cola2, a hierarchical language modeling framework in which meaning lives in a continuous latent space and a discrete decoder assembles the text. Trained on 346.03B raw tokens, a not-yet-converged checkpoint averaged 60.10 across eight tasks and outperformed Llama-3-8B on HellaSwag and StoryCloze. Llama-3's stated pretraining budget is more than 43 times larger.

What exactly the latent space compresses

The idea of hierarchical generation has been around for a while: the model first plans the meaning in broad strokes, then realizes it in tokens. The weak point is always the same. If the latents compress the wrong thing, or compress it the wrong way, the decoder either ignores them or degenerates into a simple copy machine.

Cola2 tackles this in several steps. Text is split dynamically into semantic spans rather than fixed-length chunks. Negative feedback controls the compression rate without auxiliary losses: the system adjusts its own threshold, so no extra penalty terms are added to training. To make the latents both compact and predictive, the authors mix two signals: reconstruction of the source text and ordinary generation.

A diffusion decoder and the balance between copying and independence

The decoder is a block-causal masked diffusion model. It is trained under three kinds of conditions: with clean latents, with noised latents, and with latents dropped entirely. This makes it possible to apply classifier-free guidance. By changing the guidance strength, you can shift the model's behavior between exact recovery of meaning from the latent and autonomous generation.

One engineering detail concerns gradient flow. Token-prediction gradients pass through the latents into the encoder, so the representations adapt to the decoder's needs. The Flow Matching prior, however, is decoupled from representation learning via stop-gradient. The prior learns to predict the next latents, but it does not pull the encoder along with it and does not corrupt the space the encoder builds.

There is also a subtlety at inference. The generated text is re-encoded using fixed span boundaries, so the latent history matches the actual text context rather than what the model "thought" earlier. Without this step, errors between the continuous and discrete parts would accumulate from block to block.

The numbers and their caveats

• Training budget in the main experiments: about 210B tokens, five downstream tasks.
• Scaling: the authors describe it as competitive and unsaturated.
• Large run: 346.03B raw tokens, checkpoint not yet converged.
• Average across eight tasks: 60.10.
• Beats Llama-3-8B on HellaSwag and StoryCloze with a pretraining budget more than 43 times smaller.

The comparison with Llama-3-8B should be read with care. HellaSwag and StoryCloze measure the plausibility of a continuation and commonsense reasoning, which is exactly what a model with an explicit layer of semantic structure handles well. The picture may look different on knowledge and reasoning benchmarks, and according to the paper, a win is claimed on only two of the eight tasks. On the other hand, the scaling curve does not plateau, which means the potential at larger budgets has not been exhausted.

Where this could be useful

For those building models, the main takeaway is practical: data efficiency can be sought not only in the transformer architecture or corpus curation, but also in how the intermediate semantic layer is designed. If latents compress text into spans and the decoder generates in blocks, inference cost could also drop: fewer steps per unit of meaning.

For researchers working on diffusion language models, Cola2 offers a ready-made recipe for joining the continuous and discrete parts: negative feedback for compression, three conditioning modes for the decoder, stop-gradient for the prior, and recomputation of latents during generation. Each of these elements can be tested and transferred separately.

For the industry, this is a signal that alternatives to the purely autoregressive approach are gaining weight. For now, we are talking about a model in the 8B class with unfinished training. But if the scaling holds up at tens and hundreds of billions of parameters, hierarchical latent models could become a serious candidate for the next step after classic next-token prediction.