Gorio Tech Blog search

AURORA-LM: Autoencoding Unified Representation for Continuous-Latent Diffusion Language Modeling 요약 설명

03 Aug 2026 | Paper Review Continuous Diffusion Language Models Flow Matching

목차

이번 글에서는 AURORA-LM: Autoencoding Unified Representation for Continuous-Latent Diffusion Language Modeling 논문의 핵심 포인트만 간단히 정리한다.

  • 2026년 8월 3일(Arxiv), arXiv
  • Liang, Jiajun, Liao, Yucheng, Cao, Yukang, Wei, Jiazhe, Li, Ken, Tan, Wende, Zhang, Jiankun, Cui, ZY, Yang, Jingkang, Guo, Liucheng, et al.
  • PRLab, Nanjing University, S-Lab, Nanyang Technological University, Imperial College London
  • 논문 링크
  • Github
  • Project Page

요약

  • AURORA-LM은 이산 토큰이 아니라 연속 잠재열을 생성한 뒤 토큰으로 복호화하는 확산 언어 모델이다. 기존 연속 언어 모델은 생성하기 쉬운 잠재분포를 위해 오토인코더 잠재를 압축하는 경우가 많았고, 저자들은 그 압축이 정확한 토큰 복원에 필요한 정보를 줄일 수 있다고 본다. 이에 토큰 복원용 고용량·접두사 정렬 잠재표현을 먼저 학습하고 동결한 뒤, 별도의 block-causal diffusion Transformer가 이 전체 폭 잠재분포를 flow matching으로 학습하도록 두 단계를 분리한다. 생성기는 블록 사이에서는 좌에서 우로 조건부 생성하고 각 블록 내부 위치는 병렬로 denoise한다. 저자들은 noisy 입력 경로만 저랭크로 제한하고 clean-latent 목표는 전체 폭으로 유지하며, 잠재 폭에 맞춘 노이즈 배분 및 self-trajectory consistency를 결합한다. OpenWebText 1,024토큰 무조건 생성의 Table 2에서 AURORA-LM-S는 평가된 7개 시스템 중 Gen-PPL 23.56으로 최저, MAUVE 0.890으로 최고를 기록했으며, XSum Table 3에서는 ROUGE-1/2/L이 36.6/13.4/28.9였다. 1B denoiser인 AURORA-LM-L은 Table 4의 9개 과제 macro average에서 Cola-DLM의 25.1%보다 7.5퍼센트포인트 높은 32.6%를 기록했지만, 두 모델의 학습 설정이 일치한 재학습 비교는 아니다.
    • Figure 1의 방사형 그래프는 지표별 정규화 값이며 Gen-PPL은 역수로 표시하므로, 원시 수치 비교에는 Table 2–4를 사용해야 한다.
    • 논문의 parameter 수는 기본적으로 동결 오토인코더를 제외한 block-causal denoiser 기준이다. AURORA-LM-S는 denoiser 130M, 전체 시스템은 267.8M이며, AURORA-LM-L은 denoiser 1.01B, 전체 시스템은 약 1.34B parameters다.
생성 설정과 모델 규모별 시스템 수준 평가
생성 설정과 모델 규모별 시스템 수준 평가
추론 시 생성 표현에 따른 언어 모델링 패러다임
추론 시 생성 표현에 따른 언어 모델링 패러다임

