Summary (Overview)

• AURORA-LM decouples continuous text representation learning (autoencoder) from generative distribution modeling (diffusion), enabling high-capacity, decodable latents without sacrificing generation quality.
• It constructs a prefix-ordered, high-channel latent sequence via a Query-based Encoder-Decoder, and models its distribution with a block-causal diffusion Transformer trained via flow matching.
• Three key innovations adapt the diffusion model to handle the full-width latent: a low-rank noisy-input bottleneck, noise-level calibration to latent width, and self-trajectory consistency to align predictions along the denoising trajectory.
• Under matched evaluation, AURORA-LM achieves the lowest Gen-PPL (23.56) and highest MAUVE (0.890) on OpenWebText unconditional generation, and best ROUGE scores on XSum conditional summarization among all evaluated continuous and discrete diffusion baselines.
• Scaling to 1B parameters (~1,500 EFLOPs) yields consistent gains over a larger publicly released latent-diffusion language model (1.8B parameters) across nine language benchmarks.

Introduction and Theoretical Foundation

Continuous latent spaces have become a universal interface for generative modeling in vision, video, and audio (e.g., diffusion models). However, language remains an outlier—modern language models rely on discrete token representations (Brown et al., 2020; Touvron et al., 2023). Existing continuous language models either inherit embedding spaces not designed for joint generation and decoding (Diffusion-LM, PLAID) or compress autoencoded latents to ease diffusion at the cost of token-level fidelity (LD4LG, COSMOS, Cola-DLM). This creates a tension: compact latents are easier to generate but lose information needed for accurate reconstruction; wide latents preserve fidelity but make the prior distribution more complex.

AURORA-LM challenges this design compromise. Instead of simplifying the representation to accommodate the generative model, it preserves a high-capacity, decodable text latent and designs the diffusion model to learn its distribution directly. The key insight is to separate the construction of a decodable text representation (stage 1) from the modeling of its generative distribution (stage 2).

Methodology

1. Continuous Text Latent Construction (Query-based Encoder-Decoder)

Given a token sequence $w_{1:L}$, the encoder produces a continuous latent sequence $\bar{z}_{\text{enc}} \in \mathbb{R}^{N \times D}$ with $N = \text{round}(cL)$ ($c$ is retention ratio). The encoder uses $N$ learnable queries $q \in \mathbb{R}^D$ that aggregate information from causally expanding token prefixes $E[w_{1:\lceil iL/N\rceil}]$ and preceding latent states, inducing a prefix order. The decoder uses $L$ queries to reconstruct tokens from latent prefixes, with matched visibility (position $j$ can only attend to latent prefix $z_{\text{enc},1:\lfloor (j-1)N/L\rfloor + 1}$). The autoencoder is trained with cross-entropy and regularization (token-embedding dropout $p_x=0.3$, latent dropout $p_z=0.6$). After training, the autoencoder is frozen and its outputs are standardized per-channel: $z = (z_{\text{enc}} - \mu) \oslash s$.

2. Block-Causal Modeling of Latent Distribution

The latent sequence is partitioned into $B = \lceil N/Q \rceil$ contiguous blocks of size $Q$. The prior is factorized as:

$$p_\theta(\alpha) = \prod_{b=1}^B p_\theta\bigl(\alpha^{(b)} \mid \alpha^{(<b)}\bigr), \qquad \alpha^{(<b)} = (\alpha^{(1)}, \ldots, \alpha^{(b-1)}).$$

Each block's conditional distribution is learned via flow matching. For a clean block $\alpha^{(b)}$ and noise $\varepsilon^{(b)} \sim \mathcal{N}(0,I)$, the linear path is:

$$\alpha_t^{(b)} = (1 - t) \alpha^{(b)} + t \varepsilon^{(b)}, \quad t \in [0,1].$$

The model predicts the clean endpoint $\hat{\alpha}_\theta^{(b)} = f_\theta(\alpha_t^{(b)}, t; \alpha^{(<b)})$ using a two-stream attention mask: the denoiser attends to its own noisy positions and the clean prefix but not to future clean blocks or other noisy states. This allows parallel training of all block conditionals.

3. Learning the Full-Width Latent Distribution

The frozen decoder expects a full-width clean latent $D$. To avoid processing the noisy input at full width, a low-rank input bottleneck is applied:

$$W_{\text{in}} = W_{\text{up}} W_{\text{down}}, \quad W_{\text{down}} \in \mathbb{R}^{D_b \times D}, \quad W_{\text{up}} \in \mathbb{R}^{H \times D_b},$$

with $D_b < \min(D,H)$. The output head still predicts the full $D$-dimensional clean block. The noise-level distribution is calibrated to latent width using a tan-$d$ parameterization (Eq. 20), with larger $d$ assigning more training mass to high-noise states. The flow-matching loss is:

$$\mathcal{L}_{\text{FM}} = \mathbb{E}_{z\sim q_\mathcal{E}, t\sim \pi, \varepsilon\sim \mathcal{N}(0,I)}\left[\frac{1}{D|\mathcal{J}|}\sum_{i\in\mathcal{J}}\| \hat{z}_{\theta,i} - z_i\|_2^2\right].$$

