📝 [Research Note] Memory Wall 극복을 위한 Token-level Sparsity 및 Speculative Decoding 병합 전략

1. 문제 정의 및 목표

배경: Autoregressive 디코딩의 메모리 병목 (Memory Wall) LLM의 추론은 이전 토큰들을 기반으로 다음 토큰을 하나씩 생성하는 자기회귀(Autoregressive) 방식을 따릅니다. 이 과정에서 중복 연산을 피하기 위해 이전 토큰들의 Key, Value 벡터를 메모리에 저장해두는 KV Cache를 사용합니다.

하지만 문맥(Context)이 길어질수록 KV Cache의 크기는 선형적으로 폭발하며, 매 토큰을 생성할 때마다 이 거대한 KV Cache를 메모리(DRAM)에서 연산기(SRAM)로 끌고 와야 합니다. 결국 연산 속도보다 메모리 전송 속도가 느려 전체 추론 속도의 발목을 잡는 Memory Wall 현상이 발생합니다.

해결책 제안: 투트랙(Two-Track) 접근법 이러한 I/O 병목을 해결하기 위해 두 가지 전략을 병합합니다.

Token-level Sparsity: 보관해야 할 KV Cache의 절대적인 용량 자체를 줄이는 방법.
Speculative Decoding (추측 해독): 메모리에 접근하는 횟수(I/O Frequency) 자체를 줄이는 방법.

연구 목표
본 문서는 Token-level Sparsity의 'Heavy-Hitter' 보존 알고리즘과 Speculative Decoding의 병렬 검증 아키텍처를 분석하고, 이 두 기법이 결합되었을 때 Jetson/MPU와 같은 Edge 하드웨어의 제한된 메모리 환경에서 어떠한 시너지를 내는지 고찰합니다.

2. Token-level Sparsity: Heavy-Hitter 및 Attention Sink 유지 전략

표준 Attention 메커니즘에서 KV Cache의 메모리 요구량은 시퀀스 길이 에 비례하는 $O(N)$ 으로 증가합니다. 하지만 모든 토큰이 다음 토큰 생성에 동일하게 기여하는 것은 아닙니다.

Heavy-Hitter 및 Local Window 보존 최신 연구(H2O, StreamingLLM 등)에 따르면, Attention 점수(Score)는 극소수의 특정 토큰들에 강하게 집중되는 Sparsity(희소성) 를 보입니다.
따라서 전체 토큰을 저장하는 대신 다음 두 가지만 KV Cache에 남기고 나머지는 과감히 버림(Eviction)으로써 메모리 사용량을 $O(1)$ 에 가깝게 고정할 수 있습니다.

Attention Sink (Heavy-Hitter): 문장의 시작 부분이나 문맥을 관통하는 핵심 키워드 토큰. 누적 Attention 점수가 높은 상위 $K$ 개의 토큰.
Local Window: 방금 생성된 최신 $L$ 개의 토큰 (문맥적 연속성 유지).

이 전략을 통해 모델은 무한한 길이의 문맥을 처리하면서도 KV Cache의 물리적 크기를 제한된 SRAM 용량 이내로 통제할 수 있습니다.

3. Speculative Decoding (추측 해독): 메모리 I/O 호출 최소화

KV Cache의 크기를 줄였다 하더라도, 여전히 1개의 토큰을 생성할 때마다 가중치와 KV Cache를 메모리에서 읽어와야 하는 I/O 호출 빈도의 문제는 남아있습니다. 이를 해결하는 것이 Speculative Decoding입니다.

Draft-then-Verify 아키텍처

초안 생성 (Drafting): 매우 작고 빠른 보조 모델(Draft Model)이 미리 $\gamma$ 개의 토큰을 연속으로 빠르게 추측(Speculation)하여 생성합니다. 작은 모델이므로 가중치를 불러오는 메모리 비용이 매우 적습니다.
병렬 검증 (Verification): 크고 정확한 메인 모델(Target Model)은 Draft Model이 생성한 $\gamma$ 개의 토큰을 입력으로 받아 **단 한 번의 Forward Pass(병렬 연산)**로 검증합니다.

메인 모델과 초안 모델의 예측 확률 분포 $p(x)$ 와 $q(x)$ 를 비교하여, 일치하는 토큰은 채택(Accept)하고, 틀린 토큰이 나오면 즉시 거절(Reject) 후 메인 모델이 올바른 토큰으로 정정합니다. 이론적으로 초안 모델의 적중률(Hit Rate)이 높다면, 메인 모델은 한 번의 메모리 접근(I/O)만으로 여러 개의 토큰을 얻어낼 수 있습니다.

4. Edge 하드웨어 환경에서의 투트랙 시너지 고찰

Edge 디바이스는 고성능 서버에 비해 DRAM 대역폭이 극도로 협소하며, 빠른 속도를 내는 칩 내부의 SRAM(L1/L2 Cache) 용량은 불과 수 MB 수준입니다. 여기서 두 기법의 결합은 파괴적인 시너지를 냅니다.

SRAM 점유율과 Memory Bandwidth의 완벽한 조화

Token-level Sparsity의 역할: KV Cache 용량을 지속적으로 가지치기(Pruning)하여 크기를 고정함으로써, KV Cache 전체가 느린 DRAM으로 밀려나지 않고 빠른 SRAM 내부에 상주(Reside)하도록 보장합니다.
Speculative Decoding의 역할: Target 모델의 거대한 가중치를 DRAM에서 불러와야 하는 횟수 자체를 $\frac{1}{\gamma}$ 수준으로 대폭 줄여, DRAM 대역폭 병목을 근본적으로 우회합니다.

결과적으로, **"용량이 작은 데이터(KV Cache)는 초고속 SRAM에 두고, 용량이 큰 데이터(Target 가중치)는 느린 DRAM에서 가져오되 그 횟수를 최소화한다"**는 하드웨어 메모리 계층(Memory Hierarchy) 최적화의 정석을 소프트웨어 알고리즘으로 달성하게 됩니다.

5. 결론

Token-level Sparsity를 통한 공간(Space) 최적화와 Speculative Decoding을 통한 시간 및 I/O(Time & I/O) 최적화의 병합은 제한된 자원을 가진 On-device AI 구현에 필수불가결한 전략입니다.

Trade-off 및 한계점