Continuous Text Latent Construction

  • 길이 L의 토큰열 w_1:L에 대해 Query-based Encoder는 N=round(cL)개의 위치와 채널 폭 D를 가진 연속 잠재열 z_enc을 만든다. 보존 비율 c=1은 길이를 압축하지 않는 경우이고, c가 작아지면 잠재 위치 수가 줄어든다. 인코더의 i번째 잠재 질의는 자신 및 앞선 잠재 상태와 토큰 접두사 w_1:ceil(iL/N)만 보므로, 뒤쪽 잠재 위치일수록 더 긴 텍스트 접두사를 읽는다. Query-based Decoder는 원 토큰을 입력으로 받지 않고, 출력 위치 j에서 앞선 decoder-query 상태와 z_enc,1:floor((j-1)N/L)+1만 참조해 어휘 로짓을 복원한다. 이처럼 양쪽의 가시성 범위를 맞춰 압축된 잠재열에서도 토큰 접두사와 잠재 접두사의 대응을 유지한다.
    • 인코더 입력과 디코더 출력 사영은 학습 가능한 토큰 임베딩 행렬 E를 공유하고, 디코더 최종 상태는 E의 전치로 어휘 로짓에 사영한다. 오토인코더는 원 토큰에 대한 토큰 단위 교차엔트로피로 공동 학습한다.
    • 학습 때 token-embedding dropout은 각 유효 토큰 위치의 입력 임베딩 전체를 독립적으로 제거하고, latent dropout은 잠재 좌표를 독립적으로 제거한다. 부록의 구현값은 각각 p_x=0.3, p_z=0.6이며 역 dropout 재스케일링은 사용하지 않는다. 학습 후 두 dropout을 끄고 오토인코더를 동결해 denoiser 학습용 결정론적 잠재열을 만든다.
Query-based 오토인코더와 블록 인과 잠재 확산의 전체 파이프라인
Query-based 오토인코더와 블록 인과 잠재 확산의 전체 파이프라인

Block-Causal Modeling of Continuous Text Latents

  • 동결된 인코더 출력은 채널별 평균과 표준편차로 affine 표준화한 뒤 denoiser의 학습 대상으로 사용한다. 표준화 잠재열의 N개 위치를 크기 Q의 연속 블록 B=ceil(N/Q)개로 나누고, 전체 분포를 이전 clean 블록 접두사에 조건부인 블록 분포들의 곱으로 분해한다. 대상 블록은 clean 블록과 가우시안 노이즈를 t∈[0,1]에서 선형 보간한 상태로 만들며, 모델은 현재 noisy 블록, 노이즈 수준, 이전 clean 블록을 받아 clean 블록을 예측한다. 이는 모든 잠재 위치를 한 번에 생성하는 비인과 방식과 위치 하나씩 생성하는 완전 순차 방식 사이에서, 접두사 조건부 구조와 블록 내부 병렬성을 함께 사용하는 선택이다.
    • 훈련에서는 생성된 접두사를 순차 롤아웃하지 않는다. 동결 인코더가 만든 clean 블록 접두사를 teacher-forcing 문맥으로 사용하고, 각 대상 블록은 독립적으로 오염시킨다.
    • two-stream attention mask는 블록 b의 예측이 자기 블록의 noisy 상태와 이전 clean 블록에는 접근하게 하되, 미래 clean 블록 및 다른 블록의 noisy 상태는 보지 못하게 한다. 따라서 한 번의 forward pass로 모든 블록 조건부를 계산하면서 추론 시 블록 인과성을 보존한다.

Learning the Full-Width Latent Distribution

  • 동결 decoder가 요구하는 clean 잠재는 전체 폭 D로 유지하되, 저자들은 Transformer가 noisy 입력을 읽는 경로만 D→D_b→H의 두 선형층으로 저랭크 사영한다. D_b<min(D,H)이면 noisy 입력 사영의 rank는 최대 D_b가 되지만, 출력 head는 D차원의 clean 잠재를 그대로 예측한다. 주목적함수는 유효 잠재 위치에서 예측 clean 잠재와 clean 목표 사이의 평균제곱오차인 clean-endpoint flow-matching 손실이다. 최종 시스템은 tan-d shift 7로 노이즈 수준을 표본화하고, x0 예측 및 x0 공간 손실을 사용한다.
    • self-conditioning은 이전 단계의 clean-latent 초안 예측을 추가 입력으로 제공한다. 훈련에서는 확률 p_sc로 보조 forward pass의 출력을 gradient에서 분리해 주 예측에 넣고, 나머지 경우에는 초안을 생략한다. p_sc는 AURORA-LM-S에서 0.6, L에서 0.5다.
    • self-trajectory consistency는 독립 noisy 상태로 훈련하는 흐름 맞춤과 반복 solver 궤적으로 추론하는 차이를 줄인다. 고정 16-step 스케줄에서 현재 예측으로 한 Euler update를 해 인접한 저노이즈 상태를 만들고, 온라인 모델의 현재 clean 예측과 EMA 모델의 다음 상태 예측을 맞춘다. 상태 update와 EMA target에는 gradient를 흘리지 않으며, 최종 손실은 L_FM+λ_ct L_ct이고 λ_ct=1.0이다. consistency 가중치는 S 모델 첫 10,000회, L 모델 첫 5,000회 최적화 단계 동안 선형 증가한다.
