Orthrus: Memory-Efficient Parallel Token Generation via Dual-View Diffusion

Orthrus 는 공유 KV 캐시와 정확한 합의 메커니즘을 통해 손실 없는 추론 충실도를 보장하면서 병렬 토큰 생성을 가능하게 하기 위해 경량 확산 모듈을 고정된 자기회귀 LLM 과 통합하는 메모리 효율적인 이중 아키텍처 프레임워크입니다.

핵심

긴 복잡하고 이야기를 쓰려고 한다고 상상해 보세요. 이를 수행하는 두 가지 방법이 있지만, 둘 다 치명적인 결점이 있습니다:

1. "한 단어씩 쓰는 작가"(자기회귀 모델): 이 작가는 매우 똑똑하고 정밀합니다. 한 단어를 쓰기 전에 모든 단어를 신중하게 고려하여 이야기가 완벽하게 일관되도록 합니다. 하지만 그들은 느립니다. 한 단어를 끝내고, 메모를 확인한 뒤 다음 단어를 생각하며 써야 합니다. 실수를 두려워하기 때문에 속도를 낼 수 없습니다.
2. "배치 작성자"(확산 모델): 이 작가는 한 번에 문단 전체를 쓰려고 합니다. 매우 빠릅니다! 하지만 여러 단어를 동시에 추측하면서 각각을 신중하게 검토하지 않기 때문에, 종종 논리적 오류를 범하거나 줄거리를 잃거나 터무니없는 내용을 씁니다.

Orthrus는 두 세계의 장점을 결합한 새로운 프레임워크입니다. 이는 신중한 작가의 정밀성을 잃지 않으면서 한 번에 문단 전체를 쓸 수 있게 해주는 "이중 음성" 시스템을 구축합니다.

다음은 간단한 비유를 통해 작동 방식을 설명한 것입니다:

"건축가와 시공자" 비유

AI 모델을 건축 현장으로 생각하며, 여기에는 건축가와 시공자라는 두 명의 작업자가 있습니다.

• 건축가 (동결된 LLM): 이는 원래의, 고도로 훈련된, 초지능 모델입니다. 건물이 어떻게 되어야 하는지 정확히 아는 전문가입니다. 그들은 "동결"되어 있어 이 과정에서 생각을 바꾸거나 새로운 것을 배우지 않습니다. 오직 완벽한 설계도만 제공합니다.
• 시공자 (확산 모듈): 이는 팀에 추가된 새롭고 가벼운 작업자입니다. 그들의 임무는 벽돌 (토큰) 을 빠르게 놓는 것입니다.

그들이 어떻게 협력하는가:

1. 배경 설정 (프리필링): 먼저, 건축가는 전체 프롬프트 (지시사항) 를 읽고 완벽한 고충실도 "메모리 맵"(KV 캐시라고 함) 을 구축합니다. 이 맵에는 이야기의 나머지를 구축하는 데 필요한 모든 맥락이 포함되어 있습니다.
2. 병렬 스프린트 (생성): 건축가가 한 장의 벽돌씩 놓는 대신, 시공자가 건축가의 맵을 보고 한 줄의 벽돌 (예: 32 장) 을 한 번에 놓으려 합니다.
3. 안전 점검 (합의): 이것이 마법 같은 부분입니다. 시공자의 작업이 승인되기 전에, 건축가가 시공자의 배치 작업을 즉시 점검합니다.
   • 시공자가 건축가의 완벽한 논리에 따라 다음 단어를 올바르게 추측했다면, 건축가는 "좋아! 그대로 두자!"라고 말합니다.
   • 시공자가 잘못 추측했다면, 건축가는 "아니, 그건 틀렸어"라고 말하며 즉시 그 특정 단어를 수정합니다.
   • 이 과정은 다음 배치에 대해 반복됩니다.

이것이 왜 중요한가요?