Self-conditioning feeds the clean-latent prediction from the previous sampling step back into the model with probability $p_{\text{sc}}$. Self-trajectory consistency aligns clean-latent predictions at neighboring states along the model's own sampling trajectory. Given current time $t$, the lower-noise neighbor $t' = \text{Prev}_S(t) < t$ is computed, and the updated state is:

$$\widetilde{\alpha}_{t'} = \frac{t'}{t} \alpha_t + \left(1 - \frac{t'}{t}\right) \text{sg}(\hat{z}_\theta(\alpha_t, t)).$$

The consistency loss is:

$$\mathcal{L}_{\text{ct}} = \mathbb{E}_{z\sim q_\mathcal{E}, t\sim \pi, \varepsilon\sim \mathcal{N}(0,I)}\left[ \frac{1}{D|\mathcal{J}|} \sum_{i\in\mathcal{J}} \| \hat{z}_\theta(\alpha_t, t)_i - \text{sg}(\hat{z}_{\theta_{\text{EMA}}}(\widetilde{\alpha}_{t'}, t')_i)\|_2^2 \right].$$

The total prior objective is $\mathcal{L}_{\text{train}} = \mathcal{L}_{\text{FM}} + \lambda_{\text{ct}} \mathcal{L}_{\text{ct}}$.

4. Generation and Decoding

During inference, latent blocks are generated from left to right. Each block is initialized as Gaussian noise and denoised via ODE/SDE solver, conditioned on the completed prefix. For unconditional generation, the prefix is empty; for prompt-conditioned generation, the prompt is encoded and standardized. Self-conditioning classifier-free guidance (SC-CFG) and standard classifier-free guidance (CFG) are used for unconditional and conditional generation, respectively. The generated latent $\hat{z}$ is transformed back to decoder space: $z_{\text{dec}} = \mu + s \odot \hat{z}$, and the frozen decoder recovers token logits.

Empirical Validation / Results

Controlled Ablations (OWT128)

• Latent capacity: Wider channels ($D=1024$) preserve token-recovery accuracy under stronger corruption and yield higher MAUVE after noise calibration (Fig. 5). Sequence compression causes quality degradation (Table 12).
• Noisy-input bottleneck: $D_b=128$ gives highest mean MAUVE across 16–64 step ODE sampling (Fig. 6a).
• Noise allocation: MAUVE improves with more training mass on high-noise states ($m_{0.7}$) across tan-$d$ and logit-normal families (Fig. 6b). Best: tan-$d$ with $d=7$.
• Prediction target: Clean-latent ($x_0$) prediction with $x_0$-space loss achieves highest MAUVE (Table 1).
• Self-trajectory consistency: Improves MAUVE at every step budget, especially in few-step regime (Fig. 7a).
• Block size: Smaller blocks ($Q=4$) give higher MAUVE but more sequential stages; $Q=16$ is selected as a balanced point (Fig. 7b).

System-Level Comparison

Unconditional generation (OpenWebText, 1024 tokens):

Conditional summarization (XSum):

Scaling Evaluation (1B parameters)

AURORA-LM-L (1.01B denoiser, total ~1.34B) trained on 294.7B tokens (~1,500 EFLOPs) compared to Cola-DLM (1.8B DiT, ~2,000 EFLOPs). AURORA-LM-L uses 16-step ODE, CFG scale 3.0.

Theoretical and Practical Implications

• Theoretical: The results demonstrate that continuous language generation can effectively bridge diffusion-based generative modeling and discrete token decoding through a high-capacity, causally structured, decodable text representation. The decoupling of representation construction (autoencoder) and distribution modeling (diffusion) allows each component to be optimized independently, avoiding the representational compromise that limited prior work.
• Practical: The low-rank input bottleneck, noise-level calibration, and self-trajectory consistency are general techniques that can be applied to other latent diffusion models. The block-causal formulation provides a tunable trade-off between parallelism and conditioning quality. The strong scaling results suggest that continuous latent language models can compete with and potentially surpass discrete diffusion models at scale, while offering a natural interface for multimodal generation.

Conclusion

AURORA-LM presents a unified framework for continuous-latent diffusion language modeling that decouples representation learning from generative modeling. It constructs a high-capacity, prefix-ordered latent sequence via a Query-based Encoder-Decoder, and learns its distribution with a block-causal diffusion Transformer using flow matching. Key innovations—low-rank noisy-input bottleneck, noise-level calibration, and self-trajectory consistency—enable the model to handle the full-width latent without reducing decoder-facing capacity. Under matched evaluation, AURORA-LM achieves state-of-the-art performance among continuous and diffusion-based LMs on OpenWebText and XSum, and scales to 1B parameters with results surpassing a larger latent-diffusion model. This work establishes that a decodable, high-capacity continuous text representation can serve as an effective interface between diffusion-based generation and discrete token decoding, providing a promising foundation for scaling continuous language models and for unified multimodal generation.