AURORA-LM: Autoencoding Unified Representation for Continuous-Latent Diffusion Language Modeling

Jiajun Liang, Yucheng Liao, Yukang Cao, Jiazhe Wei, Ken Li, Wende Tan, Jiankun Zhang, ZY Cui, Jingkang Yang, Liucheng Guo, Shiqi Yang, B. Yang, Caifeng Shan, Ziwei Liu, Chenyang Si

AURORA-LM decouples text representation from generative modeling to enable high-fidelity continuous language generation.

How can we model text generation in a continuous latent space using diffusion, while maintaining the causal structure required for language?

Continuous language models often compress text into low-capacity latents to simplify diffusion, which degrades the decoder's ability to recover precise tokens. AURORA-LM separates these tasks: a frozen query-based autoencoder constructs a high-capacity, prefix-aligned latent sequence, while a block-causal denoiser learns to generate these full-width latents using flow matching. This approach achieves state-of-the-art performance among continuous language models on OpenWebText and XSum, scaling effectively to 1 billion parameters.

Paper Primer

The core challenge is the tension between latent tractability and reconstruction fidelity. AURORA-LM resolves this by freezing a high-capacity autoencoder that maps text to a causal latent sequence, then training a denoiser to model that distribution directly without forcing the autoencoder to compress information for the sake of the diffusion prior.

The method hinges on a block-causal denoiser: it generates latent blocks from left to right while jointly denoising positions within each block. To handle the high-dimensional target, the model uses a low-rank projection for the noisy input pathway while maintaining full-width clean-latent prediction, calibrated by self-trajectory consistency to stabilize the denoising path.

AURORA-LM outperforms existing continuous and diffusion-based language models on free generation and conditional summarization.

Comparison against autoregressive, discrete-diffusion, and continuous embedding-based baselines on OpenWebText and XSum. Surpasses a larger publicly released latent-diffusion language model under matched evaluation protocols at 1B parameter scale.

Why does the paper use a block-causal factorization instead of standard joint denoising?

Joint denoising would discard the causal structure of language, while token-by-token generation would sacrifice the parallel processing benefits of continuous diffusion. Block-causal factorization provides an intermediate organization that supports left-to-right conditioning across blocks and parallel denoising within them.

How does this approach differ from previous latent-space language models?

Previous methods often compress the latent space to make the distribution easier to learn, which limits token-level fidelity. AURORA-LM explicitly retains the full channel capacity of the latent representation for the decoder and adapts the diffusion model to handle the resulting complex distribution.

Continuous language generation can achieve high fidelity by treating the latent space as a high-capacity interface rather than a bottleneck, provided the generative model is specifically designed to accommodate the full-width latent target.

Introduction

Why language still relies on discrete tokens while other data types move to continuous latents.

Modern generative models treat images, video, and audio as points in continuous latent spaces, yet language generation still hinges on discrete token vocabularies.

Current continuous‑latent language approaches either inherit token embeddings that were never meant for joint encoding‑decoding, or they compress autoencoded latents so diffusion can operate, but this compression erodes token‑level fidelity.

AURORA‑LM builds a decodable, high‑capacity continuous latent for text and then trains a block‑causal diffusion model to generate that latent directly.

Retaining a high‑capacity latent improves reconstruction but expands the latent distribution, making diffusion harder; AURORA‑LM mitigates this by adding noise only to the input side, calibrating noise to latent width, and enforcing self‑trajectory consistency to close the training‑inference gap.

**Figure 1.** Evaluations of AURORA-LM across various generation settings and model scales. (a) AURORA-LM achieves the best displayed results on OpenWebText and XSum; radial values are normalized per metric, with Gen-PPL shown reciprocally. (b) AURORA-LM outperforms a larger publicly released latent-diffusion language model across nine language benchmarks.

The key shift is moving from discrete token sequences to a continuous latent that can be diffused while still supporting exact text decoding.

Method Overview

We detail AURORA‑LM’s block‑causal diffusion mechanism that learns full‑width latent distributions.