• 메모리 낭비 없음: 일반적으로 두 개의 모델이 작동하면 두 세트의 메모리 노트가 필요합니다. Orthrus 는 시공자와 건축가가 정확히 동일한 메모리 맵을 공유하기 때문에 영리합니다. 시공자는 별도의 메모를 만들 필요가 없으며 건축가의 메모만 보면 됩니다. 이는 막대한 양의 컴퓨터 메모리를 절약합니다.
• 품질 손실 없음: 건축가 (원래의 똑똑한 모델) 가 모든 단어에 대해 최종 결정을 내리기 때문에, 이야기는 건축가가 한 단어씩 썼을 때와 똑같이 훌륭합니다. "드리프트"나 품질 저하가 없습니다.
• 엄청난 속도: 시공자가 한 번에 32 장의 벽돌을 놓고 즉시 점검만 받도록 함으로써, Orthrus 는 느린 한 단어씩 쓰는 방법보다 최대 7.8 배 빠릅니다.

결과

이 논문은 수학 문제 해결 (MATH-500), 코드 작성, 논리 퍼즐 해결과 같은 어려운 작업에서 이를 테스트했습니다.

• 속도: 표준 모델보다 훨씬 빨랐습니다.
• 정확도: 원래의 느린 모델과 정확도가 동일했습니다.
• 효율성: 모델 파라미터의 아주 작은 부분 (약 16%) 만 훈련하면 되어 기존 AI 시스템에 추가하기 쉽고 저렴합니다.

간단히 말해, Orthrus는 이야기의 다음 30 단어를 즉시 추측할 수 있는 속독가를 고용하는 것과 같지만, 바로 옆에는 어떤 실수든 즉시 수정하는 엄격한 편집자가 서 있습니다. 그 결과는 여전히 완벽하게 정확한 속도로 쓰인 이야기입니다.

기술 요약

기술 요약: Orthrus – 이중 뷰 확산을 통한 메모리 효율적 병렬 토큰 생성

1. 문제 제기

자기회귀 (AR) 대형 언어 모델 (LLM) 은 높은 충실도의 생성과 강력한 추론 능력으로 인해 자연어 처리 분야에서 현재 지배적인 위치를 차지하고 있습니다. 그러나 디코딩 단계에서 근본적인 비효율성을 겪습니다: 토큰 생성이 엄격하게 순차적으로 이루어지기 때문입니다. 프리필링 (pre-filling) 단계에서는 프롬프트를 병렬로 처리하지만, 생성 단계에서는 $N$개의 토큰을 생성하기 위해 $N$개의 별도 순전파 (forward pass) 가 필요합니다. 이러한 순차적 의존성은 메모리 대역폭 병목 현상을 초래하여 하드웨어 활용도 저하와 높은 추론 지연 시간을 야기합니다.

반면, 확산 언어 모델 (DLM) 은 노이즈 제거를 통해 토큰 블록을 동시에 처리함으로써 네이티브 병렬 생성을 제공합니다. 그러나 기존 DLM 은 다음과 같은 중대한 장애물에 직면해 있습니다:

• 성능 저하: "조건부 드리프트 (conditional drift)"로 인해 유사한 규모의 AR 모델보다 종종 성능이 떨어집니다. 이는 엄격한 인과적 의존성을 위반하는 조건부 독립성 가정이 문제시되기 때문이며, 특히 복잡한 추론 작업에서 두드러집니다.
• 훈련 비용: 기준선 일관성을 달성하기 위해서는 종종 막대한 훈련 데이터셋 (예: 수천억 개의 토큰) 이나 지속적인 사전 훈련이 필요합니다.
• 아키텍처의 이질성: 사전 훈련된 AR 모델을 확산 프레임워크에 적응시키는 것은 종종 기본 가중치를 변경하여 원래 모델의 정확한 예측 분포를 파괴하고, 그 추론 능력을 일치시키지 못하게 합니다.

핵심적인 과제는 AR 모델의 높은 충실도 인과적 조건부 설정과 확산 모델의 병렬 디코딩 속도를 어느 것도 희생하지 않고 통합하는 것입니다.

