PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jiajun Liang, Yucheng Liao, Yukang Cao, Jiazhe Wei, Ken Li, Wende Tan, Jiankun Zhang, ZY Cui, Jingkang Yang, Liucheng Guo, Shiqi Yang, B. Yang, Caifeng Shan, Ziwei Liu, Chenyang Si

arXiv:2608.02602v1cs.CL

TL;DR

Text generation still relies mainly on discrete tokens, while continuous autoencoded approaches often simplify latents in ways that can reduce token-level fidelity. AURORA-LM preserves a high-capacity decodable latent and models its distribution with block-causal diffusion, outperforming evaluated baselines on free generation and summarization.

  • Problem

    Language generation remains predominantly discrete, while continuous autoencoded models often simplify decoder-facing latents, potentially removing distinctions needed for exact word, syntax, and ordering recovery.

  • Method

    AURORA-LM constructs a high-capacity, prefix-aligned continuous latent with a Query-based Encoder-Decoder and models its full-width distribution using block-causal flow-matching diffusion.

  • Results

    Under matched protocols, AURORA-LM outperforms evaluated autoregressive, discrete-diffusion, and continuous baselines on OpenWebText free generation and XSum summarization, surpassing a larger latent-diffusion model at 1B parameters.

  • Takeaways & Limitations

    A high-capacity, causally structured, decodable continuous text representation can effectively interface diffusion-based generation with discrete token decoding.

Abstract

from arXiv · show

Language remains an outlier in generative modeling: while images, video, and audio are increasingly modeled in continuous latent spaces, text generation still relies predominantly on discrete tokens. Existing continuous language models either inherit embedding spaces not designed for joint generation and decoding, or compress autoencoded latents to ease diffusion, sacrificing token-level fidelity. Instead of simplifying the representation to suit the generative model, we preserve a high-capacity, decodable text latent and design the diffusion model to learn its distribution directly. We introduce AURORA-LM, a continuous-latent diffusion language model that separates the construction of a decodable text representation from the modeling of its distribution. A Query-based Encoder-Decoder organizes text into a high-capacity, prefix-aligned latent sequence, and a Block-causal Diffusion Transformer learns its distribution through flow matching, generating blocks left to right while denoising positions within each block in parallel. Because such a latent is harder for diffusion to model, AURORA-LM restricts only the noisy-input pathway while retaining the full clean-latent prediction target, accommodating full-width latents without reducing decoder-facing capacity. We further calibrate the noise-level distribution to the latent width, and introduce self-trajectory consistency to bridge independently sampled training noise and iterative denoising at inference. AURORA-LM achieves the strongest performance among evaluated continuous and diffusion-based language models on OpenWebText free generation and XSum summarization. Scaling to 1B parameters with about 1500 EFLOPs of total compute yields further gains, surpassing a larger publicly released latent-diffusion language model under a matched evaluation protocol. All experiments are conducted on Ascend NPUs.

1 Introduction

AURORA-LM addresses the mismatch between high-capacity, decodable text representations and diffusion modeling by separating representation construction from latent distribution modeling. It combines a Query-based Encoder-Decoder with block-causal flow matching, preserving decoder-facing capacity while achieving strong benchmark performance.

  • Representation construction: The Query-based Encoder-Decoder constructs a high-capacity, causally ordered latent sequence that directly interfaces with token decoding.Its causal encoder reads progressively longer token prefixes, while text reconstruction trains the latent-to-token pathway.
  • Results: 1B parameters and approximately 1,500 EFLOPs of total compute yield further gains, with AURORA-LM surpassing a larger publicly released latent-diffusion language model under matched evaluation.The model outperforms evaluated baselines on OpenWebText free generation and XSum conditional summarization.
  • Framework: AURORA-LM separates text representation learning from generative distribution modeling instead of inheriting denoising spaces from embeddings or pretrained encoders.The framework explicitly constructs the continuous representation that the diffusion model learns to generate.
  • Representation construction: AURORA-LM retains full decoder-facing latent capacity and adapts the diffusion prior to model the resulting full-width latent distribution.The representation is not reduced merely to make generation easier, preserving information for accurate token recovery.
  • Distribution modeling: The block-causal denoiser uses flow matching to generate latent blocks left to right while jointly denoising positions within each block.A low-rank noisy-input pathway, full-width clean-latent prediction, width-calibrated noise allocation, and self-trajectory consistency support modeling and inference.

