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

> AURORA-LM decouples text autoencoding from diffusion modeling, achieving state-of-the-art generation quality (Gen-PPL 23.56, MAUVE 0.890) and scaling to 1B parameters.

- **Source:** [arXiv](https://arxiv.org/abs/2608.02602)
- **Published:** 2026-08-06
- **Permalink:** https://picx.dev/p/LJkeTS
- **Whiteboard:** https://picx.dev/p/LJkeTS/image

## Summary

## 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)**:

| Model | Backbone | Gen-PPL ↓ | Entropy ↑ | MAUVE ↑ |
|-------|----------|-----------|-----------|---------|
| AR | 85M | 39.40 | 5.605 | 0.851 |
| Duo | 92M | 86.57 | 5.566 | 0.704 |
| Duo-distilled | 92M | 78.22 | 5.574 | 0.715 |
| SEDD | 92M | 119.82 | 5.644 | 0.693 |
| MDLM | 92M | 121.36 | 5.664 | 0.668 |
| ELF-B | 105M | 24.11 | 5.155 | 0.229 |
| **AURORA-LM-S** | 130M | **23.56** | 5.241 | **0.890** |

**Conditional summarization (XSum)**:

| Model | ROUGE-1 ↑ | ROUGE-2 ↑ | ROUGE-L ↑ |
|-------|-----------|-----------|-----------|
| ELF-B | 36.0 | 12.2 | 27.8 |
| AR | 30.5 | 10.2 | 24.4 |
| MDLM | 33.4 | 11.6 | 25.8 |
| Duo | 31.4 | 10.1 | 25.0 |
| E2D2 | 28.4 | 8.3 | 22.0 |
| SeqDiffuSeq | 19.3 | 1.7 | 14.1 |
| **AURORA-LM-S** | **36.6** | **13.4** | **28.9** |

### 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.

| Model | Avg ↑ | MMLU | ARC-C | OBQA | HellaSwag | WinoGrande | StoryCloze | SIQA | RACE | SQuAD EM |
|-------|-------|------|-------|------|-----------|------------|------------|------|------|----------|
| AURORA-LM-L (1B) | **32.6** | 22.2 | 21.2 | 27.8 | 18.4 | 50.3 | 54.8 | 30.2 | 30.6 | 38.2 |
| Cola-DLM (1.8B) | 25.1 | 19.6 | 20.6 | 24.2 | 5.7 | 45.3 | 33.8 | 26.8 | 24.2 | 25.7 |

## 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.

---

_Markdown view of https://picx.dev/p/LJkeTS, served by PicX — AI-generated visual whiteboard summaries of research papers._