인접 denoising 상태의 clean-latent 예측을 정렬하는 self-trajectory consistency
인접 denoising 상태의 clean-latent 예측을 정렬하는 self-trajectory consistency

Latent Generation and Text Decoding

  • 무조건 생성은 빈 접두사에서 시작하며, 조건부 생성은 동결 인코더가 프롬프트를 잠재 접두사로 변환한 뒤 학습 때와 같은 채널별 표준화를 적용한다. 새 블록은 가우시안 노이즈에서 시작해 t=1에서 0까지 반복 denoise하고, 완성한 블록을 다음 블록의 인과적 문맥에 붙인다. 생성된 표준화 잠재는 affine 표준화를 역변환한 뒤 고정 Query-based Decoder에 넣어 토큰 로짓과 출력 토큰으로 변환한다. 완료 블록의 KV cache를 재사용해 좌에서 우로 진행하는 생성을 효율화한다.
    • 무조건 생성의 SC-CFG는 self-conditioning을 준 예측과 생략한 예측을 보간하고, 조건부 생성의 CFG는 프롬프트 접두사가 있는 예측과 없는 예측을 보간한다. 두 방식은 학습 목표를 변경하지 않는 추론 단계 제어이며, scale이 1이면 각각 self-conditioned 또는 조건부 예측을 그대로 사용한다.
    • 프롬프트가 블록 중간에서 끝날 때에는 boundary 블록에서 이미 알려진 위치의 encoder 잠재를 denoising 내내 고정하는 clean repainting을 사용한다. XSum 및 공개 벤치마크의 CFG는 고노이즈 구간 σ∈[0.9,1.0]에만 적용한다.

Latent Capacity

  • OWT128 통제 실험은 채널 폭 D와 길이 보존 비율 c를 분리해 조사했다. D=128, 256, 512, 1024의 clean 잠재는 모두 원 토큰을 거의 완벽하게 복원했지만, Figure 5(a)는 D가 넓을수록 더 강한 잠재 오염에서도 토큰 복원 정확도를 더 오래 유지함을 보인다. 각 폭에서 tan-d shift d∈{1,3,5,7}를 탐색해 얻은 최고 MAUVE는 D=128에서 0.594±0.035, 256에서 0.640±0.033, 512에서 0.693±0.011, 1024에서 0.839±0.024였다. 이는 각 D마다 노이즈 스케줄을 탐색한 뒤의 비교이므로, 잠재 폭만의 인과적 효과로 해석할 수는 없다.
    • D=1024 및 D_b=128에서 잠재 위치를 100%, 90%, 80%, 70% 보존했을 때 MAUVE는 각각 0.815±0.012, 0.807±0.022, 0.726±0.020, 0.671±0.063이었다. 100% 대비 90% 보존의 절대 차이는 -0.008이고, 80%와 70%에서는 각각 -0.089, -0.144다. 저자들은 c=0.9를 효율 우선의 대안으로 제시하되, 이후 절제의 기준 구성에는 압축 열화를 피하려고 c=1을 사용했다.
    • Figure 5(c)의 0.82, 0.81, 0.73, 0.67은 반올림 표기이며, 정확한 평균과 표준편차는 Appendix Table 12에 제시되어 있다.
