SparseDecoding, 생성 단계에 맞춘 LLM 가지치기 제안

들어가며

대규모 언어 모델 추론의 병목은 항상 연산량에만 있는 것은 아니다. 디코딩 단계에서 모델은 보통 한 번에 하나의 토큰을 생성하지만, 각 단계마다 많은 가중치를 다시 읽어야 한다. 이 때문에 메모리 접근과 대역폭이 지연 시간에 큰 영향을 준다. 가지치기를 통해 0이 아닌 파라미터 수를 줄일 수 있지만, 실제 속도 향상을 얻으려면 가지치기 기준과 실행 커널이 토큰 단위 디코딩의 특성에 맞아야 한다.

Westlake ENCODE Lab 연구진은 이러한 문제를 겨냥해 SparseDecoding을 제안했다. 알고리즘 측면에서는 고정된 자연어 시퀀스 대신 밀집 모델이 자기회귀적으로 생성하는 동안의 활성값을 수집해 가지치기를 보정한다. 시스템 측면에서는 디코딩에서 중요한 희소 행렬-벡터 곱(SpMV)을 위한 N:M 희소 커널을 개발했다. Llama-3.1-8B, Llama-3.3-70B, Qwen3-14B, Qwen3-32B 등을 대상으로 한 실험에서 A100 GPU 기준 최대 1.48배의 종단 간 디코딩 가속을 보고했다.

기존 보정 방식의 분포 불일치

일부 학습 없는 가지치기 방법은 미리 수집한 자연어 시퀀스에서 Hessian 관련 정보를 계산하고, 이를 바탕으로 가중치의 중요도를 추정한다. 구현은 간단하지만 자연어 입력에서 관찰한 활성 분포가 실제 생성 상황을 대표한다는 가정이 필요하다.

자기회귀 디코딩에서는 모델이 사람이 고정한 시퀀스만 처리하지 않는다. 이전 단계에서 모델이 직접 생성한 토큰이 다음 입력의 일부가 되며, 생성이 계속되면서 은닉 상태와 계층별 활성 분포도 달라질 수 있다. 논문은 이러한 차이가 가지치기 목표와 실제 생성 동작 사이의 불일치를 만들고, 가지치기 모델의 성능을 떨어뜨릴 수 있다고 설명한다.

SparseDecoding은 밀집 모델이 자기회귀 생성을 수행하는 동안 각 계층의 활성값을 수집한다. 프리필 단계는 제외하고, 최종적으로 가속하려는 반복적인 토큰 생성 단계에 보정 데이터를 맞춘다. 따라서 특정 고정 문장에서 어떤 가중치가 중요한지를 보는 대신, 모델이 토큰을 연속적으로 생성할 때 어떤 가중치가 중요한지를 반영하려는 접근이다.

핵심 내용

• 디코딩 인식 보정: 고정 자연어가 아니라 밀집 모델의 생성 중 얻은 계층별 활성값을 사용한다.
• SpMV 중심 최적화: 디코딩을 지배하는 희소 행렬-벡터 곱을 대상으로 하며, 희소 행렬-행렬 곱(SpMM)에 치우친 기존 지원을 보완한다.
• 알고리즘과 시스템의 공동 설계: 비트마스크 인덱싱과 고정 스텝 순회를 적용한 N:M 희소 SpMV 커널을 구현했다.
• 여러 모델군 검증: Llama와 Qwen 계열의 다양한 모델을 장문 생성 벤치마크에서 평가했다.

의미와 한계

파라미터나 0이 아닌 원소를 줄이는 것만으로 실제 벽시계 시간이 자동으로 줄어들지는 않는다. 런타임이 불규칙한 벡터 단위 희소 연산을 효율적으로 처리하지 못하면, 이론적 메모리 절감이 지연 시간 개선으로 이어지기 어렵기 때문이다. SparseDecoding은 보정 데이터의 분포와 저수준 SpMV 구현을 하나의 배포 문제로 함께 다룬다는 점에서 의미가 있다.

다만 보고된 가속 수치는 특정 모델, N:M 희소 패턴, A100 GPU 환경에 기반한다. 실제 효과는 하드웨어, 배치 크기, 시퀀스 길이, 추론 프레임워크 구현에 따라 달라질 수 있다. 그럼에도 주된 사용 방식이 순차적인 토큰 생성인 시스템이라면, 생성 과정 자체에서 보정 신호를 수집하자는 관점은 단순히 희소율을 높이는 것보다 실제 서비스 조건에 가까운 최적화 방향을 제시한다.

출처: Hugging Face Daily Papers