2. 방법론: Orthrus 아키텍처

Orthrus 는 단일 Transformer 내에서 이러한 패러다임을 통합하는 이중 아키텍처 프레임워크를 제안합니다. AR 백본을 대체하는 대신, Orthrus 는 동결된 사전 훈련된 AR 모델에 경량의 훈련 가능한 확산 모듈을 추가합니다.

2.1 통합 이중 뷰 어텐션

이 아키텍처는 공유 키 - 값 (KV) 캐시 위에서 작동하는 두 가지 구별된 어텐션 경로를 도입합니다:

1. 동결된 AR 헤드 (파란색 경로): 이 경로는 엄격하게 동결된 상태로 유지됩니다. 유일한 기능은 프리필링 단계에서 컨텍스트를 처리하여 고충실도 인과적 KV 표현 ($K_{AR}, V_{AR}$) 을 구축하는 것입니다. 이는 정확한 예측 분포에 대한 "교사" 역할을 합니다.
2. 훈련 가능한 확산 헤드 (빨간색 경로): AR 어텐션 헤드 alongside 에 주입되는 경량 모듈 (AR 대응물에서 초기화됨) 입니다. 이는 고속 병렬 생성을 위해 특별히 설계되었습니다.

2.2 훈련: 이중 패스 블록 마스킹

훈련은 확산 뷰의 병렬 예측을 동결된 AR 모델의 정확한 목표 분포와 정렬하는 데 중점을 둡니다.

• 데이터 구성: 시퀀스에 대해 길이 $K$의 무작위 블록이 선택됩니다. 블록의 첫 번째 토큰은 가시적인 "앵커"로 유지되고, 후속 $K-1$개의 토큰은 <mask> 토큰으로 대체됩니다.
• 어텐션 메커니즘: 확산 헤드는 특수한 블록 마스크($M_{diff}$) 를 사용하여 이러한 손상된 블록을 처리합니다. 이 마스크는 두 가지 규칙을 강제합니다:
  1. 인과적 컨텍스트: 블록 내 위치는 블록 앵커에 선행하는 깨끗한 AR 컨텍스트에 인과적으로 어텐션합니다.
  2. 양방향 블록: 동일한 마스킹된 블록 내의 위치는 서로에게 양방향으로 어텐션하여 병렬 컨텍스트 집계 기능을 가능하게 합니다.
• 목적: 확산 헤드는 동결된 AR 헤드의 전체 예측 분포에 대한 순방향 KL 발산을 최소화합니다. 그래디언트는 확산 모듈을 통해서만 흐르며 AR 백본은 변경되지 않습니다.

2.3 추론: 모델 내 합의를 통한 정확한 분포 일치

Orthrus 는 합의 메커니즘을 통해 분포 드리프트 없이 병렬 생성을 달성합니다:

1. 병렬 프로젝션: 확산 헤드는 현재 앵커 토큰과 $K-1$개의 마스크를 받아 단일 순전파를 통해 $K$개의 후보 토큰을 동시에 프로젝션합니다.
2. 구조적 검증: 프로젝션된 블록은 즉시 동결된 AR 헤드를 통해 라우팅됩니다. AR 헤드가 완전히 채워진 블록을 보게 되므로, 단일 패스에서 모든 $K$개 위치의 정확한 목표 확률을 계산합니다.
3. 합의 및 커밋: 아키텍처는 엄격한 왼쪽에서 오른쪽으로의 평가를 수행합니다. 프로젝션된 토큰은 동결된 AR 헤드의 탐욕적 예측과 일치할 때만 수용됩니다. 인덱스 $j$에서 편차가 발생하면, 시스템은 $j-1$까지 동기화된 접두사를 커밋하고 정확한 AR 수정 토큰을 추가한 후 캐시를 잘라냅니다. 이는 손실 없는 추론을 보장하여 출력이 기본 모델의 예측 분포와 엄격하게 일치하도록 합니다.

