기술 요약: KnapSpec

문제 정의

대규모 언어 모델(LLM)은 모델 크기와 컨텍스트 길이가 증가함에 따라 상당한 추론 지연 시간과 계산 비용 문제에 직면해 있습니다. Speculative Decoding(SD)은 드래프트 모델을 사용하여 토큰을 제안하고 타겟 모델이 이를 병렬로 검증함으로써 생성을 가속화하지만, 별도의 드래프트 모델에 의존하는 기존 방식은 배포 복잡성과 분포 불일치 문제를 야기합니다. Self-Speculative Decoding(SSD)은 타겟 모델의 하위 네트워크를 드래프터로 사용하여 보조 모델 없이 이 문제를 해결합니다.

그러나 기존의 SSD 방식은 다음과 같은 결정적인 한계점을 가집니다:

1. 정적 휴리스틱(Static Heuristics): 대부분의 방법은 고정된 레이어 스킵 전략이나 정적 조기 종료(early-exit) 지점에 의존하며, 이는 긴 컨텍스트 시나리오에서 어텐션 메커니즘의 동적인 계산 오버헤드를 반영하지 못합니다.
2. 비효율적인 지표: 이전 연구들은 "레이어당 토큰 수(Tokens-per-Layer, TPL)"나 원시 수락률(acceptance rates)과 같은 대리 지표를 최적화하려 하지만, 이는 실제 벽시계 시간(wall-clock) 처리량과 강한 상관관계를 보이지 못합니다.
3. 하드웨어 불가지론(Hardware Agnosticism): 많은 접근 방식이 트랜스포머 레이어를 균일한 지연 시간을 가진 분리 불가능한 단위로 취급합니다. 이는 어텐션 레이어의 지연 시간은 컨텍스트 길이에 따라 선형적으로($O(n)$) 증가하는 반면, MLP 레이어의 지연 시간은 일정($O(1)$)하다는 사실을 간과합니다.

방법론: KnapSpec

KnapSpec은 드래프트 하위 네트워크의 선택을 **0/1 배낭 문제(0/1 Knapsack Problem)**로 재정의하여 Tokens-per-Time (TPT) 처리량을 극대화하는 학습 불필요(training-free) 프레임워크입니다.

1. 디커플링 및 적응형 지연 시간 모델링

전체 트랜스포머 블록을 건너뛰는 기존 방식과 달리, KnapSpec은 Attention과 MLP 레이어를 분리(decouple)합니다. 또한 하드웨어 특유의 지연 시간을 컨텍스트 길이에 따른 함수로 모델링합니다:

• $t_{Attn} = \Theta(n)$ (시퀀스 길이에 따라 확장)
• $t_{MLP} = \Theta(1)$ (상수)
  이를 통해 시스템은 컨텍스트 창이 확장됨에 따라 병목 현상이 MLP에서 Attention으로 이동하는 것에 맞춰 드래프트 모델을 적응적으로 재구성할 수 있습니다.

2. TPT 지표

KnapSpec은 TPL을 대체하여 **Tokens-per-Time (TPT)**를 주요 최적화 목표로 도입합니다. TPT는 다음과 같이 정의됩니다:
$$\text{TPT}(S, \gamma) = \frac{1 - \alpha_S^{\gamma+1}}{1 - \alpha_S} \cdot \frac{1}{\gamma \cdot t_{\text{Draft}}(S) + t_{\text{Target}}}$$
여기서 $\alpha_S$는 수락률, $\gamma$는 드래프트 길이, $t_{\text{Draft}}(S)$는 선택된 하위 네트워크 $S$의 지연 시간입니다. 이 지표는 최적화 목표를 실제 벽시계 속도와 직접적으로 일치시킵니다.

3. 배낭 문제 정식화 및 동적 계획법

레이어 선택 문제를 다음과 같이 배낭 문제로 매핑합니다:

• 아이템(Items): 개별 Attention 및 MLP 레이어.
• 무게(Weight): 레이어의 지연 시간 (정수로 정규화됨).
• 가치(Value): 드래프트 레이어의 은닉 상태(hidden state)와 타겟 모델의 은닉 상태 사이의 코사인 유사도 (수락률의 대리 지표 역할을 함).

이를 효율적으로 해결하기 위해 KnapSpec은 병렬 동적 계획법(Parallel DP) 알고리즘을 사용합니다:

