Uno, 이산 확산으로 LLM 추론을 무손실 가속

들어가며

대규모 언어 모델은 대체로 한 번에 다음 토큰 하나를 예측하는 자기회귀 방식으로 텍스트를 생성한다. 이 방식은 학습과 품질 측면에서 안정적이지만, 생성 단계에서는 토큰 간 의존성 때문에 계산이 순차적으로 진행된다. GPU가 병렬 연산을 지원하더라도 이 구조적 제약은 쉽게 사라지지 않는다. 이번 논문의 Uno는 이 문제를 이산 확산 기반의 병렬 초안 생성으로 완화하려는 접근이다.

Uno의 설계

Uno는 기존 AR 모델을 완전히 확산 모델로 바꾸지 않는다. 대신 AR 모델을 유지하면서 여러 토큰을 동시에 제안할 수 있는 경량 확산 파라미터를 추가한다. 핵심 구성은 다음과 같다.

• AR 가중치 유지: 기반 모델은 표준 다음 토큰 예측 목표로 학습되며, 최종 출력 분포와 후보 검증을 담당한다.
• 확산 경로 학습: Diffusion Distillation 단계에서 추가 가중치를 학습해 여러 토큰을 병렬로 초안 생성한다.
• 병렬 초안과 AR 검증: 확산 경로가 후보 시퀀스를 제안하면 AR 경로가 이를 확인한다. 받아들여지는 후보가 많을수록 순차적인 디코딩 라운드를 줄일 수 있다.
• 별도 드래프트 모델 불필요: 기존 추측 디코딩처럼 별도의 소형 언어 모델을 배포할 필요가 없다.

논문은 Ψ-Spec이라는 샘플러 계열도 제시한다. 이 샘플러는 컨텍스트 길이를 고정한 상태에서 무손실 가속과 추론 시 확장을 지원하는 것을 목표로 한다. 초록에 따르면 Uno는 평가된 모든 배치 크기에서 비교된 주요 추측 디코딩 방식보다 높은 처리량을 보였고, 기반 AR 모델 대비 최대 약 3배의 속도 향상을 달성했다. 또한 8B Uno가 비교 대상인 오픈 확산 언어 모델 DiffusionGemma를 앞섰다고 설명한다. 다만 제공된 초록은 후반부가 잘려 있으므로, 비교 모델과 조건의 전체 내용은 논문 본문에서 확인해야 한다.

무손실 가속의 의미

무손실이라는 표현은 확산 경로가 모든 토큰을 항상 정확히 예측한다는 뜻이 아니다. 최종적인 기준은 여전히 원래의 AR 모델이다. 확산 경로는 여러 후보를 병렬로 만드는 역할을 하고, AR 경로는 해당 후보가 자신의 분포와 부합하는지 검증한다. 따라서 성능 향상은 후보 수용률, 검증 비용, 배치 크기, 하드웨어의 병렬 처리 효율에 따라 달라진다.

의미와 쟁점

Uno는 추측 디코딩과 완전한 확산형 언어 모델 사이의 절충안으로 볼 수 있다. AR 모델은 품질과 분포를 보장하고, 확산 경로는 병렬 후보 생성을 담당한다. 논문은 기존 오픈 가중치 AR 모델을 확장하거나 처음부터 Uno를 학습할 수 있다고 설명한다. 이런 구조는 별도 드래프트 모델을 관리하는 부담을 줄이면서 기존 AR 모델의 장점을 활용하려는 시도다.

한편 Hugging Face 페이지에서는 기술적 선행성과 프레임워크의 유사성을 놓고 논쟁이 있었다. 한 댓글은 선행 연구인 Orthrus 역시 AR 백본을 고정하고 학습 가능한 확산 경로를 추가하며, KV 캐시를 재사용하고, 토큰을 병렬로 초안 생성한 뒤 AR 가중치로 검증한다고 지적했다. Uno 연구진은 구현 차이를 강조했다. Uno는 기존 어텐션 구조를 유지하면서 LoRA 같은 경량 방식을 활용하는 반면, Orthrus는 확산 어텐션 헤드를 추가하고 확산 블록에서 양방향 어텐션을 사용한다는 설명이다.

이 논쟁과 별개로 Uno는 LLM 서빙의 한 방향을 보여준다. AR과 확산 중 하나를 선택하기보다, AR 모델에는 품질 보장을 맡기고 보조 확산 경로에는 병렬화를 맡기는 방식이다. 다양한 모델과 실제 부하 조건에서 보고된 효과가 재현된다면, 이러한 혼합형 디코딩은 LLM 추론 가속의 실용적인 선택지로 자리 잡을 가능성이 있다.

출처：Hugging Face Daily Papers