2 Related Work

Prior language-generation paradigms operate over discrete tokens, token-derived embeddings, or continuous encoder–decoder latents, each trading off sequential structure, denoising flexibility, or representation capacity. AURORA-LM adopts latent-space generation with a frozen autoencoder that preserves a high-capacity, decoder-facing representation rather than compressing it for diffusion tractability.

  • Autoregressive Language Models: Autoregressive models generate tokens strictly left to right, leveraging causal factorization and next-token prediction but constraining generation order.These properties have supported large-scale language-modeling progress while limiting parallel generation.
  • Discrete Diffusion Language Models: Discrete diffusion models corrupt and iteratively denoise token sequences, including masking-based approaches such as D3PMs and MDLM.This paradigm remains entirely within the discrete token vocabulary.
  • Embedding-space Language Models: Embedding-space methods denoise continuous sequences derived directly from token representations without a separately trained encoder–decoder.Diffusion-LM introduced controllable generation, DiffuSeq extended the formulation to conditional generation, and PLAID achieved competitive benchmark likelihood.
  • Latent-space Language Models: Latent-space models first map text into continuous representations and then model those representations with diffusion or flow matching.Existing approaches include compact fixed-length latents from compressed encoder features and full-length contextual representations without compression.
  • AURORA-LM: AURORA-LM constructs a high-capacity, decoder-facing latent for accurate token recovery instead of compressing the sequence, then freezes the autoencoder before denoiser training.Freezing fixes the target latent distribution and separates representation construction from diffusion-side design choices.

3 The AURORA-LM Framework

AURORA-LM separates decodable continuous-text representation construction from latent-distribution modeling using a Query-based Encoder-Decoder and a block-causal denoiser trained with flow matching. Its framework preserves full-width decoder-facing latents while generating latent blocks left to right and denoising positions within each block in parallel.

  • Framework overview: AURORA-LM combines a Query-based Encoder-Decoder with a block-causal denoiser to construct and model a decodable continuous text latent.The autoencoder is frozen after representation learning, and the denoiser models standardized latent outputs through blockwise flow matching.
  • Continuous text representation: The encoder’s progressively expanding token prefixes organize latent positions causally, aligning each latent prefix with the context used for block-causal conditioning.The query-based encoder aggregates variable-length token sequences into ordered continuous positions.
  • Continuous text representation: The latent sequence has N = round(cL) positions in D channels, where c = 1 gives no sequence compression and smaller c produces shorter latents.D and N jointly determine representation capacity through channel width and sequence length.
  • Block-causal denoising: Block-causal factorization generates latent blocks left to right while jointly denoising all positions within each block.This intermediate organization preserves causal order while retaining parallelism across positions inside a block.
  • Training objectives: The denoiser predicts each block’s full D-dimensional clean representation, using a low-rank bottleneck only for noisy inputs while preserving the full-width prediction output.Self-trajectory consistency addresses errors that can accumulate when successive denoising states receive inconsistent clean-latent estimates.

4 Experiment