1. 1단계 (DP 탐색): 코사인 유사도를 최대화하면서 다양한 지연 시간 예산에 대한 최적의 레이어 부분 집합을 계산합니다. 이 단계에서는 탐색 공간을 줄이기 위해 유사도 임계값($\tau = 0.5$)을 기반으로 한 가지치기(pruning)를 활용합니다.
2. 2단계 (그리드 탐색): DP를 통해 생성된 후보 집합으로부터 TPT 지표를 최대화하는 쌍 $(S^*, \gamma^*)$를 선택하기 위해 그리드 탐색을 수행합니다.

4. 이론적 정당성

본 논문은 은닉 상태 간의 코사인 유사도가 타겟 모델의 토큰 수락률을 위한 수학적으로 타당한 대리 지표임을 입증하는 엄격한 이론적 분석을 제공합니다. Lemma 4.1은 드래프트와 타겟 임베딩 사이의 코사인 유사도가 특정 경계(타겟의 상위 토큰 마진에 의존)를 초나면, 그리디 디코딩 예측이 반드시 동일함을 보여줍니다. 이는 코사인 유사도를 배낭 문제의 "가치"로 사용하는 것에 대한 정당성을 부여합니다.

주요 기여

1. 배낭 기반 레이어 선택: 드래프트 선택을 배낭 문제로 재정의하여, 적응형 구성을 위해 Attention과 MLP 레이어를 분리할 수 있게 한 최초의 프레임워크입니다.
2. 컨텍스트 인지 최적화: 지연 시간을 시퀀스 길이의 함수로 모델링함으로써, KnapSpec은 컨텍스트 창이 확장되어도 높은 가속도를 유지하며 실제 하드웨어의 변화하는 병목 현상에 대응합니다.
3. TPT 지표: 수락률이나 TPL보다 실제 처리량과 더 강력하게 상관되는 하드웨어 인지형 지표를 도입했습니다.
4. 이론적 토대: SSD 맥락에서 코사인 유사도와 토큰 수락률 사이의 관계를 연결하는 최초의 엄격한 증명을 제공하여, 이를 최적화 목표로 사용할 수 있는 공식적인 근거를 마련했습니다.
5. 학습 불필요 효율성: 추가적인 학습, 증류(distillation) 또는 보조 파라미터가 필요 없는 플러그 앤 플레이 방식입니다.

실험 결과

Qwen3 및 Llama3 모델을 대상으로 다양한 규모(1B ~ 70B)와 긴 컨텍스트(최대 32K 토큰) 작업(추론 및 요약)에서 실험을 수행했습니다.

• 성능: KnapSpec은 다양한 규모의 모델에서 기존의 SOTA 학습 불필요 베이스라인들(LayerSkip, Kangaroo, DEL, ASD, Draft & Verify, SWIFT, CLaSp)을 일관되게 능가합니다.
• 가속도: (예: Llama3.1-70B의 GovReport 요약 작업에서) 표준 자기회귀 디코딩 대비 최대 1.47배의 벽시계 시간 가속을 달erm합니다.
• 강건성: 그리디 디코딩과 확률적 샘플링(Nucleus sampling, $T=0.7$) 모두에서 우수한 성능을 유지합니다.
• 절제 연구(Ablation): 분석 결과, TPT가 수락률보다 처리량의 더 나은 예측 변수임을 확인했습니다 (피어슨 상관계수 0.837 vs 0.538). 또한, KnapSpec은 컨텍스트 길이가 증가함에 따라 건너뛸 Attention 레이어의 비율을 동적으로 늘리는 반면, 정적 방법들은 이에 적응하지 못함을 확인했습니다.

의의 및 주장

저자들은 KnapSpec이 별도의 드래프트 모델을 학습시키는 오버헤드 없이 차세대 생성형 AI를 가속화할 수 있는 실용적인 솔루션을 제공한다고 주장합니다. 레이어 선택을 하드웨어 특유의 지연 시간과 컨텍스트 길이를 고려한 제약 최적화 문제로 다룸으로써, KnapSpec은 처리량을 극대화하기 위해 추론의 변화하는 병목 지점을 공략합니다. 이 연구는 SSD가 정적 휴리스틱을 넘어 적응형이며 수학적으로 근거가 있는 구성 선택 방식으로 나아가는 새로운 패러다임을 제시하며, 긴 시퀀스에 대해서도 타겟 모델의 출력 분포를 보존하면서 고속 추론을 보장합니다.