Technical Summary: AURORA-LM

Problem Statement

Language generation remains an outlier in modern generative modeling. While images, video, and audio are increasingly modeled in continuous latent spaces using diffusion processes, text generation still relies predominantly on discrete tokens. Existing attempts to apply continuous diffusion to language face a fundamental design compromise:

1. Inherited Representations: Methods like Diffusion-LM or TEncDM inherit embedding spaces from token embeddings or pre-trained encoders. These spaces are not explicitly designed for the joint requirements of accurate text reconstruction and effective continuous generation.
2. Compression Trade-offs: Autoencoder-based approaches (e.g., LD4LG, COSMOS) map text to continuous latents but often compress these representations to make the diffusion distribution easier to learn. This compression sacrifices token-level fidelity, creating a bottleneck where improving generative tractability degrades decoding accuracy, and vice versa.

The core challenge is to construct a continuous text representation that supports faithful, accurate decoding while designing a generative model capable of learning its distribution without sacrificing that representational capacity.

Methodology: AURORA-LM Framework

AURORA-LM (Autoencoding Unified Representation for Continuous-Latent Diffusion Language Modeling) addresses these challenges by decoupling representation construction from distribution modeling into two distinct stages.

1. Continuous Text Latent Construction (The Autoencoder)

Instead of inheriting representations, AURORA-LM explicitly constructs a high-capacity, decodable text latent using a Query-based Encoder-Decoder:

• Query-based Encoder: Organizes discrete text into a high-capacity, prefix-aligned latent sequence. It uses $N$ latent queries to aggregate token prefixes of increasing length. As the query index $i$ increases, the encoder attends to progressively longer token prefixes ($w_{1:\lceil iL/N \rceil}$), inducing a causal order where latent positions align with textual prefixes.
• Query-based Decoder: Maps latent prefixes back to token logits. Crucially, the decoder is restricted to attending only to the corresponding latent prefix, ensuring a strict causal correspondence between latent and token sequences.
• High-Capacity Design: Unlike prior works that compress latents for generative ease, AURORA-LM retains a high channel width ($D$) and minimal sequence compression ($c \approx 1$). This preserves the information necessary for exact token recovery, treating the latent space as a faithful interface to text rather than a simplified prior.
• Training: The autoencoder is trained via token-level cross-entropy and then frozen. This fixes the target latent distribution, allowing the generative model to be studied independently of representation learning.

2. Block-Causal Latent Diffusion (The Generative Model)

A Block-Causal Diffusion Transformer learns the distribution of the frozen latent sequences using Flow Matching:

• Block-Causal Factorization: To balance causal generation with parallel denoising, the latent sequence is partitioned into blocks. The model generates blocks from left to right (causal conditioning on previous blocks) but denoises all positions within a block in parallel.
• Full-Width Latent Prediction: The model predicts the full-width clean latent ($x_0$) required by the decoder, rather than a compressed version.
• Noisy-Input Bottleneck: To handle the difficulty of learning a high-dimensional distribution from noisy inputs, the model applies a low-rank projection only to the noisy input pathway ($W_{down} \in \mathbb{R}^{D_b \times D}$), while the output head predicts the full-width clean target. This allows the model to process corrupted inputs efficiently without reducing the decoder-facing capacity.
• Noise Calibration: The noise-level distribution is calibrated to the latent width. Empirical results show that wider latent representations benefit from noise schedules that place greater emphasis on high-noise inputs.
• Self-Trajectory Consistency: To bridge the gap between training on independently sampled noisy states and inference via iterative denoising, the model introduces a consistency loss. This aligns clean-latent predictions at neighboring states along the denoising trajectory, stabilizing the generation process and enabling high-quality results with fewer sampling steps.

3. Inference

During inference, latent blocks are generated sequentially from left to right. The generated standardized latents are inverted to the original encoder space and passed to the frozen decoder to recover the final token sequence. The framework supports both unconditional generation and prompt-conditioned generation (using classifier-free guidance or self-conditioning guidance).

Key Contributions

1. Unified Framework: AURORA-LM separates text representation learning from generative distribution modeling, explicitly constructing a continuous representation that the diffusion model learns to generate, rather than inheriting one from token embeddings.
2. High-Capacity Interface: It retains the full channel capacity of the latent representation for accurate token decoding, adapting the diffusion prior to model the resulting complex distribution rather than simplifying the representation.
3. Block-Causal Flow Matching: It introduces a block-causal denoiser trained with flow matching that generates blocks left-to-right while denoising positions within blocks in parallel.
4. Architectural Innovations: The combination of a low-rank noisy-input pathway, noise-level calibration to latent width, and self-trajectory consistency allows the model to effectively learn full-width latent distributions.

Experimental Results

Experiments were conducted on Ascend NPUs.

• OpenWebText (Unconditional Generation): AURORA-LM-S (130M parameters) achieved the lowest Generation Perplexity (Gen-PPL: 23.56) and highest MAUVE (0.890) among evaluated autoregressive, discrete-diffusion, and continuous-embedding baselines.
• XSum (Conditional Summarization): AURORA-LM-S outperformed all baselines (including ELF-B, AR, and MDLM) on ROUGE-1, ROUGE-2, and ROUGE-L metrics.
• Scaling: AURORA-LM-L (1.01B parameters, ~1,500 EFLOPs compute) demonstrated consistent performance gains and surpassed a larger publicly released latent-diffusion model (Cola-DLM, ~1.8B parameters) under a matched evaluation protocol.
• Ablation Studies:
  • Latent Width: Wider latents ($D=1024$) provided better robustness to corruption and higher generative performance.
  • Input Bottleneck: A moderate bottleneck width ($D_b=128$) for noisy inputs was optimal, allowing full-width prediction without processing corrupted inputs at full dimensionality.
  • Trajectory Consistency: Self-trajectory consistency significantly improved few-step sampling quality.

Significance and Claims

The paper claims that continuous language generation can effectively bridge diffusion-based generative modeling and discrete token decoding. The authors argue that the prevailing design compromise—simplifying representations to accommodate generative models—is unnecessary. Instead, by preserving a high-capacity, decodable text latent and designing the diffusion model to learn its distribution directly (via specific architectural choices like block-causal factorization and trajectory consistency), it is possible to achieve state-of-the-art performance in both unconditional and conditional text generation tasks. The results suggest that continuous spaces can serve as a general interface for language, retaining semantic fidelity and syntactic regularity while leveraging the benefits of diffusion modeling.