Speculative Decoding과 양자화(Quantization)를 결합한 기법인 QSPEC에 대해, 개인적인 공부를 위해 그 요점이 되는 내용을 정리했습니다.

QSPEC은 저정밀도의 빠른 **W4A4를 드래프트 생성(Draft Generation)**에, 고정밀도의 **W4A16을 검증(Verification)**에 이용하는 Speculative Decoding 기법입니다. 양쪽 모두에서 양자화된 Weight와 KV 캐시를 공유함으로써, 메모리 사용량을 억제하면서도 W4A16에 상응하는 정밀도와 추론 효율화를 동시에 달성했다고 보고되었습니다.

Juntao Zhao, Wenhao Lu, Sheng Wang, Lingpeng Kong, and Chuan Wu, "QSPEC: Speculative Decoding with Complementary Quantization Schemes", 2025.

논문의 배경 지식이 되는 기술 용어

모델의 가중치(Weight)만을 저정밀도로 만드는 방식. 예를 들어, W4A16에서는 모델의 가중치를 4비트로 양자화하여 저장하고, 활성화(Activation)에는 FP16 등의 16비트 형식을 사용한다.

가중치와 활성화 모두를 동시에 양자화하는 방식. 고정밀도로의 역양자화(Dequantization)를 수행하지 않고, 저정밀도 커널(Kernel)로 연산한다. 저정밀도 연산에 대응하는 하드웨어를 활용하여 고속으로 추론할 수 있다.

W4A16은 모델의 가중치를 INT4 등으로 저장하고, 활성화는 FP16이나 BF16으로 다룬다. 실제 구현에서는 4비트 가중치를 GPU로부터 읽어 들인 뒤, 계산 시에는 FP16 등으로 역양자화하여 행렬 곱(Matrix Multiplication)을 수행하기도 한다.

W4A4는 가중치와 활성화 모두를 4비트로 양자화한다.

행렬 곱의 입력은 4비트이지만, 곱셈 누적 연산(MAC)의 결과는 값이 커지기 때문에 통상적으로 INT32 등의 높은 정밀도로 누적한다. 그 후, 스케일(Scale)을 적용하여 FP16이나 INT4 등으로 변환한다.

가중치의 역양자화(Dequantization)란, INT4 등으로 작게 저장된 가중치를 행렬 곱에서 다룰 수 있는 FP16/BF16 등의 값으로 되돌리는 처리이다.

예: W4A16에서의 역양자화

W4A16에서는 가중치는 4비트, 활성화는 FP16/BF16이 된다.

Y = X_{FP16} \hat W_{FP16}

여기서,

\hat W_{FP16} = s \big(W_{INT4} - z\big)

로 하여, 4비트 가중치를 FP16 등으로 역양자화한 뒤 행렬 곱을 수행한다.

INT4, INT8, FP8 등의 데이터 타입에 맞춰 실행 효율화를 도모한 커널. 구체적으로는 다음과 같은 이유로 가속화된다.

• 메모리에서 읽어오는 데이터를 줄임

FP16 가중치를 W4로 만들면 이론적인 가중치 크기는 1/4이 되어, 메모리에서 전송하는 데이터량을 줄일 수 있다. - 캐시(Cache)나 레지스터(Register)를 유효하게 활용함

데이터를 작게 만듦으로써 캐시나 공유 메모리(Shared Memory)에 더 많은 요소를 둘 수 있다. - 전용 연산기의 처리량(Throughput)을 높일 수 있음

GPU나 NPU에는 저정밀도 행렬 곱을 고속으로 처리하는 전용 연산기가 있다. 일반적으로 저정밀도가 될수록 1개 명령어로 처리할 수 있는 요소의 수를 늘릴 수 있다.

가장 큰 이유는 1개 토큰을 생성하기 위해 거대한 가중치와 KV 캐시를 읽어 들여야 하는 반면, 소규모 배치(Small Batch)에서는 그것들을 충분히 재사용할 수 없기 때문이다.

Prefill 단계에서는 프롬프트 내의 다수 토큰을 병렬 계산하여 토큰을 병렬로 계산할 수 있지만, Decode 단계에서는 원칙적으로 1스텝씩 순차적으로 생성해야 한다.

Speculative Decoding은 작고 빠른 Draft 모델이 여러 개의 후보 토큰을 먼저 생성하고, 큰 Target 모델이 그것들을 한꺼번에 검증함으로써 출력 품질을 유지하면서 Decode를 가속화하는 기법이다.

이전의 문맥 $ h $에 대하여, Draft 모델의 확률 분포를 $ q(x|h) $, Target 모델의 확률 분포를 $ p(x|h) $라고 했을 때, Draft가 후보 $ x $를 생성했을 때 기본적인 Speculative Sampling에서는 다음 확률로 수락한다.

a(x) = min \big( 1, \frac{p(x|h)}{q(x|h)}\big )