3. 주요 기여

• 새로운 이중 아키텍처 프레임워크: Orthrus 는 표준 AR Transformer 내에 병렬 확산 모듈을 내장하여, 두 뷰가 중복된 역사적 KV 캐시 저장 없이 공유 KV 캐시 위에서 작동할 수 있게 합니다.
• 손실 없는 추론 보장: 모델 내 합의 메커니즘을 사용하여 Orthrus 는 기본 LLM 의 정확한 예측 분포를 유지하며, 이전 확산 적응 방법보다 우수한 엄격한 손실 없는 생성을 보장합니다.
• 중대한 추론 가속: 확산 헤드를 토큰 병렬 생성에 네이티브로 활용함으로써 Orthrus 는 순차적 병목 현상을 깨뜨려 최대 7.8 배의 속도 향상을 제공합니다.
• 극도의 파라미터 및 메모리 효율성: 통합은 경량화되어 있습니다. 병렬 기능은 총 모델 파라미터의 약 16% 만 미세 조정하여 10 억 개 미만의 토큰 (단일 8xH200 노드에서 24 시간 미만 소요) 으로 주입됩니다.

4. 실험 결과

저자들은 수학 추론 (GSM8K, MATH-500, AIME) 및 코드 생성 (HumanEval, MBPP) 벤치마크에 걸쳐 Qwen3 모델 계열 (1.7B, 4B, 8B 파라미터) 에서 Orthrus 를 평가했습니다.

• 효율성: Orthrus 는 8B 모델에서 평균 순전파당 토큰 수 (TPF) 5.39 를 달성하여 작업 및 온도 설정에 따라 3.07 배에서 7.83 배까지의 속도 향상을 실현했습니다.
• 정확도: 성능 저하를 겪는 적응 방법과 달리, Orthrus 는 기본 Qwen3-8B 모델의 정확한 제로샷 정확도를 달성했습니다. 예를 들어, MATH-500 에서 Orthrus 는 86.2% 정확도에 도달한 반면, Fast-dLLM-v2 와 같은 최첨단 확산 적응 방법은 11.1 포인트 하락 (75.1% 대 86.2% 기준선) 을 겪었습니다.
• 특정 디코딩과의 비교: 외부 특정 디코딩 방법 (EAGLE-3, DFlash) 과 비교하여 Orthrus 는 drafter 모델을 위한 별도의 중복 KV 캐시 유지가 필요하지 않기 때문에 훨씬 더 높은 평균 수용 길이 (MATH-500 에서 11.7 대 DFlash 의 7.9 및 EAGLE-3 의 3.5) 를 달성했습니다.

5. 중요성 및 주장

이 논문은 Orthrus 가 자기회귀 생성의 충실도와 확산 기반 병렬성 사이의 트레이드오프를 근본적으로 조화시킨다고 주장합니다.

• 구조적 통합: 병렬 생성을 순차적 제약에서 분리하면서도 고정된 고충실도 AR 표현에 기반을 둠으로써, Orthrus 는 다른 확산 접근법이 겪는 "분포 드리프트"를 제거합니다.
• 확장성 및 플러그 - 앤 - 플레이: 이 프레임워크는 고품질 기존 오픈소스 AR 모델에 원활하게 적응하여 엘리트 추론 능력을 희생하지 않고 병렬 처리량을 해제할 수 있는 매우 확장 가능한 솔루션으로 제시됩니다.
• 생산 타당성: $O(1)$ 메모리 캐시 오버헤드와 최소한의 파라미터 추가로 Orthrus 는 대규모 확산 모델을 처음부터 다시 훈련하는 계산 비용을 피하면서 고투과율 LLM 배포를 위한 실용적이고 메모리 효율적인 경로를 제공합니다.

저자들은 Orthrus 가 엄격한 손실 없는 추론 가속을 제공하며, 병렬 생성 충실도에 대한 새로운 최첨단 상태를 제시한다고 결론지었습니다.