Uno로 확산 기반 병렬 토큰 생성 가속

왜 중요한가

LLM의 Autoregressive 구조는 이전 토큰을 입력으로 삼아 한 번에 하나씩 생성하므로 긴 출력에서 지연이 누적됩니다. Uno는 AR 가중치와 경량 확산 가중치를 분리하고, 별도 draft model 없이 여러 토큰을 동시에 샘플링해 기본 AR 모델의 품질을 유지하면서 추론 처리량을 높입니다.

핵심 기여

AR 분포를 유지하는 확산 증강 구조

AR 가중치는 표준 Next-Token Prediction 목표로 학습하고, 경량 확산 가중치는 여러 토큰을 동시에 생성하도록 학습하는 분리 구조를 사용합니다. 이를 통해 확산 샘플링의 병렬성을 기존 AR 모델의 분포와 결합합니다.

Diffusion Distillation 기반 추가 학습

기존 LLM 학습 파이프라인에 Diffusion Distillation 단계를 추가해 여러 토큰을 병렬 생성하는 확산 가중치를 학습합니다. 논문은 이 추가 과정이 기존 학습에 negligible overhead를 더한다고 설명합니다.

별도 draft model 없는 Ψ-Spec

Ψ-Spec은 고정된 context length에서 손실 없는 가속과 inference-time scaling을 지원하는 sampler 계열입니다. Speculative Decoding처럼 별도 draft model을 준비하지 않고 AR 모델 분포에서 여러 토큰을 처리합니다.

Uno 모델과 공개 체크포인트

Uno는 처음부터 학습하거나 기존 open-weight AR LLM을 확장해 구축할 수 있는 모델 계열입니다. 논문은 코드와 체크포인트를 공개 저장소 및 프로젝트 페이지에서 제공합니다.

핵심 아이디어 이해하기

기존 LLM의 Autoregressive 생성은 현재까지 나온 토큰을 입력으로 받아 다음 토큰 하나의 확률 분포를 계산하고, 선택된 토큰을 다시 문맥에 붙이는 과정을 반복합니다. 각 단계의 출력이 다음 단계의 입력이 되므로 여러 위치를 한 번에 계산하기 어렵고, 출력 길이가 늘어날수록 순차 지연이 쌓입니다.

Uno는 AR 모델이 정의한 토큰 분포 자체를 버리지 않고, 그 분포에서 여러 토큰을 동시에 뽑는 확산 경로를 추가합니다. AR weights는 표준 Next-Token Prediction으로 원래 분포를 학습하고, diffusion weights는 해당 분포를 사용해 여러 위치의 토큰을 병렬 생성하도록 별도로 학습합니다.

Diffusion Distillation은 이 경량 diffusion weights를 기존 LLM 학습 뒤에 붙이는 과정입니다. 입력 문맥에서 AR 모델이 제공하는 생성 분포를 기준으로 확산 가중치를 학습하고, 추론에서는 Ψ-Spec sampler가 여러 토큰을 함께 처리해 순차적인 토큰별 계산을 줄입니다.

이 구조는 d-LLM처럼 생성 방식만 바꾸면서 기본 AR 모델의 품질을 희생하는 접근이 아니라, AR 분포를 유지한 채 샘플링 단계를 병렬화하는 방식입니다. 따라서 Uno는 별도 draft model을 검증하는 Speculative Decoding과도 다르게 동작하며, 논문은 고정된 context length에서 손실 없는 가속을 목표로 합니다.

방법론

모델 파라미터를 두 집합으로 분리합니다. AR weights는 입력 문맥에서 다음 토큰 확률을 예측하는 표준 Next-Token Prediction 목표로 학습하고, diffusion weights는 같은 AR 모델 분포를 바탕으로 여러 토큰을 동시에 생성하도록 학습합니다. 입력 문맥은 AR 경로에서 토큰 분포로 변환되고, diffusion 경로는 이 분포를 이용해 여러 출력 위치의 토큰을 병렬 샘플링합니다.

Diffusion Distillation 단계에서는 기존 LLM 학습 파이프라인에 경량 diffusion weights 학습을 추가합니다. 이 과정은 AR weights 전체를 다시 학습하는 대신 확산 생성에 필요한 별도 가중치를 학습하므로, 논문은 기존 학습 대비 추가 비용이 negligible하다고 설명합니다. Uno는 이 절차로 처음부터 학습하거나 기존 open-weight AR LLM에 확산 가중치를 추가해 구성할 수 있습니다.

추론 단계에서는 Ψ-Spec sampler가 고정된 context length에서 여러 토큰을 함께 샘플링합니다. AR 모델이 정의한 분포를 입력으로 받아 병렬 토큰 후보를 생성하므로, Speculative Decoding처럼 별도 draft model이 후보를 만든 뒤 대상 모델이 검증하는 절차가 필요하지 않습니다. 논문은 이 방식으로 inference-time scaling과 손실 없는 가속을 지원한다고 설명합니다.

주요 결과

Uno는 평가한 모든 batch size에서 주요 Speculative Decoding 방법보다 높은 throughput을 기록했고, 기본 AR 모델 대비 최대 3배 speedup을 달성했습니다. 이 속도 향상은 장치가 지원하는 가장 큰 batch size에서도 관찰되었습니다.

8B Uno 모델은 agentic tool use, coding, long-context reasoning에 관한 모든 평가 benchmark에서 주요 open d-LLM인 26B DiffusionGemma와 proprietary Mercury 2보다 높은 성능을 기록했습니다. 초록에는 각 benchmark의 개별 점수는 제시되지 않았지만, 세 평가 영역 전체에서 8B Uno가 두 비교 모델을 앞섰다고 명시되어 있습니다.

