Cola2: Co-Adaptive Hierarchical Language Modeling with Continuous Semantic Compression and Balanced Discrete Generation
Hongcan Guo, Yian Zhao, Feng Wang, Shen Nie, Guoqiang Wei, Yan Zeng, Hengshuang Zhao
cs.CL
2026-10-22
Cola2 compresses text into continuous latents and decodes with masked diffusion. Unconverged at 346B tokens, it averages 60.10 on eight tasks and beats Llama-3-8B on two.
Hierarchical language models separate semantic planning from the tokens that realize it. A continuous latent is meant to hold what a span means, and a discrete decoder writes the words. Cola DLM already samples a latent prior and then decodes, but splitting the two spaces does not give them a useful division of labor. A shorter latent sequence is not automatically a predictive code. Reconstruction alone does not force the latent to fill in hidden or future tokens. Freeze the encoder before the decoder has learned, and the split of information is fixed too early. Let the prior-fitting gradient update the encoder, and the latent drifts toward states the prior can match easily.
Cola2 settles three choices in one training run: what the latent keeps, how much residual text the decoder must supply, and which losses may change the representation.
Content-based segmentation cuts the target into contiguous spans. Each span is one diagonal-Gaussian latent. The compression rate ρ = K/T is the number of latents divided by the number of tokens. The default is about two tokens per latent, width 64, maximum span 8. A gradient-free negative-feedback bias holds ρ near the target. There is no auxiliary rate loss competing with representation learning. Attention is bidirectional inside a span and causal across latents.
The decoder is a block-causal masked diffusion model. Past blocks are visible, the current block is denoised in both directions, and it cannot read its own clean tokens. A duration head predicts how many tokens a latent covers, then the latent is expanded into a token-aligned condition. A block-causal Diffusion Transformer fits the latent sequence with Flow Matching. Stop-gradient sits on that path, so the prior loss cannot rewrite the encoder. Semantics are sampled on the short latent axis. Wording stays in token space.
Each latent block is clean, noised, or dropped with probabilities 0.2, 0.4, and 0.4. A generative loss hides part of the encoder input, so the latent has to predict tokens it cannot copy. Token-loss gradients pass through the latents into the encoder. Stage 1 learns compression and this balance. Stage 2 adds the DiT. The prior gradient stays blocked, and the decoder-to-encoder path stays open.
Inference proposes a latent block, decodes tokens, re-encodes them on the span layout already chosen, and only then commits the latent. Defaults are 8 DiT steps and 4 unmasking steps, with classifier-free guidance on both sides. The 346B-token run initializes the decoder and the encoder embedding from Qwen3-4B-Base, keeps causal self-attention, and feeds [MASK] at hidden positions.
Ablations train on OLMo 2 and report LAMBADA accuracy. At ρ = 1/2 and maximum span 4, dynamic segmentation scores 25.0%, 4.4 points above random spans and 8.2 above fixed-length chunks. Maximum span 8 reaches 29.4%, above the uncompressed model at 28.7%. The auxiliary-free controller scores 25.0%, and DLCM scores 18.4%. A four-layer 0+2+2 encoder, with no token layer, two pooling layers, and two latent layers, scores 29.6%. Six-layer variants land between 20.9% and 26.3%. Removing pooling falls to 15.7%. Intra-span bidirectional attention scores 26.7%. Fully causal attention, endpoint pooling, and mean pooling all stay at or below 19.0%. Width 64 scores 32.3%. Width 512 falls to 14.5%. Longer spans still cost accuracy: the best setting drops from 32.3% at span 2 to 25.2% at span 4.
Pure reconstruction scores 15.7%. Span-AR and a BERT-style objective reach 25.2% and 26.5%. With an MDM decoder, a frozen encoder scores 47.6%, joint updates score 53.5%, and joint updates plus the generative loss score 58.2%. Sending the DiT gradient into the encoder leads by 2.2 points at 20k steps, 23.6% versus 21.4%, and trails by 1.0 point at 140k, 32.1% versus 33.1% with the gradient detached. On 8 denoising steps, guidance scale 14 scores 37.1%, against 12.2% at scale 1. From 8 steps to 128 steps the spread is at most 0.3 points.
Autoregressive multi-token prediction and MDM never differ by more than 1.4 points. At 150k steps they score 44.1% and 45.5%. Drop, noise, and clean probabilities of 0.4, 0.4, and 0.2 score 44.7% on LAMBADA, while reconstruction stays above 99.5%. A decoder trained on that mix peaks at wdec = -0.9, again 44.7%. A clean-heavy decoder peaks at 40.3%. With 3B parameters in total, a 2B DiT plus a 1B decoder scores 44.4%, 3.3 to 4.4 points above the other splits. Shuffling latents drops accuracy from 46.5% to 34.3%. Zeroing them drops it only to 41.2%.
A matched scale-up to about 210B tokens covers RACE, OBQA, Story Cloze, LAMBADA, and SQuAD. Uncompressed Cola-MDM averages about 34.5%, autoregression about 33%, and about 64% on LAMBADA, roughly 15 points above Cola DLM. Compressed Cola2-MDM stays near autoregression on LAMBADA, about 61%, and leads on RACE, about 10.5% versus 4.5%. It is weaker on SQuAD and Story Cloze. After about 150B tokens its average gap to LLaDA shrinks from about 4 points to about 1.5, and the curve has not flattened.
A separate unconverged checkpoint uses 346.03B source tokens. The data are OLMo 2 retokenized with the Qwen3-4B tokenizer at context 4096. The three stages use 52.43B, 125.83B, and 167.77B source tokens. The last stage freezes the VAE and trains only an approximately 4B DiT. The eight tasks are exact match on generated answers, and each Cola2 score is the mean of three samples. Recipes, tokenizers, token accounting, and compute are not matched.
| Model | Train tokens | 8-task avg | HellaSwag | StoryCloze | MMLU |
| Llama-3-8B | 15T+ | 66.25 | 39.21 | 79.49 | 63.19 |
| Qwen2.5-7B | 18T | 75.08 | 75.67 | 95.09 | 70.16 |
| Qwen3-8B-Base | 36T | 75.87 | 68.13 | 97.44 | 72.31 |
| LLaDA-8B-Base | 2.3T | 72.11 | 69.65 | 95.94 | 61.27 |
| Cola2 | 346.03B | 60.10 | 41.75 | 80.42 | 51.32 |
The 60.10 average sits 6.15 behind Llama-3-8B. Cola2 leads only on HellaSwag, 41.75 versus 39.21, and StoryCloze, 80.42 versus 79.49. On the same sheet Qwen3-8B-Base averages 75.87 and LLaDA-8B-Base averages 72.11.
Changing the token decoder barely moves the number. MDM and autoregressive multi-token prediction stay within 1.4 points, so the latent interface is not tied to one generator. The latent is what moves the score. Widening it from 64 to 512 drops accuracy from 32.3% to 14.5%, and shuffling it costs 12.2 points. Prior gradients that rewrite the encoder look helpful early and hurt later. At a fixed 3B, 2B for the DiT and 1B for the decoder is the best split. Pushing parameters further to either side is worse.
346.03B source tokens is about 1/43 of the 15T+ reported for Llama-3-8B, and the eight-task average is still 6.15 points lower. The two leads are 2.54 and 0.93 points, from a decoder initialized with Qwen3-4B-Base, under generative exact match. Next to Qwen3-8B-Base and LLaDA-8B-Base, this is an early run that has not finished training.
The latent DiT is far from converged. Larger models and longer training are still missing, so the ceiling is unknown. Compression cuts FLOPs on the latent prior in theory. End-to-end speed also depends on the decoder, the CFG branches, and re-encoding, and no speed measurement is reported. How an autoregressive decoder and an MDM should split the work, and whether the same latent can meet vision, is left open.
The eight-task table is not a matched-compute comparison. The decoder starts from Qwen3-4B-Base, and the metric is exact match on generated text. That is why Llama-3-8B scores 39.21 on HellaSwag, a figure that does not line up with the usual multiple-choice leaderboard. The five-task curves at about 210B tokens are still in the mid-30s and do not continue into the 60.10 run, which changes context length, initialization, and the three-stage schedule. How tightly the rate controller actually holds ρ is not reported.