Experiments evaluate AURORA-LM through controlled ablations, matched baseline comparisons, and scaling to a 1.01B-parameter model. Results support wide, full-width latents, calibrated high-noise training, self-trajectory consistency, and small denoising blocks as effective design choices.

  • Latent representation: Wider latent widths preserve token-recovery accuracy under stronger corruption and achieve their best calibrated generation performance at D = 1024.All widths reconstruct nearly perfectly without corruption, but wider representations remain robust over a broader corruption range; best MAUVE improves consistently with width.
  • Latent representation: Using the full latent sequence with D = 1024 avoids compression-related degradation, while c = 0.9 remains a practical efficiency alternative.Compression reduces positions and blockwise denoising stages while preserving causal prefix structure, but can degrade generation quality.
  • Training design: Db = 128 achieves the highest mean MAUVE when the clean prediction target remains fixed at D = 1024.The noisy-input bottleneck controls how much corrupted state the Transformer reads, while the clean target retains full decoder-facing width.
  • Training design: MAUVE consistently improves as schedules increase high-noise mass, whereas v-space and v-prediction objectives perform poorly in this setting.The experiments indicate that wider latent targets benefit from stronger high-noise calibration rather than a particular schedule parameterization.
  • Sampling and architecture: Self-trajectory consistency improves MAUVE at every tested point, while block-size increases reduce MAUVE from 0.909 at Q = 4 to 0.631 at Q = 64.Self-trajectory consistency addresses mismatch between independently trained noisy states and iterative inference; smaller blocks improve conditioning but require more sequential stages.
  • Benchmark results: AURORA-LM achieves Gen-PPL 23.56 and MAUVE 0.890, while AURORA-LM-L reaches a macro average of 32.6 versus 25.1 for Cola-DLM.AURORA-LM leads evaluated systems on unconditional generation and ROUGE metrics, and AURORA-LM-L performs better than Cola-DLM on all nine benchmarks.

5 Conclusion

AURORA-LM separates decodable representation learning from full-width latent distribution modeling using a causally structured continuous text representation. Controlled studies validate its design choices, supporting the representation as an interface between diffusion generation and discrete token decoding.

  • Core formulation: AURORA-LM decouples representation learning from distribution modeling in a continuous-latent diffusion language model.The approach separates construction of a decodable text representation from modeling its distribution.
  • Core formulation: A Query-based Encoder-Decoder constructs a high-capacity, causally ordered latent sequence for accurate token recovery and blockwise left-to-right generation.The autoencoder is frozen before a block-causal diffusion Transformer models the resulting full-width latent distribution through flow matching.
  • Design validation: Controlled studies support wider latents, noise calibration, a moderate noisy-input bottleneck, clean-latent prediction, high-noise training, self-trajectory consistency, and block-causal generation.Wider representations retain token-level information more robustly under corruption, while the bottleneck preserves the full-width clean prediction target.
  • Implications: The results establish high-capacity, causally structured, decodable continuous text representations as an effective interface between diffusion generation and discrete token decoding.The formulation is presented as a basis for scaling to longer contexts, broader language capabilities, and unified models spanning language and other continuous modalities.

Appendix · A Additional Method and Experimental Details

The appendix provides additional method and experimental details, covering implementation, evaluation, controlled-study results, and qualitative examples. It proceeds from method specifications through datasets, configurations, training, evaluation settings, and concludes with quantitative results and generated samples.

  • A Additional Method and Experimental Details: The appendix supplements the main text with implementation details.
  • A Additional Method and Experimental Details: It documents the evaluation protocols used in the experiments.
  • A Additional Method and Experimental Details: It reports complete controlled-study results.
  • A Additional Method and Experimental Details: It includes qualitative examples alongside the quantitative analyses.
  • A Additional Method and Experimental Details: The appendix presents datasets, model configurations, and training procedures.
  • A Additional Method and Experimental Details: It concludes with full quantitative results and generated samples.

A.1 Method and Implementation Details · A.1.1 Autoencoder Training-Time Masking

