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

Core Idea§

AURORA-LM keeps a high-capacity, decodable continuous latent for text and makes the diffusion model adapt to it, rather than compressing the latent to simplify generation. It splits the problem into two independently designed components: a query-based encoder-decoder that builds a decodable latent sequence, and a block-causal diffusion transformer that learns to generate that sequence through flow matching.

Latent Representation: Query-based Encoder-Decoder (QED)§

Instead of vector-quantized tokens or heavily downsampled autoencoder latents, QED produces a latent sequence $Z = [z_1, \dots, z_N]$ with the same width as the model hidden state. Learnable query vectors attend to the source text via cross-attention and are organized as blocks. The latents are prefix-aligned: each block only depends on previous blocks, so the model can generate blocks autoregressively while preserving full token-level fidelity for later decoding.

Generation: Block-causal Diffusion Transformer (BC-DiT)§

Diffusion operates directly on these full-width latents using flow matching. The forward process interpolates between a clean latent $z_0$ and noise $z_1$:

$$ z_t = (1 - t) z_0 + t z_1, \quad t \in [0, 1]. $$

The network predicts the velocity field $v_\theta(z_t, t)$, trained with:

$$ L_{FM} = \mathbb{E}_{t, z_0, z_1} [\| v_\theta(z_t, t) - (z_1 - z_0) \|^2]. $$

The transformer is block-causal: previously generated blocks are clean and used as context, while positions within the current block are denoised in parallel. This hybrid of autoregressive and diffusion decoding balances quality and latency.

To avoid the usual trick of reducing latent width for tractable diffusion, AURORA-LM restricts only the noisy-input pathway—noisy latents pass through a lightweight adapter before entering the transformer—while the output head predicts the full clean latent. This preserves decoder-facing capacity and avoids token-level reconstruction loss.

Additional Training Techniques§

• Noise-level calibration: the noise-time distribution $p(t)$ is adjusted according to latent width, preventing extremely high-dimensional latents from being dominated by near-pure-noise steps.
• Self-trajectory consistency: during training, the model's own one-step denoised prediction is used to synthesize a subsequent noisy point, mimicking iterative inference and reducing the gap between single-step noise sampling and multi-step denoising.

Illustrative Training Loop§

def train_step(model, encoder, decoder, x, t_dist):
    z0 = encoder(x)          # [B, N, D] decodable latent
    z1 = torch.randn_like(z0)
    t = t_dist.sample(z0.shape[-1])  # width-calibrated noise level
    zt = (1 - t) * z0 + t * z1

    v_pred = model(zt, t)             # noisy input through adapter
    loss = F.mse_loss(v_pred, z1 - z0)

    # self-trajectory consistency
    z0_hat = zt - t * v_pred
    t_next = (t + delta_t).clamp(max=1.0)
    zt_next = (1 - t_next) * z0_hat + t_next * z1
    v_pred_next = model(zt_next, t_next)
    loss = loss + lam * F.mse_loss(v_pred_next, z1 - z0_hat)
    return loss

Results§

On OpenWebText free-form generation and XSum summarization, AURORA-LM achieves the strongest performance among evaluated continuous and diffusion-based language models. Scaling to 1B parameters with roughly 1500 EFLOPs of total compute improves quality further and surpasses a larger publicly released latent-diffusion language model under a matched evaluation protocol. All experiments were run on Ascend NPUs.