One-sentence Summary

Researchers from Nanjing University, Nanyang Technological University, and Imperial College London propose AURORA-LM, a continuous-latent diffusion language model that uses a Query-based Encoder-Decoder to create high-capacity, prefix-aligned latent sequences and a Block-causal Diffusion Transformer with flow matching, noise-level calibration, and self-trajectory consistency, achieving the strongest performance among evaluated continuous and diffusion-based language models on OpenWebText free generation and XSum conditional summarization while scaling to 1B parameters.

Key Contributions

• A Query-based Encoder-Decoder constructs a prefix-aligned, high-capacity continuous text latent, frozen before diffusion training to decouple representation learning; wider latents retain token-level information robustly under corruption.
• A Block-causal Diffusion Transformer models the full-width latent distribution via flow matching, generating blocks left-to-right while denoising positions within each block in parallel.
• To learn the high-dimensional distribution without reducing decoding capacity, a low-rank bottleneck is applied only to the noisy input, the noise-level distribution is calibrated to latent width, and self-trajectory consistency aligns training with iterative denoising; this yields the strongest performance among evaluated continuous and diffusion-based language models on OpenWebText free generation and XSum conditional summarization, and a 1B-parameter model surpasses a larger publicly released latent-diffusion language model.

Introduction

Continuous latent spaces serve as a powerful backbone for modern generative models, enabling high-quality synthesis in vision and audio. However, language modeling still relies on discrete tokens, creating an asymmetry in multimodal systems. Prior attempts to generate text in continuous space either repurpose token embeddings or pretrained encoder features, which are not explicitly optimized for both accurate reconstruction and tractable generation. Autoencoder-based approaches can learn a tailored latent, but they face a trade-off: compressing the representation simplifies the generative task at the cost of decoding fidelity, while preserving detail complicates latent modeling. The authors introduce AURORA-LM, a framework that decouples the construction of a text representation from the learning of its distribution. They first build a high-capacity, causally ordered latent space with a query-based encoder-decoder that faithfully supports token recovery, then freeze it and train a block-causal diffusion model via flow matching. To handle the resulting full-width latent distribution without compression, they combine a low-rank input projection, width-aware noise allocation, and a self-trajectory consistency regularizer. This design achieves strong generation and conditional summarization performance, outperforming prior continuous and discrete baselines while scaling effectively.

Method

The AURORA-LM framework is organized around two primary learning objectives: constructing a continuous text representation that can be accurately decoded back into discrete tokens, and modeling the generative distribution of this continuous latent space. The architecture comprises a Query-based Encoder-Decoder and a block-causal denoiser, which is a Transformer trained with flow matching.

The first stage focuses on continuous text latent construction. The authors leverage a Query-based Encoder-Decoder to organize a discrete text sequence into a high-capacity, causally ordered latent interface. Given a token sequence of length $L$L, the encoder produces a continuous latent sequence $z_{\mathrm{enc}} \in \mathbb{R}^{N \times D}$zenc​∈RN×D, where $D$D is the channel width and $N = \mathrm{round}(cL)$N=round(cL) is the number of latent positions determined by the latent retention ratio $c$c. To construct this ordered sequence, $N$N latent queries aggregate the variable-length token sequence. Each query is initialized from a shared learnable vector and distinguished by the context visible at its position. At layer $\ell$ℓ, query $i$i attends to preceding latent states and the corresponding token prefix, with RoPE applied to encode positions. The layer updates are computed as:

The decoder mirrors this structure, initializing $L$L output queries from a second shared vector. Output position $j$j attends to the latent prefix $z_{\mathrm{enc}, 1: \lfloor (j - 1) N / L \rfloor + 1}$zenc,1:⌊(j−1)N/L⌋+1​ and preceding decoder states to reconstruct token logits. The autoencoder is trained jointly using token-level cross-entropy loss, regularized by token-embedding and latent dropout. Once trained, the autoencoder is frozen to fix the latent space for the subsequent generative modeling stage.

The second stage involves block-causal modeling of the continuous text latents. The outputs of the frozen encoder are standardized using per-channel affine normalization to induce an empirical distribution. The block-causal denoiser models this distribution by factorizing the latent sequence into contiguous blocks of size $Q$Q. This block-causal factorization enables left-to-right causal generation across blocks while jointly denoising all positions within a block. For a clean block $\alpha^{(b)}$α(b) and Gaussian noise $\varepsilon^{(b)} \sim \mathcal{N}(0, I)$ε(b)∼N(0,I), a linear probability path is defined as:

The denoiser is trained to predict the clean endpoint $\alpha^{(b)}$α(b) from the noisy state, noise level $t$t, and clean prefix $\alpha^{(<b)}$α(<b).

To efficiently learn the full-width latent distribution, the authors introduce a noisy-latent input bottleneck. While the frozen decoder requires a full-width clean block of dimension $D$D, the denoiser processes the noisy input through a lower-rank pathway with bottleneck dimension $D_b < \min(D, H)$Db​<min(D,H), where $H$H is the Transformer hidden width. This forces the model to extract a compact representation before processing. Parallel blockwise training is achieved by using clean prefixes from the frozen encoder and applying a two-stream attention mask, allowing all block conditionals to be evaluated in a single forward pass. The clean-endpoint flow-matching loss is defined as:

To bridge the gap between independent training samples and the sequential nature of inference, the framework incorporates self-conditioning and self-trajectory consistency. Self-conditioning feeds the clean-latent prediction from the previous sampling step back into the model as an additional input during training. Self-trajectory consistency aligns the clean-latent predictions at neighboring states along the sampling trajectory by minimizing the distance between the current model's prediction and an exponential moving average model's prediction at a lower-noise level $t'$t′:

The complete training objective combines these losses: $\mathcal{L}_{\mathrm{train}} = \mathcal{L}_{\mathrm{FM}} + \lambda_{\mathrm{ct}} \mathcal{L}_{\mathrm{ct}}$Ltrain​=LFM​+λct​Lct​.

During inference, latent generation proceeds block by block from left to right. For prompt-conditioned generation, the prompt is encoded and standardized to initialize the prefix. The remaining blocks are generated by denoising Gaussian noise from $t=1$t=1 to $t=0$t=0, conditioned on the completed prefix. The generated standardized latent sequence is then un-standardized and passed through the frozen query decoder to recover the output token sequence. The authors also employ guidance mechanisms at inference time, such as Self-Conditioning Classifier-Free Guidance for unconditional generation and standard Classifier-Free Guidance for prompt-conditioned generation, to control the trade-off between generation quality and diversity.

Experiment

The experiments first conduct controlled ablations on OpenWebText to validate key design choices, showing that wide latent representations (D=1024), a narrow noisy-input bottleneck (128), a high-noise training schedule, x₀ prediction with x₀-space loss, and blockwise generation with self-trajectory consistency collectively optimize performance. System-level comparisons then demonstrate that AURORA-LM-S outperforms autoregressive, discrete-diffusion, and other continuous-generation baselines on both unconditional text generation and conditional summarization, confirming the advantage of modeling a learned continuous representation. Finally, scaling AURORA-LM-L to 1.01B parameters yields consistent improvements over a larger state-of-the-art latent diffusion model across nine diverse language tasks, establishing the design's effectiveness at larger scales.

Direct regression of the clean latent (target x0, loss x0) achieves the highest MAUVE scores across both noisy-input widths, substantially outperforming all velocity-based alternatives. The combination of a clean target with a velocity loss collapses to near-zero quality, while velocity targets paired with either loss show moderate but lower performance. These trends remain consistent whether the noisy input is bottlenecked at 128 dimensions or kept at the full 1024-dimensional width. Clean-latent regression (x0 target, x0 loss) yields MAUVE above 0.80 for both bottleneck widths, far exceeding the next-best configuration. Using a velocity loss with a clean-latent target drives MAUVE below 0.06, indicating a fundamental mismatch between the prediction target and loss space. Velocity-prediction targets (v target) paired with an x0 loss attain moderate MAUVE (0.73 at width 128) but drop sharply to 0.33 at full width. The relative ranking of target–loss combinations is preserved across noisy-input widths, but the performance gap between clean and velocity targets widens at the full 1024-dimensional width.

On unconditional generation of 1,024-token OpenWebText samples, the continuous-latent AURORA-LM-S achieves the best overall quality, posting the lowest Gen-PPL (23.56) and highest MAUVE (0.890) among all compared systems. It surpasses the previous continuous model ELF-B, which suffered from extremely low MAUVE despite competitive perplexity, and substantially outperforms autoregressive and discrete-diffusion baselines. AURORA-LM-S reduces Gen-PPL by about 40% relative to the autoregressive baseline and by over 70% relative to discrete-diffusion models, while raising MAUVE to 0.890. ELF-B achieved a similar Gen-PPL (24.11) but its MAUVE collapsed to 0.229, showing that AURORA-LM-S uniquely combines low perplexity with high distributional similarity to real text. Discrete-diffusion approaches (SEDD, MDLM) yield Gen-PPL above 119 and MAUVE below 0.7, indicating much weaker long-text generation quality than the continuous-latent model.

On the XSum conditional generation task, AURORA-LM-S obtains the top ROUGE-1, ROUGE-2, and ROUGE-L scores, overtaking all evaluated baselines including the prior best model ELF-B. The strong performance extends the model's continuous-latent formulation from unconditional to prompt-conditioned generation. AURORA-LM-S leads all compared models on XSum across all three ROUGE metrics. It improves over the next best model ELF-B by more than one point on both ROUGE-2 and ROUGE-L.

The smaller 1B-parameter AURORA-LM-L outperforms the larger 1.8B-parameter Cola-DLM on all nine evaluated benchmarks and achieves a higher macro average, demonstrating that the model's advantages persist at scale and across diverse language tasks. AURORA-LM-L attains a macro average of 32.6 versus Cola-DLM's 25.1, outperforming on every task despite having roughly half the parameters. The largest absolute gaps appear on story and commonsense reasoning tasks such as HellaSwag (18.4 vs. 5.7) and StoryCloze (54.8 vs. 33.8).

Ablation studies on latent regression targets and losses confirm that directly predicting the clean latent with a clean-target loss yields the highest generation quality, while velocity-based formulations sharply degrade performance, especially when mismatched. In unconditional long-text generation, the continuous-latent AURORA-LM model attains the best combination of low perplexity and high distributional similarity, far surpassing autoregressive and discrete-diffusion baselines that either suffer from poor coherence or collapsed similarity. AURORA-LM extends this strong generative capability to prompt-conditioned summarization, where it achieves top ROUGE scores, and to few-shot reasoning benchmarks, where a smaller 1B-parameter version outperforms a 1.8B discrete-diffusion alternative on all tasks, confirming efficient scaling across diverse language evaluations.