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

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.

한 줄 요약

AURORA-LM은 연속 잠재 공간에서 텍스트 생성을 위한 새로운 확산 언어 모델로, 블록 인과 트랜스포머와 쿼리 기반 인코더-디코더를 결합하여 높은 정확도를 달성한다.

핵심 기여도

쿼리 기반 인코더-디코더(Query-based Encoder-Decoder)를 도입하여 prefix-aligned, high-capacity 텍스트 잠재 시퀀스를 구성.
블록 인과 확산 트랜스포머(Block-causal Diffusion Transformer)를 통해 flow matching 기반의 분포 학습을 수행.
노이즈 입력 경로를 제한하면서 전체 clean-latent 예측 타겟을 유지, 데코더 측 용량 저하 없이 full-width 잠재 공간 처리.
1B 파라미터, 1500 EFLOPs 계산량에서 기존 latent-diffusion 모델을 상회하는 성능 달성.

핵심 아이디어

기존 연속 언어 모델은 잠재 공간을 단순화하거나 별도의 임베딩 공간을 사용해 토큰 수준의 정확도를 희생하는 경향이 있다. AURORA-LM은 대신 높은 용량의 디코더 가능한 텍스트 잠재 공간을 유지하고, 확산 모델이 이를 직접 학습하도록 설계한다. 이는 텍스트 생성 과정에서 잠재 표현의 복잡성과 생성 정확도를 동시에 달성하려는 시도이다. 핵심 아이디어는 잠재 표현의 구성과 분포 모델링을 분리하는 동시에, 블록 단위 생성과 병렬 디노이징을 통해 확산 모델의 효율성을 극대화하는 것이다.

기술적 접근법

**Query-based Encoder-Decoder**: 텍스트를 prefix-aligned, high-capacity 잠재 시퀀스로 인코딩.
**Block-causal Diffusion Transformer**: 블록 단위로 왼쪽에서 오른쪽으로 생성하며, 각 블록 내 위치는 병렬 디노이징.
**Flow matching**: 확산 모델이 잠재 분포를 학습하는 방식.
**Self-trajectory consistency**: 학습 시 독립적으로 샘플링된 노이즈와 추론 시 반복 디노이징 간 일관성 유지.
**Noise-level calibration**: 잠재 공간의 너비에 따라 노이즈 분포를 조정.
**하드웨어**: 모든 실험은 Ascend NPU에서 수행됨.

주요 결과

**OpenWebText** 데이터셋에서 자유 생성 성능이 기존 연속 및 확산 기반 언어 모델 중 가장 우수.
**XSum** 요약 작업에서 1B 파라미터, 1500 EFLOPs 계산량으로 기존 공개된 latent-diffusion 모델을 상회.
**모델 규모 대비 성능**: 1B 파라미터 모델이 더 큰 모델보다 높은 성능을 보임 (명시되지 않음).

의의 및 한계

AURORA-LM은 텍스트 생성에서 연속 잠재 공간의 활용을 확장하고, 확산 모델이 복잡한 잠재 분포를 학습할 수 있음을 보여준다. 특히, 블록 인과 트랜스포머와 self-trajectory consistency는 확산 모델의 효율성과 일관성을 동시에 향상시키는 중요한 기술적 기여이다. 그러나, 모델의 복잡성은 추론 시간 및 계산 자원을 증가시키며, 실제 대규모 배포 시 성능-자원 균형이 중요한 한계로 작용할 수 있다.

실용적 활용

AURORA-LM은 대규모 텍스트 생성, 요약, 번역 등 NLP 분야에서 높은 정확도와 일관성을 요구하는 상황에 적용 가능하다. 특히, 블록 단위 생성과 병렬 디노이징 구조는 실시간 또는 대량 텍스트 처리 시스템에 유리하다.