AURORA-LM: 연속 잠재 확산 언어모델

왜 중요한가

텍스트 생성은 여전히 이산 토큰 위주인데 연속 잠재로 직접 생성하면 이미지·오디오 분야에서 얻은 잠재 공간 이점들을 언어로 옮길 수 있다는 점에서 의미가 큽니다. AURORA-LM은 디코더가 복원할 수 있는 고용량 잠재를 보존하면서 확산 모델이 그 분포를 직접 학습하도록 구조를 분리해 실용적 제약을 완화합니다. 이런 설계는 토큰 충실도를 유지하면서 확산 기반 생성의 병렬성·스케일링 이점을 활용할 수 있어 자유 생성과 요약 같은 downstream 작업에서 성능 향상을 기대하게 합니다.

핵심 기여

AURORA-LM 아키텍처 도입

논문은 디코더-정렬(decoder-facing) 용량을 유지하는 고용량 연속 잠재를 생성하는 전체 파이프라인을 구축합니다. Query-based Encoder-Decoder가 텍스트를 프리픽스 정렬된 잠재 시퀀스로 구성하고, 그 잠재 분포를 Block-causal Diffusion Transformer가 흐름 매칭(flow matching)으로 학습하도록 분리해 설계적 이점을 확보합니다. 이 분리는 잠재 표현의 복원 가능성을 유지하면서 확산 모델의 학습 목표를 명확히 분리해 처리 효율과 샘플 품질 사이 균형을 맞추는 근거가 됩니다.

블록-인과 확산 Transformer와 흐름 매칭

모델은 시퀀스를 블록 단위로 왼쪽에서 오른쪽으로 생성하되 각 블록 내부의 위치는 병렬로 디노이즈하도록 처리합니다. 이 처리 흐름을 흐름 매칭으로 학습해 노이즈 추가-제거 경로를 직접 모델링하며 샘플링 시 병렬성과 순차성 모두를 확보합니다. 블록 단위 생성은 긴 시퀀스에서 처리량을 높이면서도 순서성을 유지하려는 구현적 타협을 제공하므로 실험적 목적에서 중요합니다.

노이즈 경로 보정과 자기-경로 일관성

AURORA-LM은 노이즈 입력 경로에만 제약을 두고 깨끗한 잠재 예측(target)을 온전히 유지하도록 설계해 고차원 잠재를 보존합니다. 노이즈 레벨 분포를 잠재 폭(latent width)에 맞춰 보정하고 self-trajectory consistency를 도입해 학습 시 독립적으로 샘플된 노이즈와 추론 시 반복적 디노이징의 불일치를 줄입니다. 이러한 조치는 학습-추론 간 갭을 줄여 샘플 품질과 안정성을 개선하려는 메커니즘입니다.

실험적 성과와 확장성 검증

초록에 따르면 OpenWebText 자유 생성과 XSum 요약에서 평가된 연속·확산 계열 모델들 사이에서 최고 성능을 달성했습니다. 1B 파라미터 규모로 확장했고 전체 계산량 약 1500 EFLOPs에서 추가 성능 향상이 관찰되며, 동일 평가 프로토콜에서 공개된 더 큰 잠재-확산 언어모델을 능가한 결과가 보고되었습니다. 모든 실험은 Ascend NPU에서 수행되어 하드웨어 조건과 스케일 관점에서의 검증 근거를 제공합니다.

핵심 아이디어 이해하기

텍스트를 연속 잠재로 모델링하면 이미지·오디오에서 쓰이는 잠재 기반의 생성 장점을 언어에 적용할 수 있지만, 기존 접근은 디코더가 바로 쓸 수 있는 고용량 잠재를 희생하거나 임베딩을 재사용해 생성-복원 요구를 충족시키지 못했습니다. AURORA-LM은 잠재의 구성(encode)과 분포 학습(generate)을 명확히 분리해 Query-based Encoder-Decoder로 디코더-정렬 고용량 잠재를 확보하고, Block-causal Diffusion Transformer로 그 분포를 흐름 매칭으로 학습하는 방식으로 문제를 해결합니다. 이렇게 하면 노이즈 경로에만 제약을 두면서도 깨끗한 잠재를 온전히 예측 대상으로 유지할 수 있어 토큰 수준의 충실도를 유지하면서 확산 모델의 병렬성·스케일링 이점을 얻을 수 있습니다.

방법론