기술 상세

Uno의 핵심은 AR 모델의 분포를 생성 모델의 기준으로 유지하면서 diffusion weights를 별도 학습하는 파라미터 분리입니다. AR weights는 표준 NTP 목표를 통해 문맥에서 다음 토큰 분포를 계산하고, diffusion weights는 이 분포에서 여러 토큰을 동시에 생성하도록 최적화됩니다. 입력은 기존 AR 문맥과 토큰 분포이고, 처리 과정은 확산 가중치의 병렬 샘플링이며, 출력은 여러 위치의 토큰입니다.

Ψ-Spec은 이 병렬 샘플링을 수행하는 sampler 계열입니다. 고정된 context length에서 AR 모델 분포를 바탕으로 여러 토큰을 함께 생성하며, 별도 draft model이 후보를 만들고 대상 모델이 검증하는 Speculative Decoding과 구조가 다릅니다. 논문은 이 과정이 기본 AR 모델의 품질을 희생하지 않는 손실 없는 가속과 inference-time scaling을 지원한다고 설명합니다.

Uno는 두 가지 구축 경로를 제공합니다. 하나는 AR weights와 diffusion weights를 함께 처음부터 학습하는 방식이고, 다른 하나는 기존 open-weight AR LLM에 경량 diffusion weights를 추가하는 방식입니다. Diffusion Distillation은 후자의 확장과 기존 학습 파이프라인에 적용할 수 있는 추가 단계로 제시되며, 논문은 학습 오버헤드가 negligible하다고 밝힙니다.

실험에서는 batch size를 바꾸며 throughput을 측정하고, 기본 AR 모델 및 leading speculative-decoding 방법과 비교했습니다. Uno는 평가한 모든 batch size에서 더 높은 throughput을 기록했으며, 장치가 지원하는 최대 batch size에서도 기본 AR 모델보다 최대 3배 빠른 결과를 냈습니다. 8B Uno는 agentic tool use, coding, long-context reasoning benchmark에서 26B DiffusionGemma와 Mercury 2를 비교 대상으로 삼았습니다.

실무 활용

Uno는 순차 토큰 생성이 병목이 되는 LLM 추론 환경에서 여러 토큰을 병렬 샘플링해 처리량을 높이는 데 활용할 수 있습니다. 코드와 체크포인트가 공개되어 있어 기존 open-weight AR LLM을 확장하거나 Uno를 처음부터 학습하는 실험에 사용할 수 있습니다.

agentic tool use 시스템에서 도구 호출 계획과 응답 생성을 처리하는 LLM 추론
coding 작업에서 긴 코드 출력을 생성하는 모델의 처리량 개선
long-context reasoning 애플리케이션에서 고정된 context length를 유지한 병렬 생성
기존 open-weight AR LLM에 diffusion weights를 추가하는 추론 가속 실험

코드 공개 여부: 공개

코드 저장소 보기

키워드

Next-Token Prediction(다음 토큰 예측)Autoregressive(자기회귀)Diffusion-Augmented LLM(확산 증강 LLM)Diffusion Distillation(확산 증류)Ψ-SpecSpeculative Decoding(추측 디코딩)Uno

용어 해설

자기회귀 생성(Autoregressive Generation): — 앞서 생성한 토큰을 다음 입력으로 다시 사용해 한 번에 하나씩 토큰을 생성하는 방식입니다. 각 단계가 이전 단계의 결과에 의존하므로 생성 품질은 유지하기 쉽지만, 긴 문장을 만들 때 순차 처리로 인해 지연 시간이 누적됩니다.
다음 토큰 예측(Next-Token Prediction): — 문맥에 등장한 토큰들을 입력으로 받아 다음에 올 토큰의 확률 분포를 학습하는 목표입니다. LLM은 이 분포를 반복해서 계산하며 문장을 생성하고, Uno의 AR 가중치도 기존 NTP 목표로 학습됩니다.
확산 증류(Diffusion Distillation): — 확산 기반 생성기가 기존 모델의 출력 분포를 모방하도록 별도의 경량 가중치를 학습하는 과정입니다. 이 논문에서는 여러 토큰을 동시에 생성하는 확산 가중치를 추가 학습하며, 기존 LLM 학습 파이프라인에 드는 추가 비용을 작게 유지합니다.
추측 디코딩(Speculative Decoding): — 별도의 draft model이 여러 후보 토큰을 먼저 생성하고, 대상 모델이 이를 검증해 순차 생성 횟수를 줄이는 추론 기법입니다. Uno는 이 구조와 달리 별도 draft model 없이 AR 분포에서 여러 토큰을 병렬로 샘플링합니다.
확산 언어 모델(Diffusion LLM): — 토큰을 순차적으로 확정하는 대신 확산 과정을 이용해 여러 토큰을 함께 생성하는 언어 모델입니다. Uno는 확산 방식으로 생성 속도를 높이면서도 기본 AR 모델의 품질을 유지하도록 설계되었습니다.

Uno로 확산 기반 병렬 토큰 생성 가속

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 기사

Diffusion 언어 모델의 속도와 품질 사이 균형

서로 다른 어휘로 LLM 추론을 최대 2.8배 가속

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

실무 활용

키워드

관련 토론

댓글

관련 기사

Diffusion 언어 모델의 속도와 품질 사이 균형

서로 다른 어휘로 LLM 추론을 최대 2.8배 가속