Existing language models either tie representation learning to the generative prior or sacrifice capacity to make diffusion tractable. AURORA‑LM separates these concerns, first building a rich continuous latent and then learning its distribution with a block‑causal diffusion model.

Instead of denoising a single long sequence at once, the model processes fixed‑size blocks left‑to‑right while jointly cleaning all positions inside each block, and it learns a velocity field that continuously transports noisy latents to the clean target.

Block 1 (positions 1‑2) receives noisy inputs z₁, z₂ and the model predicts clean outputs x̂₁, x̂₂.

Self‑trajectory consistency enforces that the prediction for position 2 at the current step aligns with the prediction for position 2 one step later after a small additional noise injection.

Block 2 (positions 3‑4) is processed only after Block 1’s clean predictions are available, preserving left‑to‑right causality.

The low‑rank projection reduces the noisy‑input dimension from 4 to 2 before the velocity field is applied, cutting compute while the clean‑latent prediction remains 4‑dimensional.

After the ODE integration finishes, the concatenated clean predictions [x̂₁, x̂₂, x̂₃, x̂₄] form the full‑width latent ready for decoding.

Self‑trajectory consistency ties adjacent diffusion states together, preventing drift and allowing the model to reach high‑quality latents in fewer ODE steps.

**Figure 2.** Language modeling paradigms organized by the representation generated at inference. Autoregressive models generate discrete tokens one at a time; masked discrete diffusion iteratively denoises masked tokens; embedding-space diffusion denoises continuous token embeddings; and latent-space diffusion generates encoder-decoder text latents before decoding them into tokens. AURORA-LM adopts the latent-space formulation.

The AURORA-LM Pipeline

How AURORA‑LM builds a continuous, prefix‑ordered latent space for diffusion.

Discrete token spaces make diffusion inefficient and break the smoothness needed for continuous generative modeling. AURORA‑LM therefore introduces a continuous latent interface that preserves token information while enabling blockwise generation.

**Figure 3. Overview pipeline of AURORA-LM.** (a) The Query-based Encoder-Decoder constructs an ordered text latent sequence from causally expanding input prefixes and maps it back to token logits via a query decoder. (b) A block-causal denoiser parameterizes the latent prior and is trained with flow matching. (c) At inference, latent blocks are generated from left to right with KV caching of the completed prefix.

It writes each token prefix into its own latent slot using shared queries, so the latent sequence grows in lockstep with the text and can be generated block by block.

How does this differ from a conventional encoder that outputs a single pooled vector?

Standard encoders collapse the entire sequence into one representation, losing the prefix ordering needed for blockwise generation. The Query‑based Autoencoder keeps a distinct latent for each prefix, preserving causality and enabling the diffusion model to generate blocks sequentially.

Encoder query 1 attends to token prefix w₁ (first token) and receives no prior latents, producing z₁ = [0.3, ‑0.1].

Encoder query 2 now sees tokens w₁‑w₃ (⌈2·4/2⌉ = 4, but we cap at L = 4) and the previous latent z₁; after MHA+FFN it yields z₂ = [0.7, 0.2].

Decoder query dec is initialized as `q_dec` = [0, 0]. For output token 1 it can attend only to z₁, producing hidden h₁ = [0.25, ‑0.05] and logits o₁ = h₁ Eᵀ.

For output token 2 the decoder sees latents z₁‑z₂, yielding h₂ = [0.60, 0.10] and logits o₂ = h₂ Eᵀ.

Tokens 3 and 4 follow the same pattern, each using the full latent prefix available at their position.

Even with only two latent slots the decoder can reconstruct all four tokens because each latent aggregates progressively larger prefixes, demonstrating how compression (c < 1) preserves enough information for exact token recovery.

Block-Causal Modeling

Learn a latent distribution by generating blocks left‑to‑right while denoising each block in parallel.

Generating an entire latent sequence at once discards the causal ordering learned by the autoencoder, while generating one position at a time wastes the parallelism that continuous denoising affords. Block‑causal modeling strikes a middle ground: it proceeds left‑to‑right across blocks but denoises all positions inside a block together — like writing a page at a time while drafting the sentences on that page in parallel.

We generate latent representations block by block, preserving the left‑to‑right causal prefix while exploiting parallel denoising inside each block.