잠재 채널 폭과 시퀀스 압축률에 따른 잠재 용량 평가
잠재 채널 폭과 시퀀스 압축률에 따른 잠재 용량 평가

Modeling the Full-Width Latent Distribution

  • D=1024 full-width clean 목표에서 noisy-input 병목 폭 D_b를 바꾼 실험은 16·32·64-step ODE의 평균 MAUVE를 D_b=32, 128, 256, 512, 1024에서 각각 0.786, 0.808, 0.760, 0.748, 0.790으로 보고한다. Figure 6(a)의 이 평균 기준으로 D_b=128이 이 탐색 범위에서 가장 높다. 모든 조건의 clean 목표 폭은 D=1024로 동일하므로, 결과는 decoder에 전달되는 목표 폭을 줄이지 않고 noisy 입력 경로 폭만 조절한 비교다.
    • 노이즈 배분은 tan-d와 두 logit-normal 계열을 고노이즈 질량 m_0.7=Pr(σ>0.7)으로 맞춰 비교했다. Table 14에서 tan-d d=1은 m_0.7=0.258, MAUVE 0.229±0.023이고 d=7은 0.795, 0.839±0.024이다. 세 schedule 계열의 탐색점에서는 고노이즈 질량이 클수록 MAUVE가 증가하는 양상이 관찰됐지만, 이는 D=1024 OWT128 통제 설정의 결과다.
    • Table 1에서 x0 예측과 x0 공간 손실은 D_b=128에서 0.815±0.012, D_b=1024에서 0.807±0.033으로 각 병목 폭의 네 조합 중 가장 높았다. x0 예측을 v 공간에서 평가하면 clean-latent 오차에 1/σ²가 곱해져 저노이즈 상태에 감독이 집중되며, 저자들은 이것이 고노이즈 배분 결과와 맞지 않는 이유라고 해석한다.
noisy-input 병목 폭과 고노이즈 배분의 full-width 잠재 모델링 평가
noisy-input 병목 폭과 고노이즈 배분의 full-width 잠재 모델링 평가
예측 목표와 손실 공간 조합별 MAUVE 비교
예측 목표와 손실 공간 조합별 MAUVE 비교

Efficient Blockwise Generation

  • self-trajectory consistency는 OWT128에서 시험한 모든 step budget의 MAUVE를 높였다. Appendix Table 16의 consistency off→on 수치는 8 step에서 0.116±0.006→0.672±0.038, 16 step에서 0.310±0.039→0.786±0.034, 32 step에서 0.589±0.062→0.815±0.012, 64 step에서 0.660±0.023→0.832±0.014이다. 절대 차이는 각각 0.556, 0.476, 0.226, 0.172로 적은 step에서 더 크며, 128·256 step에서도 적용 모델이 높지만 격차는 좁아진다.
    • 별도 훈련한 Q=4, 8, 16, 32, 64 모델을 guidance 없는 32-step ODE로 평가한 Table 16의 MAUVE는 각각 0.909±0.012, 0.857±0.011, 0.815±0.012, 0.681±0.023, 0.631±0.027이다. 작은 블록은 완성된 선행 문맥을 더 자주 제공하지만 순차 생성 단계가 많고, 큰 블록은 내부 병렬성을 늘리지만 조건부 문맥은 줄인다.
    • 수치 충돌이 있다. 4.1.3 본문은 Q=16의 MAUVE를 0.816으로 서술하지만, Figure 7은 반올림해 0.82로 표시하고 정확 수치를 제시한 Appendix Table 16은 0.815±0.012이다. 여기서는 반복 평가의 평균과 표준편차를 직접 제시한 Table 16을 채택한다. Q=4 대비 Q=16은 순차 블록 단계를 32에서 8로 줄이는 대신 MAUVE가 0.094 낮아진다.
self-trajectory consistency와 블록 크기에 따른 블록별 생성 평가
self-trajectory consistency와 블록 크기에 따른 블록별 생성 평가