항상 최대 확률의 토큰을 선택하는 Greedy Decoding (탐욕적 디코딩)에서는, Target 모델의 최대 확률 토큰이 Draft 후보와 일치하면 수락된다.

수락률이 높을수록 Target 모델의 1회 검증으로 확정할 수 있는 토큰 수가 늘어나기 때문에, 기본적으로 가속화율도 높아진다.

단, Speculative Decoding (추측적 디코딩)에서는 첫 번째 거절(rejection)이 일어난 시점에서 그보다 뒤의 후보를 채택할 수 없으므로, 선두로부터의 토큰 수락률이 중요하다.

1 라운드에서 수락되는 Draft 토큰 수 $ A $의 기대값은 다음과 같이 나타낼 수 있다.

E[A] = \alpha + \alpha^2 + ... + \alpha^K = \frac{\alpha(1-\alpha^K)}{1-\alpha}

표준적인 Speculative Sampling (추측적 샘플링)에서는 Draft 후보가 도중에 거절된 경우, Target 모델의 분포로부터 보정 토큰을 1개 생성한다. 또한, 모든 후보가 수락된 경우에도 Target이 계산한 추가의 1토큰을 사용할 수 있다. 따라서 1라운드에서 확정되는 토큰 수 $ L $은,

L = A + 1

으로 생각할 수 있으며, 그 기대값은 다음과 같이 나타낼 수 있다.

E[L] = 1 + \alpha + \alpha^2 + ... + \alpha^K = \frac{1-\alpha^{K+1}}{1-\alpha}

일반적인 Target 모델이 1토큰을 생성하는 시간을 $ T_{target} $이라 하자.

Speculative Decoding의 1라운드에 필요한 시간을 다음과 같이 정의한다.

T_{round} = T_{draft} + T_{verify} + T_{overhead}

이 경우의 가속화율은 다음과 같이 나타낼 수 있다.

S \approx \frac{E[L]T_{target}}{T_{draft}+T_{verify}+T_{overhead}}

여기서는 논문에 기재된 다음 두 가지 관찰 결과에 대해 기술한다.

• W4A4는 다단계 추론 태스크에서 최종 성능이 크게 저하된다.
• 반면, 토큰 단위로 보면 W4A4와 W4A16의 예측은 매우 유사하다.

Atom이나 QuaRot 등의 Weight-Activation 양자화 기법은 일반적인 벤치마크에서 Weight-only 양자화에 가까운 성능을 유지하면서 추론을 가속화할 수 있다고 보고되어 있다.

하지만 QSpec의 논문 저자들은 기존 평가에서 사용되어 온 태스크만으로는 활성화 양자화(activation quantization)에 의한 성능 저하를 충분히 포착할 수 없다고 말한다. 또한 이를 보여주는 벤치마크로서, Llama-3-8B-instruct에 Atom 기반 양자화를 적용하여 다음 태스크에서 W16A16, W4A16, W4A4를 비교하였다.

• WikiText-2: 문장에 대한 예측 문제
• PIQA: 물리적 상식 추론
• GSM8K: 여러 단계가 필요한 산수 문제
• MBPP: Python 프로그램 생성 문제

결과적으로 WikiText-2나 PIQA에서는 W4A4의 성능 저하가 비교적 작지만, GSM8K나 MBPP에서는 W4A4의 성능이 평균 약 30% 저하되는 반면, W4A16은 약 4% 저하에 그친다.

이 결과로부터 W4A4의 활성화 양자화는 GSM8K나 MBPP와 같은 다단계 추론 태스크에 영향을 미친다고 볼 수 있다.

논문의 저자들에 따르면, 토큰 단위의 예측에서는 W4A4와 W4A16의 Top-1 예측 사이에 유사성이 있다. 이 성질을 확인하기 위해 GSM8K 테스트셋을 사용하여 다음과 같은 실험을 수행하였다.

• Atom 기반의 A4W16을 사용하여 Greedy Decoding을 수행하고, 각 문제에 대한 답변 토큰을 생성한다.
• Atom 기반의 W4A16이 선택한 Top-1 토큰과 그 예측 확률을 기록한다.
• 문제문과 W4A16이 생성한 답변을 연결하여, W4A4로 순전파(forward pass)를 1회 실행한다.
• 각 토큰 위치에서의 W4A4와 W4A16의 예측을 비교한다.

논문의 그림 2를 통해 다음 두 가지 점을 확인할 수 있다.

• W4A4와 W4A16 모두 토큰 예측 확률의 상당 부분이 80%를 넘는다.
• 거절되는 토큰은 수락되는 토큰에 비해 매우 적다.

즉, W4A4는 최종적인 태스크 성능은 크게 저하되지만, 생성하는 토큰의 대부분은 W4A16과 일치합니다.

