기술 요약: AURORA-LM

문제 정의

언어 생성은 현대적 생성 모델링에서 여전히 이례적인 위치에 머물러 있습니다. 이미지, 비디오, 오디오는 확산 프로세스(diffusion processes)를 사용하여 연속적인 잠재 공간(continuous latent spaces)에서 모델링되는 경우가 점점 늘어나고 있는 반면, 텍스트 생성은 여전히 주로 이산적 토큰(discrete tokens)에 의존합니다. 연속적 확산을 언어에 적용하려는 기존의 시도들은 근본적인 설계상의 타협에 직면해 있습니다:

1. 상속된 표현(Inherited Representations): Diffusion-LM이나 TEncDM과 같은 방식은 토큰 임베딩이나 사전 학습된 인코더로부터 임베딩 공간을 상속받습니다. 이러한 공간은 정확한 텍스트 재구성(reconstruction)과 효과적인 연속적 생성이라는 두 가지 요구사항을 동시에 충족하도록 명시적으로 설계되지 않았습니다.
2. 압축 트레이드오프(Compression Trade-offs): 오토인코더 기반 접근 방식(예: LD4LG, COSMOS)은 텍스트를 연속적 잠재 공간으로 매핑하지만, 확산 분포를 학습하기 쉽게 만들기 위해 종종 이러한 표현을 압축합니다. 이러한 압축은 토큰 수준의 충실도(fidelity)를 희생시키며, 생성의 용이성을 높이면 디코딩 정확도가 저하되고, 그 반대의 경우도 마찬가지인 병목 현상을 만들어냅니다.

핵심 과제는 정확한 디코딩을 지원할 수 있는 연속적 텍스트 표현을 구축하는 동시에, 표현 능력을 희생하지 않으면서 그 분포를 학습할 수 있는 생성 모델을 설계하는 것입니다.

방법론: AURORA-LM 프레임워크

AURORA-LM(Autoencoding Unified Representation for Continuous-Latent Diffusion Language Modeling)은 표현 구축과 분포 모델링을 두 개의 별도 단계로 분리함으로써 이러한 과제를 해결합니다.

1. 연속적 텍스트 잠재 공간 구축 (오토인코더)

기존의 표현을 상속받는 대신, AURORA-LM은 **쿼리 기반 인코더-디코더(Query-based Encoder-Decoder)**를 사용하여 고용량의 디코딩 가능한 텍xt 잠재 공간을 명시적으로 구축합니다:

• 쿼리 기반 인코더: 이산적 텍스트를 고용량의 접두사 정렬(prefix-aligned) 잠재 시퀀스로 구성합니다. 이는 $N$개의 잠재 쿼리를 사용하여 점진적으로 길어지는 토큰 접두사를 집계합니다. 쿼리 인덱스 $i$가 증가함에 따라, 인코더는 점진적으로 더 긴 토큰 접두사($w_{1:\lceil iL/N \rceil}$)에 주의(attend)하며, 이를 통해 잠재 위치가 텍스트 접두사와 정렬되는 인과적 순서(causal order)를 유도합니다.
• 쿼리 기반 디코더: 잠재 접두사를 다시 토큰 로짓(logits)으로 매핑합니다. 결정적으로, 디코더는 오직 대응하는 잠재 접두사에만 주의하도록 제한되어, 잠재 시퀀스와 토큰 시퀀스 사이에 엄격한 인과적 대응 관계를 보장합니다.
• 고용량 설계: 생성의 용이성을 위해 잠재 공간을 압축하는 이전 연구들과 달리, AURORA-LM은 높은 채널 폭($D$)과 최소한의 시퀀스 압축($c \approx 1$)을 유지합니다. 이는 잠재 공간을 단순화된 사전(prior)이 아닌, 텍스트에 대한 충실한 인터페이스로 취급하여 정확한 토큰 복구에 필요한 정보를 보존합니다.
• 학습: 오토인코더는 토큰 수준의 교차 엔트로피(cross-entropy)를 통해 학습된 후 **동결(frozen)**됩니다. 이를 통해 목표 잠재 분포를 고정함으로써, 생성 모델이 표현 학습과 독립적으로 연구될 수 있도록 합니다.

2. 블록-인과적 잠재 확산 (생성 모델)

**블록-인과적 확산 트랜스포머(Block-Causal Diffusion Transformer)**는 **플로우 매칭(Flow Matching)**을 사용하여 동결된 잠재 시퀀스의 분포를 학습합니다:

• 블록-인과적 인수분해(Block-Causal Factorization): 인과적 생성과 병렬 디노이징 사이의 균형을 맞추기 위해 잠재 시퀀스를 블록으로 분할합니다. 모델은 왼쪽에서 오른쪽으로 블록을 생성(이전 블록에 대해 인과적 조건부 설정)하지만, 하나의 블록 내부의 모든 위치는 병렬로 디노이징합니다.
• 전체 폭 잠재 예측(Full-Width Latent Prediction): 모델은 디코더가 요구하는 전체 폭의 깨끗한 잠재값($x_0$)을 예측하며, 압축된 버전을 예측하지 않습니다.
• 노이즈 입력 병목(Noisy-Input Bottleneck): 노이즈가 섞인 입력으로부터 고차원 분포를 학습하는 어려움을 해결하기 위해, 모델은 노이즈가 섞인 입력 경로에 대해서만 저차원 투영(low-rank projection)($W_{down} \in \mathbb{R}^{D_b \times D}$)을 적용하고, 출력 헤드는 전체 폭의 깨끗한 타겟을 예측합니다. 이를 통해 모델은 디코더가 마주하는 용량은 줄이지 않으면서도 오염된 입력을 효율적으로 처리할 수 있습니다.
• 노이즈 캘리브레이션(Noise Calibration): 노이즈 레벨 분포는 잠재 폭에 맞춰 조정됩니다. 실험 결과, 더 넓은 잠재 표현은 고노이즈 입력에 더 큰 비중을 두는 노이즈 스케줄로부터 이득을 얻는 것으로 나타났습니다.
• 자기 궤적 일관성(Self-Trajectory Consistency): 독립적으로 샘플링된 노이즈 상태에서의 학습과 반복적 디노이징을 통한 추론 사이의 간극을 메우기 위해, 모델은 일관성 손실(consistency loss)을 도입합니다. 이는 디노이징 궤적을 따라 이웃한 상태에서의 깨끗한 잠재 예측을 정렬하여, 생성 과정을 안정화하고 적은 샘플링 단계로도 고품질의 결과를 낼 수 있게 합니다.

3. 추론

추론 과정에서 잠재 블록은 왼쪽에서 오른쪽으로 순차적으로 생성됩니다. 생성된 표준화된 잠재 값들은 원래의 인코더 공간으로 역전사(invert)되어, 동결된 디코더를 통해 최종 토큰 시퀀스로 복구됩니다. 이 프레임워크는 무조건적 생성(unconditional generation)과 프롬프트 조건부 생성(classifier-free guidance 또는 self-conditioning guidance 사용)을 모두 지원합니다.

주요 기여

1. 통합된 프레임워크: AURORA-LM은 텍스트 표현 학습과 생성적 분포 모델링을 분리하여, 토큰 임베딩으로부터 상속된 것이 아니라 직접 구축된 연속적 표현을 생성 모델이 학습하도록 합니다.
2. 고용량 인터페이스: 정확한 토큰 디코딩을 위해 잠재 표현의 전체 채널 용량을 유지하며, 표현을 단순화하는 대신 확산 사전(diffusion prior)이 그 복잡한 분포를 학습하도록 설계했습니다.
3. 블록-인과적 플로우 매칭: 왼쪽에서 오른쪽으로 블록을 생성하면서도 블록 내 위치는 병렬로 디노이징하는 블록-인과적 디노이저를 플로우 매칭과 함께 도입했습니다.
4. 아키텍처 혁신: 저차원 노이즈 입력 경로, 잠재 폭에 맞춘 노이즈 레벨 캘리브레이션, 자기 궤적 일관성의 결 조합을 통해 모델이 전체 폭의 잠재 분포를 효과적으로 학습할 수 있도록 했습니다.

실험 결과

실험은 Ascend NPU에서 수행되었습니다.

• OpenWebText (무조건적 생성): AURORA-LM-S (130M 파라미터)는 평가된 자기회귀(autoregressive), 이산 확산(discrete-diffusion), 연속 임베딩(continuous-embedding) 베이스라인들 중 가장 낮은 생성 퍼플렉서티(Gen-PPL: 23.56)와 가장 높은 MAUVE(0.890)를 달성했습니다.
• XSum (조건부 요약): AURORA-LM-S는 ROUGE-1, ROUGE-2, ROUGE-L 지표에서 모든 베이스라인(ELF-B, AR, MDLM 포함)을 능가했습니다.
• 스케일링(Scaling): AURORA-LM-L (1.01B 파라미터, ~1,500 EFLOPs 연산량)은 일치된 평가 프로토콜 하에서 일관된 성능 향상을 보여주었으며, 더 큰 규모로 공개된 잠재 확산 모델(Cola-DLM, ~1.8B 파라미터)을 넘어섰습니다.
• 절제 연구(Ablation Studies):
  • 잠재 폭(Latent Width): 더 넓은 잠재 폭($D=1024$)이 오염에 대한 더 나은 강건성과 더 높은 생성 성능을 제공했습니다.
  • 입력 병목(Input Bottleneck): 노이즈 입력에 대해 중간 정도의 병목 폭($D_b=128$)을 사용하는 것이 최적이었으며, 이를 통해 오염된 입력을 전체 차원으로 처리하지 않고도 전체 폭의 예측이 가능했습니다.
  • 궤적 일관성(Trajectory Consistency): 자기 궤적 일관성은 적은 단계의 샘플링 품질을 유의미하게 개선했습니다.

의의 및 주장

본 논문은 연속적 언어 생성이 확산 기반 생성 모델링과 이산 토큰 디코딩 사이를 효과적으로 연결할 수 있다고 주장합니다. 저자들은 생성 모델을 수용하기 위해 표현을 단순화하는 기존의 설계적 타협이 불필요하다고 주장합니다. 대신, 고용량의 디코딩 가능한 텍㇑스트 잠재 공간을 보존하고, 특정 아키텍처적 선택(블록-인과적 인수분해 및 궤적 일관성 등)을 통해 확산 모델이 이를 직접 학습하도록 설계함으로써, 무조건적 및 조건부 텍스트 생성 작업 모두에서 최첨단(SOTA) 성능을 달이트는 것이 가능하다는 것을 보여줍니다. 결과적으로 연속적 공간이 언어의 의미적 충실도와 구문적 규칙성을 유지하면서 확산 모델링의 이점을 활용할 수 있는 일반적인 인터페이스 역할을 할 수 있음을 시사합니다.