Unconditional Generation

  • 시스템 수준 무조건 생성은 OpenWebText 1,024토큰 시퀀스로 학습한 denoiser 130M의 AURORA-LM-S를 평가한다. 각 시스템은 입력 프롬프트 없이 1,000개 표본을 생성했고, GPT-2 Large를 고정 외부 평가기로 쓴 Gen-PPL, token-unigram entropy, held-out OpenWebText와의 MAUVE를 측정했다. Table 2에서 AURORA-LM-S는 64-step SDE-DPM++ 및 SC-CFG 4 조건으로 Gen-PPL 23.56, entropy 5.241, MAUVE 0.890을 기록했다. 같은 표의 7개 시스템을 모두 대조하면 Gen-PPL은 최저이고 MAUVE는 최고지만, entropy 최고값은 MDLM의 5.664다.
    • 비교 대상은 AR 85M, Duo·Duo-distilled·SEDD·MDLM 각 92M, ELF-B 105M이다. Table 2의 모든 행은 저자들의 in-house 재현이지만, 각 시스템의 sampler와 guidance 설정은 다르므로 이 표를 동일 추론 비용 비교로 해석할 수는 없다.
    • 선택된 구성은 다섯 evaluation seed마다 1,000개 표본을 만든 Table 7 sweep에서 평균 Gen-PPL이 가장 낮은 64-step SDE-DPM++, SC-CFG 4다. 이 구성의 MAUVE는 0.8896±0.0222였고, sweep의 최고 MAUVE는 32-step SDE-DPM++, SC-CFG 3의 0.9516±0.0055였다. 선택 기준은 Gen-PPL이므로, 선택 구성이 해당 sweep에서 MAUVE 최고라는 뜻은 아니다.
OpenWebText 1,024토큰 무조건 생성 비교
OpenWebText 1,024토큰 무조건 생성 비교

Conditional Generation

  • XSum 조건부 요약에서 AURORA-LM-S는 원문 기사를 잠재 접두사로 인코딩하고 32-step ODE와 CFG scale 2.0으로 요약을 생성했다. Table 3의 ROUGE-1/2/L은 36.6/13.4/28.9이며, 표의 ELF-B 36.0/12.2/27.8, AR 30.5/10.2/24.4, MDLM 33.4/11.6/25.8, Duo 31.4/10.1/25.0, E2D2 28.4/8.3/22.0, SeqDiffuSeq 19.3/1.7/14.1보다 세 지표 모두 높다. 따라서 Table 3에 표시된 비교 시스템 범위에서는 세 ROUGE 지표 모두 가장 높다.
    • Table 3의 기준선 수치에는 ELF 논문에서 수집한 값이라는 †가 붙는다. Appendix A.2.7은 XSum 기준선 지표를 ELF에서 직접 가져왔고 기준선을 종단간 재학습하지 않았다고 명시한다. 따라서 표는 저자들이 제시한 매칭 프로토콜의 수치를 채택한 비교이지만, 모든 기준선을 본 연구가 동일 조건에서 재평가한 결과는 아니다.
    • Table 8의 guidance 탐색에서 w=2.0은 평균 ROUGE-1 36.56, ROUGE-L 28.86으로 가장 높고, ROUGE-2는 w=1.5의 13.47이 w=2.0의 13.43보다 0.04 높다. 저자들은 R-1과 R-L을 기준으로 w=2.0을 선택했다. w=2.0에서 32 step의 R-2 13.43은 64 step의 13.43과 같으며, R-1과 R-L은 32 step이 더 높다.
XSum 조건부 요약의 ROUGE 비교
XSum 조건부 요약의 ROUGE 비교