대부분의 토큰이 일치함에도 불구하고 W4A4의 성능이 크게 저하되는 이유에 대해, 저자들은 소수의 중요한 토큰에서 발생하는 예측의 차이 때문이라고 설명합니다. 자기회귀(Autoregressive) 방식의 LLM은 그때까지 생성한 모든 토큰으로부터 다음 토큰을 예측하기 때문에, 어떤 추론 단계에서 잘못된 토큰을 생성하면 그 오류를 포함한 문맥을 바탕으로 후속 토큰이 생성되어 이후의 추론 전체가 잘못된 방향으로 진행될 가능성이 있기 때문입니다.

앞선 실험 결과로부터, 대부분의 토큰을 빠른 W4A4로 생성하고 소수의 중요한 오류만을 W4A16으로 검출 및 수정한다는 QSPEC의 착안을 얻을 수 있습니다.

여기서는 QSPEC의 구성에 대해 설명합니다.

QSPEC에서는 다음 $\gamma$개의 토큰 $\hat T_{i+1:i+\gamma}$와 관련된 분포 $\hat {p}_{i+1:i+\gamma}(t)$를 예측하기 위해, W4A4를 통한 Weight-Activation 양자화와 가중치를 공유하는 양자화 스킴(Quantization Scheme)을 채택하고 있습니다. 활성화(Activation)의 정밀도를 낮춤으로써 빠른 토큰 생성이 가능해집니다.

W4A4의 성능 저하를 보완하기 위해, 검증(Verify) 단계에서는 고정밀 가중치 전용 양자화를 사용합니다.

구체적으로는, 고정밀 양자화 모델 $M_h$가

\begin{align}
&T_{\le i} \\
&\hat{T}_{i+1:i + \gamma}
...

을 연결한 것을 입력으로 받아, 고품질의 예측 확률 $p_{i+1:i+\gamma+1}(t)$를 출력합니다.

QSPEC에서는 W4A4와 W4A16의 계산에서 별도의 가중치를 저장하지 않고, 동일한 4비트 양자화 가중치를 사용합니다. 즉, W4A4와 W4A16의 차이는 가중치가 아니라 주로 활성화의 정밀도와 계산 커널(Kernel)입니다.

QSPEC에서는 수락된 토큰에 대해, W4A4에 의한 저정밀 KV 캐시(KV Cache)를 W4A16에 의한 캐시로 교체합니다. 이를 통해 후속 디코딩(Decoding) 처리에서 고품질의 문맥 정보를 이용할 수 있도록 합니다.

단일 모델 내에서 가중치를 공유하고 KV 캐시를 재사용함으로써, QSPEC는 이중 캐시 관리를 불필요하게 만들며 정밀도를 희생하지 않고도 메모리 사용량을 줄일 수 있습니다.

논문의 Table 3에서는 Atom과 QuaRot 두 종류의 양자화 기법에 대해 W16A16, W4A16, QSPEC, W4A4의 정밀도를 비교하고 있습니다. 결과적으로 W4A4는 GSM8K, MATH, MBPP 등 복잡한 추론 태스크에서 정밀도가 크게 저하되지만, QSPEC는 W4A16과 동등한 수준의 정밀도를 유지함을 보여줍니다.

QSPEC는 W4A4의 결과를 W4A16의 예측과 비교하여 검증하기 때문에, 원리적으로 W4A16과 동등한 정밀도를 얻을 수 있습니다. 다만, PyTorch의 비결정적 연산(Non-deterministic operation)이나 수치 오차로 인한 영향으로 완전히 일치하지는 않습니다.

논문의 Table 4에서는 3B, 7B, 13B Llama 모델에서 배치(Batch) 수를 8, 16, 32로 변화시키며 W4A16, W4A4, QSPEC의 추론 속도를 비교하고 있습니다. 결과적으로 W4A16에 비해 평균 1.4배, 최대 1.64배의 성능 향상을 보였습니다.

이는 다음과 같은 이유 때문인 것으로 생각됩니다.

• W4A4의 저정밀 커널을 통한 빠른 드래프트(Draft)
• W4A16을 통한 복수 드래프트 토큰의 병렬 검증

논문의 Table 5에서는 Llama-2-7b-chat-hf에서의 QSPEC와 EAGLE의 비교를 보여줍니다. EAGLE은 단일 시퀀스 입력(배치 사이즈 1)에서 최적의 성능을 발휘하지만, 배치 사이즈가 증가함에 따라(8 및 16) 효율이 저하됩니다. 또한 EAGLE의 드래프트 모델은 KV 스토리지 사용량이 늘어나기 때문에, 16 배치에서 OOM(Out of Memory)이 발생했습니다. 반면 QSPEC는 우수한 확장성(Scalability)과 메모리 효율성을 보여준다고 명시되어 있습니다.

본 기사에서는 Speculative Decoding에 양자화를 결합한 QSPEC 논문을 읽고 그 핵심 내용을 정리했습니다. 다음으로는 논문에서 제시된 사항들에 대해 실제로 재현 검증을 진행해 보고자 합니다.