The autoencoder uses two masking operators during training: token-embedding dropout on valid input positions and latent dropout before decoding. Latent dropout masks latent coordinates independently with pz = 0.6 and applies no inverted-dropout rescaling.

  • A.1.1 Autoencoder Training-Time Masking: The implementation applies token-embedding dropout and latent dropout during autoencoder training.These are the two masking operators specified for the implementation.
  • A.1.1 Autoencoder Training-Time Masking: Token-embedding dropout independently zeros the input vector at each valid token position.For each position j, the mask is sampled as mj ∼ Bernoulli(px).
  • A.1.1 Autoencoder Training-Time Masking: Latent dropout independently masks each coordinate of the encoded latent zenc before decoding.The encoded latent has shape zenc ∈ R^N×D.
  • A.1.1 Autoencoder Training-Time Masking: pz = 0.6 is used for latent dropout.The keep-mask variables satisfy rik ∼ Bernoulli(1 − pz).
  • A.1.1 Autoencoder Training-Time Masking: The latent mask is applied elementwise as ezenc = zenc ⊙ r.Masking spans all latent positions i ∈ {1, . . . , N} and coordinates k ∈ {1, . . . , D}.
  • A.1.1 Autoencoder Training-Time Masking: Latent dropout uses no inverted-dropout rescaling.The masked latent is therefore retained without rescaling after coordinate masking.

A.1.2 Loss-Space Conversion for Clean-Latent Prediction

The section converts flow-matching notation to clean-latent prediction by expressing noisy latents through the interpolation coefficient σ. Evaluating clean-latent error in velocity space scales it by 1/σ^2, emphasizing lower-noise samples.

  • Notation and conversion: Flow matching uses σ as the interpolation coefficient, with x_σ = (1 − σ)z + σε and velocity v = ε − z.Here, z is the clean latent and ε ∼ N(0, I).
  • Loss-space effect: 1/σ^2 scales clean-latent prediction error when evaluated in velocity space, giving greater weight to lower-noise samples.This is the loss-space effect examined in Section 4.1.2.

A.1.3 High-Noise Mass of the Schedule Families

This section compares schedule families using the high-noise mass mσ⋆ = Pr(σ > σ⋆) at a threshold σ⋆ ∈ (0, 1). It applies this coordinate to tan-d and logit-normal noise-level distributions, including the shared comparison point σ⋆ = 0.7.

  • High-noise mass: The comparison metric is the high-noise mass mσ⋆ = Pr(σ > σ⋆) for σ⋆ ∈ (0, 1).This metric is used to compare different noise-level distributions.
  • Tan-d schedule: The main-text tan-d schedule uses a shifted cosine parameterization with τ sampled uniformly from (0, 1).The resulting high-noise mass is evaluated under the tan-d mapping σ(τ; d).
  • Logit-normal schedule: For logit-normal sampling, ℓ follows N(µ, s2) and σ = sigmoid(ℓ), so monotonicity of sigmoid determines the corresponding high-noise mass.The sigmoid mapping is monotone, enabling the distributional comparison in high-noise-mass coordinates.
  • Schedule comparison: At σ⋆ = 0.7, m0.7 provides the common coordinate for comparing tan-d and logit-normal sweeps.This is the comparison coordinate used in Sec. 4.1.2.

A.1.4 Self-Trajectory Consistency Implementation

Self-trajectory consistency is implemented as an auxiliary objective for the block-causal denoiser trained with LFM, aligning clean-latent predictions across successive model-induced sampling states. The reported configuration uses a 16-step sampling budget and linearly warms the consistency weight from zero.

  • Consistency objective: Self-trajectory consistency aligns clean-latent predictions at successive states along the model-induced sampling trajectory without training a separate consistency model.It is implemented as an auxiliary objective for the block-causal denoiser trained with LFM.
  • Sampling schedule: S = 16 fixes the sampling budget, with the scheduler selecting t′ = PrevS(t) for sampled flow times t > 0.The scheduler uses the same schedule parameterization as sampling.
  • Sampling schedule: The current model predicts ˆzθ(αt, t), while its stop-gradient prediction produces the neighboring state eαt′.These predictions define the successive-state consistency relation.
  • Weight warmup: The consistency weight increases linearly from zero to its final value λct over the first Twarm optimization steps.The reported runs use Twarm = 10,000 steps for AURORA-LM-S and Twarm = 5,000 steps for AURORA-LM-L.

