AURORA‑LM: Continuous‑Latent Diffusion Language Model on Ascend NPU

Background

Following the ELF (Embedded Language Flows) work by He Kaiming’s team, the PRLab at Nanjing University, together with collaborators from Nanyang Technological University and Imperial College London, introduced AURORA‑LM (Autoencoding Unified Representation for Continuous‑Latent Diffusion Language Modeling). All experiments were conducted on Ascend NPU, and the model was scaled to roughly 1 B parameters to demonstrate feasibility on a domestic AI‑accelerator platform.

Why Model Language in Continuous Space?

Language is inherently discrete, composed of tokens. Even when generation moves to a continuous latent space, the final output must be mapped back to tokens. The latent sequence therefore needs to retain enough information for accurate decoding while remaining learnable. Constructing such a latent representation is the core challenge.

Two‑Step Solution

AURORA‑LM decomposes the problem into:

Step 1: Train a fully re‑trained auto‑encoder that converts text into a semantically rich, decodable continuous vector sequence.

Step 2: Use a specially designed “block‑causal diffusion model” to learn the generation distribution of the encoded vector sequence.

The model also incorporates designs for noisy‑state inputs, training procedures, and few‑step sampling to ensure stable and efficient modeling in a larger semantic space.

Model Architecture

The first stage builds a prefix‑ordered continuous latent sequence. As the latent position moves forward, the readable token prefix grows, mirroring left‑to‑right text expansion. The encoder–decoder is trained to reconstruct the original tokens.

In the second stage, the latent sequence is divided into consecutive blocks. Blocks attend causally from left to right, while intra‑block attention is bidirectional. For each noisy latent block, the model predicts the corresponding clean latent; concatenated predictions are compared against the full clean latent to compute training loss.

High‑Capacity Latent Vectors

Increasing the dimensionality of each latent vector preserves richer semantic information under noise, helping decoding of words, syntax, and local order. Experiments show that narrowing the noisy‑input dimension (a bottleneck) and raising the probability of high‑noise inputs both improve generation quality, though higher dimensionality also makes the diffusion task more complex.

Self‑Trajectory Consistency for Few‑Step Denoising

To mitigate error accumulation when using few denoising steps, AURORA‑LM introduces self‑trajectory consistency. During training, the model predicts a clean latent from a noisy input, then advances to a lower‑noise state and enforces consistency between successive clean latent estimates. This consistency loss is weighted by a hyper‑parameter and jointly optimized with the flow‑matching loss.

Experimental Results

Free‑generation experiments on OpenWebText compare AURORA‑LM with autoregressive, discrete diffusion, and other continuous models. AURORA‑LM achieves a Gen‑PPL of 23.56 and a MAUVE score of 0.890, the best among all methods.

For conditional summarization on XSum, AURORA‑LM attains ROUGE‑1/2/L scores of 36.6, 13.4, and 28.9 respectively, again the highest reported.

Scaling the model to ~1 B parameters (≈1500 EFLOPs of training) and evaluating on nine language benchmarks (including commonsense reasoning and reading comprehension) yields an average score of 32.6, surpassing a publicly available larger continuous language model on all tasks.

Conclusion

AURORA‑LM first learns high‑capacity continuous latents that can be accurately mapped back to tokens, then models their distribution with a block‑causal diffusion model. The combined input, training, and sampling designs enable continuous diffusion generation to connect with discrete token decoding, offering a promising path toward longer contexts and larger‑scale continuous language models.