Scaling Evaluation

  • 확장 평가의 AURORA-LM-L은 denoiser 1.01B parameters로, 저자 본문은 총 학습 연산량을 약 1,500 EFLOPs라고 서술한다. Table 6의 상세 추정치는 오토인코더 약 157 EFLOPs와 denoiser 약 1,406 EFLOPs로 합계 약 1,563 EFLOPs이며, 이 수치와 본문의 약 1,500 EFLOPs가 어떻게 대응하는지는 명시되지 않았다. 학습 말뭉치는 294.7B토큰으로 구성한 영어 자료 혼합물이며, 이는 최적화 동안의 총 token exposure가 아니라 구성한 코퍼스의 토큰 수다. Qwen3 tokenizer로 토큰화해 512토큰 시퀀스로 패킹했으며, instruction fine-tuning은 하지 않았다.
    • Table 4에서 AURORA-LM-L의 9개 과제 macro average는 32.6%, 공개 checkpoint Cola-DLM은 25.1%로, 25.1%를 기준으로 한 절대 차이는 7.5퍼센트포인트이고 상대 변화율은 (32.6-25.1)/25.1×100=29.9%다. AURORA-LM-L은 표의 아홉 개 개별 과제에서도 높다. 예를 들어 StoryCloze는 50.3% 대 45.3%, SIQA는 54.8% 대 33.8%, SQuAD EM은 38.2% 대 25.7%다.
    • Cola-DLM은 약 1.8B-parameter diffusion Transformer의 공개 checkpoint이며, 저자들은 이를 공식 추론 구현과 기본값인 16-step Euler·CFG 7.0으로 평가했다. AURORA-LM-L은 16-step ODE·CFG 3.0을 썼다. 두 모델은 같은 평가 사례, prompt, truncation, 답안 매칭 절차를 공유하지만, Cola-DLM을 재학습하지 않았으므로 학습 조건은 일치하지 않는다. 과제별로 최대 1,000개 사례를 남기고 두 모델의 문맥 예산을 모두 만족하는 사례만 사용했으며, 두-shot prompt와 최대 32개 새 토큰을 썼다.
1B 모델과 공개 latent-diffusion 모델의 9개 벤치마크 비교
1B 모델과 공개 latent-diffusion 모델의 9개 벤치마크 비교

부록

  • 부록은 오토인코더 학습 때의 token-embedding 및 latent dropout, clean-latent 예측을 velocity 공간 손실로 환산할 때의 가중 효과, 노이즈 schedule의 고노이즈 질량 계산, self-trajectory consistency의 warmup을 설명한다. 또한 데이터 전처리, 모델 구조, 학습과 샘플링 설정, 기준선 재현 및 few-shot 평가 프로토콜을 제공한다. 마지막으로 본문의 OWT128 통제 실험에 대한 전체 수치, OpenWebText·XSum 생성 예시, 아홉 벤치마크의 two-shot demonstration을 수록한다.

짧은 생각

이 논문의 핵심 가설은 decoder의 토큰 복원 능력에 필요한 clean 잠재의 폭을 줄이는 대신, denoiser가 noisy 입력을 읽는 통로만 좁히면 고용량 표현과 잠재분포 모델링을 분리할 수 있다는 것이다. 이를 검증한 Table 13은 동일한 D=1024 clean 목표에서 D_b=128의 평균 MAUVE 0.808이 D_b=1024의 0.790보다 0.018 높음을 보이므로, 적어도 이 OWT128·16–64-step ODE 평균 조건에서는 비대칭 설계를 뒷받침한다. 다만 Table 2의 선택된 무조건 생성 구성은 Gen-PPL 최저를 기준으로 고른 반면, 같은 Table 7 sweep에는 MAUVE가 더 높은 구성이 존재한다. 더구나 부록의 OpenWebText 예시는 문법적으로 이어지는 구간이 있어도 허구의 고유명사, 불안정한 사실 진술, 반복을 포함한다. 따라서 Gen-PPL 중심의 구성 선택이 인간이 판단하는 사실성·일관성·선호와 얼마나 일치하는지 확인하려면, 선택용 검증 세트와 최종 보고 세트를 분리하고 인간 평가 또는 사실성 평가를 추가할 필요가 있다.

Related Posts