A.2 Experimental Protocol and System Configuration

This section defines the experimental setup and evaluation procedures, covering preprocessing, architecture, training, generation, and benchmarking. It also describes the inference sweep for selecting the OpenWebText configuration and the procedures for reproducing baselines and evaluating public benchmarks.

  • The protocol covers data preprocessing, model architecture, training configuration, and generation and evaluation procedures.
  • An inference sweep selects the system-level OpenWebText configuration used in the experiments.
  • The evaluation procedures include baseline reproduction and public-benchmark assessment.

A.2.1 Training-Data Preprocessing … A.6 Few-Shot Benchmark Demonstrations

The appendix specifies data, architecture, training and evaluation protocols, then provides controlled analyses, qualitative examples, and few-shot demonstrations. It selects inference settings, documents baseline reproduction and scoring, and reports supporting numerical and qualitative evidence across the evaluated tasks.

  • A.2.1 Training-Data Preprocessing: OpenWebText data use GPT-2 BPE tokenization, with 128-token controlled-analysis sequences, 1,024-token packed system-comparison sequences, and a 294.7B-token AURORA-LM-L mixture.The mixture contains 161.3B web tokens (54.7%), 50.3B books and encyclopedia tokens (17.1%), and 83.1B academic, mathematical, and knowledge tokens (28.2%).
  • A.2.2 Model Architecture: The architecture freezes the trained autoencoder before denoiser training and uses a Transformer denoiser with RMSNorm, SwiGLU, RoPE, QK normalization, and timestep conditioning.The timestep signal uses sinusoidal embeddings projected through a two-layer MLP and added to every Transformer layer.
  • A.2.3 Training Configuration: Training configurations for AURORA-LM-S and AURORA-LM-L define optimization, batching, EMA, and compute accounting, with model-size labels excluding the autoencoder.Training compute uses 3ND accounting, counting one multiply–accumulate as one operation.
  • A.2.4 Generation and Evaluation Protocols: Generation protocols use task-specific samplers and guidance: 64-step SDE-DPM++ with wsc = 4 for OpenWebText, 32-step ODE with w = 2.0 for XSum, and 16-step benchmark sampling.Prefix-conditioned generation encodes a clean latent prefix, uses complete blocks as context, and applies clean repainting within a boundary block.
  • A.2.5 System-Level Inference-Protocol Selection: The OpenWebText sweep selects 64-step SDE-DPM++ with wsc = 4 because it achieves the lowest mean Gen-PPL while retaining MAUVE of 0.890.The sweep varies denoising steps, solver, and SC-CFG scale at the fixed 250k checkpoint, using 1,000 samples per seed across five seeds.
  • A.2.6 XSum Inference-Protocol Selection: For XSum, w = 2.0 gives the highest mean ROUGE-1 and ROUGE-L, while 32 ODE steps match 64 steps in ROUGE-2 after rounding, so that setting is selected.Increasing guidance beyond the moderate range degrades all three ROUGE metrics.
  • A.2.7 Reproduction Protocol for Baselines: Baseline comparisons evaluate released Duo and ELF checkpoints without end-to-end retraining, while AURORA-LM-S uses a 130M-parameter denoiser versus approximately 92M discrete-diffusion backbones.AURORA-LM-S matches baseline hidden size 768 and depth 12, with the parameter difference primarily arising from SwiGLU and conditioning projections.
  • A.2.8 Scoring and Answer Extraction; A.2.9 Few-Shot Benchmark Protocol: Scoring uses fixed GPT-2 Large Gen-PPL and held-out OpenWebText references for MAUVE, ROUGE for XSum, and shared truncation and answer-matching procedures for benchmarks.Few-shot tasks use two-shot prompts, textual multiple-choice answers, and parsers accepting either option letters or option text.
Loading 2608.02602v1…