먼저 Query-based Encoder-Decoder가 입력 텍스트를 프리픽스 정렬된 고용량 연속 잠재 시퀀스로 변환하며, 이 잠재는 디코더가 복원 가능한 형태로 설계됩니다. 다음으로 Block-causal Diffusion Transformer가 흐름 매칭(flow matching) 목표로 해당 잠재 분포를 학습하는데, 생성은 블록 단위로 왼쪽에서 오른쪽으로 진행하면서 블록 내부 위치들은 병렬로 디노이즈합니다. 학습 과정에서는 노이즈-입력 경로의 제약만 적용하고 깨끗한 잠재 예측 목표는 유지하며, 노이즈 레벨 분포를 잠재 폭에 맞춰 보정하고 self-trajectory consistency로 학습 시 샘플링 분포와 추론 시 반복적 디노이징의 격차를 줄입니다.

주요 결과

초록 기준으로 OpenWebText 자유 생성과 XSum 요약에서 동일 계열의 연속 잠재 및 확산 기반 언어모델들보다 우수한 성능을 기록했습니다. 또한 모델을 1B 파라미터로 확장했을 때 전체 계산량 약 1500 EFLOPs 조건에서 추가 성능 향상이 보고되었고, 동일 평가 프로토콜 하에서 공개된 더 큰 잠재-확산 모델을 능가한 결과가 언급되어 스케일링 관점의 이득을 시사합니다. 모든 실험은 Ascend NPU에서 수행되어 하드웨어 제약과 실험 환경을 명확히 제시한 점이 근거로 제공됩니다.

기술 상세

핵심 기술은 Query-based Encoder-Decoder로 생성된 프리픽스 정렬 고용량 잠재와 Block-causal Diffusion Transformer의 흐름 매칭 학습입니다. 생성 절차는 블록 단위로 왼쪽에서 오른쪽으로 진행하면서 각 블록 내부 위치를 병렬로 디노이즈하는 구조로 병렬성과 순차성 사이의 균형을 맞추며, 학습 신호는 깨끗한 잠재 예측을 목표로 유지됩니다. 추가로 노이즈 레벨 분포를 잠재 폭에 맞춰 보정하고 self-trajectory consistency를 도입해 학습 시 독립 샘플링 노이즈와 추론 시 반복 디노이징 절차의 불일치를 완화하는 보정 메커니즘을 사용합니다.

실무 활용

연속 잠재를 직접 생성하도록 설계되어 자유로운 텍스트 생성과 요약과 같은 자연어 생성 과제에서 활용 가능성이 있습니다. 프리픽스 정렬 잠재를 유지하므로 디코더-접근 방식이 요구되는 파이프라인과 호환성이 높습니다. 공개된 GitHub 리포지토리가 있어 구현·재현 가능성 측면의 실무 적용 진입 장벽을 낮춥니다.

대형 언어모델을 이용한 자유 텍스트 생성(free generation) 실험과 비교 연구
추상적 요약이나 조건부 텍스트 생성에서 잠재 기반 샘플링 전략 검증
연속 잠재 표현을 필요로 하는 downstream 파이프라인의 모델 대체 실험

코드 공개 여부: 공개

코드 저장소 보기

키워드

연속 잠재 (continuous latent)확산 모델 (diffusion model)흐름 매칭 (flow matching)Query-based Encoder-DecoderBlock-causal Diffusion Transformerself-trajectory consistency

용어 해설

연속 잠재(Continuous Latent): — 텍스트를 이산 토큰 대신 실수 벡터 공간의 연속 표현으로 인코딩한 것이며, 디코더가 복원 가능한 고용량(latent width)을 유지하면 토큰 수준의 충실도를 보존하면서 확산 모델로 분포를 학습할 수 있다.
플로우 매칭(Flow Matching): — 확산 과정의 확률 흐름을 학습하는 방법으로, 노이즈가 추가된 입력에서 목표 확률 흐름(또는 벡터장)을 추정해 샘플링 시 연속 경로를 따라 노이즈를 제거하도록 학습을 유도하는 기법이다.
블록-인과성(Block-causal): — 시퀀스를 블록 단위로 왼쪽에서 오른쪽으로 생성하면서 각 블록 내부의 위치들은 병렬로 디노이즈하는 처리 방식으로, 전체 순서성은 보존하되 병렬성을 확보하는 구조적 제약을 의미한다.
프리픽스 정렬 잠재(Prefix-aligned Latent): — 입력 텍스트의 접두사(prefix)와 정렬된 위치를 유지하는 고용량 잠재 시퀀스로, 디코더가 해당 잠재를 바로 복원 대상으로 사용할 수 있게 구성되어 디코더-정렬성(decoder alignment)을 확보한다.
자기-경로 일관성(Self-trajectory Consistency): — 학습 시 독립적으로 샘플된 노이즈 경로와 추론 시 반복적인 디노이징 경로 사이의 차이를 줄이기 위해 훈련 신호를 보정하는 기법으로, 샘플링 절차와 훈련 분포의 불일치를 완화하도록 설계된다.