Normalize the encoder output $z_{\text{enc}}$ using the per‑channel statistics $\mu$ and $s$ to obtain $z$.

Partition $z$ into the three blocks defined above.

Generate $\alpha^{(1)}$ by sampling a Gaussian noise $\epsilon^{(1)}$, forming the path $\alpha^{(1)}_t$, and applying the learned flow‑matching inverse at $t=0$.

Condition on $\alpha^{(1)}$ to generate $\alpha^{(2)}$ using its own noise $\epsilon^{(2)}$ and the same flow‑matching procedure.

Finally generate the single‑position block $\alpha^{(3)}$ conditioned on the concatenated prefix $\alpha^{(<3)}=(\alpha^{(1)},\alpha^{(2)})$.

Smaller $Q$ yields more sequential steps (more conditioning) but less parallel work per step; larger $Q$ does the opposite, so $Q$ is a tunable trade‑off between latency and contextual fidelity.

How does block‑causal modeling differ from standard autoregressive diffusion over tokens?

Autoregressive diffusion treats each token (or latent position) as a separate time step, discarding the parallel denoising capacity of continuous diffusion. Block‑causal modeling keeps the left‑to‑right causal prefix (like autoregression) but groups $Q$ positions together so the diffusion network can process the whole group simultaneously, preserving most of the parallel speedup while still respecting order.

Learning the Latent Distribution

Full-width latents are learned via a low‑rank bottleneck, parallel block training, and trajectory‑consistent conditioning.

Training a diffusion prior for full‑width latents is costly because the decoder expects a $D$‑dimensional clean block while the denoiser need not process every noisy channel at that width. The section resolves this tension with three tricks: a low‑rank bottleneck on the noisy input, parallel blockwise training, and a consistency loss that ties successive sampling steps together.

The model predicts a complete $D$‑dimensional clean block for each latent segment, even though the decoder downstream never sees the intermediate noisy representation.

Only the noisy pathway is squeezed through a low‑rank projection, forcing the diffusion model to compress the corruption before the Transformer.

Instead of rolling out the sequential generation process during training, the model sees all blocks at once while still respecting the causal mask that will be used at inference.

The model is given its own previous clean‑latent estimate as an extra conditioning signal, letting it refine predictions across denoising steps.

Beyond feeding the previous estimate, the model is explicitly penalized if its predictions at two consecutive diffusion steps disagree.

Down‑project: $W_{\text{down}}$ maps $\alpha(b)$ to a $2$‑dim vector $[0.5,\, -0.1]$.

Up‑project: $W_{\text{up}}$ lifts the $2$‑dim vector to hidden width $8$, yielding $h = [0.5,\, -0.1,\, 0,\, 0,\, 0,\, 0,\, 0,\, 0]$.

Transformer processes $h$ and outputs a full‑width clean estimate $\hat{z}_\theta = [0.48,\, -0.12,\, 0.03,\, 0.95]$.

Self‑conditioning (with $p_{\text{sc}}=1$) feeds $\hat{z}_\theta$ back as an extra token for the next denoising step.

Self‑trajectory consistency: with $t=0.6$, $S=16$, the previous step $t' = 0.55$. The Euler update produces $\alpha_{t'} = 0.92\,\alpha_t + 0.08\,\hat{z}_\theta$.

The bottleneck forces the model to compress the noisy signal before expansion, while the consistency step guarantees that the clean estimate used for the next update stays aligned with the EMA target.

**Figure 4. Overview of self-trajectory consistency.** Self-trajectory consistency aligns neighboring clean-latent predictions along a model-induced update.

Generation and Decoding

Generate latent blocks autoregressively, then decode them back to text.

Generating an entire latent sequence in one shot is memory‑intensive and prevents incremental conditioning; the trick is to generate the latents block‑by‑block, each step anchored to the already‑produced prefix.

Think of writing a paragraph one phrase at a time: each new phrase is drafted while looking at the words already written, and the draft is refined before being locked in.

Initialize block 1: sample $\epsilon_1 = (0.3,\,-0.7)$.

Apply one Euler step: $\hat{\alpha}(1) = \epsilon_1 - 0.5\,\epsilon_1 = (0.15,\,-0.35)$ (self‑conditioning not yet available).

Append $\hat{\alpha}(1)$ to the prefix; now prefix = $\{(0.15,\,-0.35)\}$.

Initialize block 2: sample $\epsilon_2 = (1.2,\,-0.4)$.

Self‑condition with $\hat{\alpha}(1)$: the denoiser adds a small bias toward the prefix, yielding $\hat{\alpha}(2) \approx (0.6,\,-0.2)$ after two Euler steps.

Append $\hat{\alpha}(2)$; prefix now contains two blocks.

Repeat for block 3, obtaining $\hat{\alpha}(3) \approx (0.05,\,-0.01)$.

Even with a tiny toy example, the prefix quickly stabilizes the later blocks because each new block is denoised while constantly looking back at the already‑cleaned context.

Invert the standardization from Eq 5: compute $z_{\text{dec}} = \mu + s \odot \hat{z}$.

Feed $z_{\text{dec}}$ into the frozen query decoder, which produces vocabulary logits $o_{1:L}$.

Sample or arg‑max the logits to obtain the final token sequence.

SC‑CFG blends a “self‑aware” prediction (which already knows its own previous estimate) with a “self‑blind” prediction, letting a scalar weight steer how much the self‑conditioning signal influences the next update.

How does SC‑CFG differ from standard classifier‑free guidance?

Standard CFG swaps between a conditioned and an unconditioned model; SC‑CFG keeps the same conditioning (the prefix) but toggles the self‑conditioning signal, allowing the model to bias its own internal state rather than an external prompt.

CFG interpolates between a generation that sees the prompt prefix and one that ignores it, letting a weight $w_{\text{cond}}$ amplify the prompt’s influence.

Why does increasing $w_{\text{cond}}$ make the generated text follow the prompt more closely?

Because the term $w_{\text{cond}}(\hat{\alpha}(b)_{\text{cond}}-\hat{\alpha}(b)_{\text{uncond}})$ is added to the unconditioned estimate; a larger $w_{\text{cond}}$ magnifies the prompt‑driven direction, pulling the latent block toward the region of latent space that the prompt occupies.

Ablation Studies

We isolate each design choice to see how its removal harms performance.

The AURORA‑LM framework decouples text encoding from diffusion‑based generation via a query‑based autoencoder and a block‑causal diffusion model.

We run controlled ablations on OpenWebText (max length 128) by varying a single component at a time, generating 1,000 samples per setting with a 32‑step ODE sampler and measuring MAUVE against held‑out references.

Noisy‑input bottleneck width $D_b=128$ yields the highest MAUVE (0.808) among the tested widths.

Mean MAUVE across 16–64‑step ODE samplers peaks at $D_b=128$; narrower ($D_b=32$) under‑feeds the denoiser and wider ($D_b=1024$) offers no gain.

Direct clean‑latent regression ($x_0$ target with $x_0$‑space loss) achieves the best MAUVE (0.815) for both $D_b=128$ and $D_b=1024$.

Table 1 shows $0.815$ (for $D_b=128$) and $0.807$ (for $D_b=1024$), outperforming all $v$‑based alternatives.

**Figure 5. Evaluations of latent capacity.** Latent capacity improves decoder robustness and MAUVE. (a) Wider channels preserve token-recovery accuracy under stronger latent corruption; the broken horizontal axis compresses $\sigma \in (0.05, 0.20)$, where all widths remain near perfect recovery. (b) For each $D$, we report the highest MAUVE over the same $d \in \{1, 3, 5, 7\}$ search; annotations indicate the selected shift. (c) MAUVE over the practical sequence-compression range, expressed as the percentage of latent positions retained; the full stress-test sweep is reported in Appendix A.3.1.

**Figure 6.** Evaluations of noisy-input bottleneck and noise allocation. Moderate noisy-input width and high-noise calibration improve modeling of the full-width latent. (a) Bars show mean MAUVE over 16-, 32-, and 64-step ODE sampling; whiskers show the corresponding min–max range. The clean prediction target remains $D = 1024$ while only $D_b$ varies. (b) Tan-$d$ sweeps $d \in \{1, 3, 5, 7\}$ and logit-normal sweeps $\mu \in \{0, 1, 2\}$ at $s \in \{0.5, 1.0\}$; all settings are expressed by the common high-noise mass $m_{0.7} = \text{Pr}(\sigma > 0.7)$; the corresponding noise distributions are shown in Appendix A.3.2. MAUVE increases with high-noise allocation across all three sweeps; filled markers denote the largest tested value ($d = 7$ and $\mu = 2$).

**Figure 7. Evaluation over blockwise generation.** Self-trajectory consistency improves few-step sampling, while smaller blocks trade parallelism for higher generation quality. (a) MAUVE as a function of denoising-step budget, with and without self-trajectory consistency. Consistency improves MAUVE at every budget, with the largest gains in the few-step regime. (b) MAUVE for separately trained models across block sizes $Q \in \{4, 8, 16, 32, 64\}$, evaluated with 32-step ODE sampling. Smaller blocks provide more sequential conditioning stages and yield higher MAUVE; the filled marker denotes the selected $Q = 16$ operating point.

System-Level Comparison

AURORA‑LM‑S outperforms prior models on both free‑form and conditional generation.

Recall that AURORA‑LM separates text encoding from latent diffusion, using a query‑based autoencoder to obtain continuous latents and a block‑causal diffusion model to generate them.

These models emit tokens one after another, each step conditioned on every previously generated token.

Generation is cast as a finite sequence of denoising steps on a discrete latent space, gradually turning random symbols into coherent text.

AURORA‑LM‑S sets new state‑of‑the‑art scores on both free‑form and summarization tasks, achieving the lowest generation perplexity (23.56) and highest MAUVE (0.890) among all baselines, and the best ROUGE‑1/2/L (36.6/13.4/28.9) on XSum.

Table 2 reports Gen‑PPL 23.56 and MAUVE 0.890; Table 3 reports ROUGE‑1 36.6, ROUGE‑2 13.4, ROUGE‑L 28.9, all surpassing the next‑best systems.

The table compares the performance of various models using ROUGE-1, ROUGE-2, and ROUGE-L metrics.

**Table 3.** XSum conditional-generation comparison. AURORA-LM-S is evaluated under the matched protocol described in Appendix A.2.4. $^\dagger$Values collected by Hu et al. [23].

Scaling Evaluation

Scaling shows AURORA‑LM‑L beats a larger baseline across diverse tasks.

AURORA‑LM‑L achieves a macro average of 32.6 % on nine benchmarks, surpassing Cola‑DLM’s 25.1 %.

Table 4 reports the per‑task percentages and the macro average for both models.

Both systems were evaluated with identical prompting protocols: two‑shot examples where supported and the same inference code. AURORA‑LM‑L’s training budget of $1{,}500$ EFLOPs is comparable to Cola‑DLM’s reported 2{,}000 EFLOPs, yet its model contains fewer parameters and uses a block‑causal diffusion architecture.

**Table 4.** AURORA-LM-L achieves a macro average of 32.6, compared with 25.1 for Cola-DLM, and performs better on all nine benchmarks. These results indicate that the gains of AURORA-LM extend beyond the controlled research-scale setting: despite using a smaller diffusion Transformer, AURORA-LM-L consistently outperforms the larger publicly released latent-diffusion model across diverse language tasks.

Latent Bottleneck Analysis

Component ablations reveal how each design choice affects MAUVE quality.

We isolate each architectural or training choice and measure its effect on the MAUVE score, the primary quality metric for generated text.

Retaining the full latent representation yields the highest MAUVE; dropping to 70 % cuts quality by 0.144, confirming that the full‑width latent is essential. A high‑noise schedule (tan‑d, s=1.0, d=7) improves MAUVE by 0.029 over weaker schedules, showing that allocating more mass to high‑noise timesteps benefits generation. Direct regression of the clean latent ($x_0$) outperforms other target‑loss pairings by up to 0.086 MAUVE, indicating that matching the clean latent is the most effective supervision. Enabling self‑trajectory consistency raises MAUVE from 0.116 to 0.672 (+0.556), demonstrating its critical role in stabilizing denoising trajectories. Finally, smaller block sizes (4) achieve 0.909 MAUVE, whereas larger blocks degrade performance (down 0.278), suggesting that over‑parameterized blocks hurt sample quality.

Implementation Details

Implementation specifics for masking and loss‑space conversion in the autoencoder.

This appendix records the low‑level training tricks that the main text omits: how we corrupt inputs during autoencoder training and how we translate clean‑latent errors into velocity space.

Randomly zero out individual token embeddings before they enter the encoder, forcing the model to rely on surrounding context.

After encoding, randomly zero out individual dimensions of the latent tensor $z_{\text{enc}}$ to regularize the decoder.

Express the clean‑latent reconstruction error in velocity space so that low‑noise samples receive higher weight during training.

Noise Schedule Details

All hyper‑parameters, model sizes, and schedule details needed to reproduce the experiments.

This appendix records every configuration choice that underlies the results in the main paper: schedule‑mass metrics, the self‑trajectory consistency loss, data preprocessing, model dimensions, and training hyper‑parameters.

It is the fraction of diffusion‑time samples whose noise level exceeds a chosen threshold, used to compare how “noisy” different schedule families are.

Self‑Trajectory Consistency (STC) is an auxiliary loss applied to the block‑causal denoiser. It forces the model’s prediction at step $t$ to agree with the prediction at the preceding step $t'=\operatorname{PrevS}(t)$ after an Euler update, using an EMA‑averaged target. The weight $\lambda_{ct}$ ramps linearly from $0$ to $1.0$ over the first $T_{\text{warm}}$ optimization steps (10 000 steps for the S model, 5 000 for the L model), and the sampling budget is fixed at $S=16$.

Training‑Data Preprocessing follows the OpenWebText pipeline with the GPT‑2 BPE tokenizer. The controlled analysis uses 128‑token sequences (OWT128); the system‑level comparison packs 1 024‑token sequences, each ending with an EOS token and separating adjacent documents by EOS. For the large model we assemble a 294.7 B‑token mixture (161.3 B web, 50.3 B books, 83.1 B academic) and pack into 512‑token blocks using Qwen3 tokenization.

The architecture consists of a query‑based encoder‑decoder (six layers each) followed by a block‑causal denoiser built from a Transformer with RMSNorm, SwiGLU, RoPE, and per‑head QK normalization. Timestep conditioning uses sinusoidal embeddings projected through a two‑layer MLP and added to every Transformer layer.

**Table 5.** Architectures of AURORA-LM-S and AURORA-LM-L. AURORA-LM-S is used for the system-level comparisons, whereas AURORA-LM-L is used for the scaling evaluation. Model-size labels refer to the block-causal denoiser and exclude the autoencoder.

**Table 6.** Training configurations for AURORA-LM-S and AURORA-LM-L in Table 5. The underlying objectives are defined in §3 and Appendix A.1. Optimization steps, warmup durations, and EMA frequencies refer to parameter updates; global batch sizes are reported after gradient accumulation. Training compute is estimated using 3ND accounting, with one multiply-accumulate counted as one operation.

**Figure 8.** Training-time noise distributions used in the schedule ablation. Each panel shows the density over the flow-matching noise level $\sigma \in [0, 1]$. The dashed reference marks the high-noise region $\sigma > 0.7$ used in Fig. 6(b).

Training Configurations

Numerical settings and inference‑protocol details for AURORA‑LM experiments.

Latent standardization uses $\epsilon$ = 1 × 10⁻⁶, a decay to 1 × 10⁻⁵ after 1 000 steps, and a tan‑d shift of 7.0; instruction fine‑tuning runs for 10 000 steps (OWT) or 5 000 steps (XSum) with learning rates 4 × 10⁻⁴ (peak) and 1 × 10⁻⁶ (minimum). Training compute is ≈ 155 EFLOPs for OWT and ≈ 23 EFLOPs for XSum, while the larger XSum model requires ≈ 1 406 EFLOPs.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers