기술 요약: 트랜스포머 가속화를 위한 게이트 서브스페이스 추론

1. 문제 제기

대규모 언어 모델 (LLM) 추론, 특히 배치 크기가 1 인 경우의 주요 병목 현상은 계산 처리량이 아닌 메모리 대역폭입니다. 지붕선 모델 (roofline model) 에 의해 공식화되었듯이, 트랜스포머의 순전파 (forward pass) 는 선형 레이어 연산 ($y = Wx$) 의 산술 강도 (arithmetic intensity) 가 약 1 FLOP/바이트이므로 완전히 메모리 대역폭에 종속됩니다. 고대역폭 메모리 (HBM) 에서 가중치 행렬을 읽는 작업이 지연 시간을 지배합니다.

기존 가속화 전략들은 다음과 같은 중대한 한계에 직면해 있습니다:

• 양자화 (Quantization) 는 가중치 요소의 크기를 줄이지만 읽는 요소의 수를 줄이지는 않습니다.
• 프루닝 (Pruning) 은 요소의 수를 줄이지만 일반적으로 재학습이나 파인튜닝을 필요로 합니다.
• 정적 저랭크 분해 (Static Low-Rank Factorization) (예: LoRA, SVD-LLM) 는 오프라인에서 가중치를 분해하지만 추론 시 특정 입력 분포에 적응하지 못합니다.
• 활성화 희소성 (Activation Sparsity) 방법들 (예: CATS, PowerInfer) 은 이산적인 뉴런 부분집합이나 요소별 희소성을 활용하지만, 활성화 매니폴드의 연속적인 저차원 구조를 활용하지는 못합니다.

본 논문은 임의의 레이어에서 활성화 벡터가 근사적으로 저차원 서브스페이스 ($k \ll d$) 에 위치한다고 주장합니다. 결과적으로 가중치와 활성화 간의 상호작용은 낮은 유효 랭크를 가집니다. 입력 분포가 활성화 공간의 모든 방향을 방문하지 않기 때문에, 모든 토큰에 대해 전체 가중치 행렬을 읽는 것은 낭비적입니다.

2. 방법론: 게이트 서브스페이스 추론 (GSI)

GSI 는 재학습, 아키텍처 수정, 또는 어텐션 근사 없이 가중치 - 활성화 상호작용의 낮은 유효 랭크를 활용하는 손실 없는 가속화 방법입니다. 이는 MLP 레이어에서 Skyline Subspace Inference (SSI) 프레임워크를 전체 트랜스포머 (QKV 및 출력 프로젝션 포함) 로 확장한 것입니다.

2.1. 핵심 분해

가중치 행렬 $W \in \mathbb{R}^{d_{out} \times d}$와 활성화 $x \in \mathbb{R}^d$를 가진 선형 레이어의 경우, 활성화는 서브스페이스 성분과 잔차로 정확히 분해됩니다:
$$x = V_k V_k^\top x + r$$
여기서 $V_k \in \mathbb{R}^{d \times k}$는 활성화 서브스페이스를 나타내는 랭크-$k$ 직교 정규 기저이며, $r$은 잔차입니다. 출력 $y = Wx$는 다음과 같이 계산됩니다:
$$y = (W V_k)(V_k^\top x) + Wr = Mg + Wr$$
여기서 $M = W V_k$는 사전 계산되고 캐시된 "저랭크 가중치 이미지"이며, $g = V_k^\top x$는 활성화가 서브스페이스로 투영된 값입니다.

2.2. 게이트 잔차 패스스루 (Gated Residual Passthrough)

핵심 혁신은 잔차 보정 $Wr$을 계산할지 여부를 결정하는 토큰별 게이트입니다.

• 고속 경로: 잔차 비율 $\rho_t = \|r\| / \|x\|$가 임계값 $\epsilon$ 미만인 경우, 출력은 $y \approx Mg$로 근사됩니다. 이는 전체 $d$개의 열 대신 가중치 행렬의 $k$개의 열 (또는 캐시된 이미지 $M$) 만 읽으면 됩니다.
• 저속 경로: $\rho_t \ge \epsilon$인 경우, 정확성을 보장하기 위해 완전한 계산 $y = Wx$가 수행됩니다.

이 게이트는 출력 분포가 제어 가능한 허용 오차 내에서 유지되도록 보장합니다. 이론적 분석 (Theorem 4) 은 고속 경로에서의 토큰별 오차를 $\|W\|_2 \cdot \epsilon \cdot \|x\|$로 제한합니다. 중요하게도, 논문은 잔차를 완전히 폐기하는 것 (정적 투영) 은 치명적인 출력 저하를 초래하는 반면, 게이트 패스스루는 손실 없는 품질을 유지한다고 입증합니다.

2.3. 깊이 캐스케이드 및 보정

이 방법은 "깊이 일관성 (depth coherence)"을 활용합니다. 즉, 연속된 레이어의 활성화 서브스페이스가 매우 높은 상관관계를 가진다는 것을 관찰합니다 (더 깊은 레이어에서 평균 코사인 유사도 > 0.90).

• 캐스케이드 초기화: 레이어 $l+1$의 기저 $V_k$는 레이어 $l$의 기저에서 초기화되며, 모든 레이어에 대한 완전한 특이값 분해 (SVD) 가 필요한 것이 아니라 약간의 랭크 -1 보정 (DGKS 업데이트) 만 필요합니다.
• 보정: 기저는 대표 보정 시퀀스에서의 단일 순전파를 통해 한 번 계산됩니다. 6B 모델의 경우 이는 약 30 초가 소요됩니다.
• 저장 오버헤드: 캐시된 이미지 $M$과 기저 $V_k$는 모델의 메모리 사용량을 4% 미만으로 증가시킵니다 (13GB 모델의 경우 약 469MB).

3. 주요 기여

1. 게이트 서브스페이스 추론 (GSI): 학습된 저랭크 서브스페이스로 활성화를 투영하고 조건부로 잔차 보정을 건너뛰어 트랜스포머의 모든 선형 맵 (QKV, 출력, MLP 프로젝션) 을 가속화하는 방법.
2. 게이팅을 통한 손실 없는 품질: 게이트 잔차 패스스루가 필수적임을 논문이 입증합니다. 이를 사용하지 않으면 정적 투영으로 인해 생성이 붕괴됩니다. 이를 사용하면 특정 작동 지점에서 기준 모델과 문자 단위로 동일한 출력을 달성합니다.
3. 깊이 캐스케이드: 레이어 간 서브스페이스 일관성을 활용하여 보정 비용을 96% 줄이는 메커니즘으로, 네트워크가 깊이를 통해 전파되는 일관된 저차원 매니폴드에서 작동함을 보여줍니다.
4. 직교 가속화: GSI 는 숨겨진 차원 ($d$-축) 에서 작동하여 토큰 차원 ($T$-축) 에서 작동하는 ADA 와 같은 기존 어텐션 가속화 방법을 보완합니다. 함께 작동하면 전체 순전파를 커버합니다.

4. 실험 결과

실험은 AMD MI300X GPU 를 사용하여 GPT-2(124M), GPT-J(6B), OPT(6.7B) 모델로 수행되었습니다.

• 가속도: 이 방법은 선형 레이어 가중치 읽기에서 3.0 배에서 10.5 배의 유효 가속도를 달성합니다.
  • GPT-J 6B 작동 지점 $k=256, \epsilon=0.10$: 15.6 배 가속도 (99.8% 의 토큰이 고속 경로를 사용).
  • OPT 6.7B $k=256, \epsilon=0.10$: 10.5 배 가속도.
  • GPT-2 124M $k=256, \epsilon=0.10$: 3.0 배 가속도.
• 품질 보존:
  • 퍼플렉시티: 비율은 일관되게 1.00 미만입니다 (예: GPT-J 의 경우 0.991). 이는 가속화된 모델이 통계적으로 기준 모델과 구별되지 않거나 약간 더 우수함을 나타냅니다.
  • 토큰 일치: 모든 구성에서 Top-1 토큰 일치율이 98% 를 초과합니다.
  • 생성 충실도: 목표 작동 지점 ($k=256, \epsilon=0.05$ 또는 $0.10$) 에서 GPT-J 와 GPT-2 는 50 개의 탐욕적 (greedy) 토큰에 대해 기준 모델과 문자 단위로 100% 동일한 출력을 생성합니다.
• 유효 파라미터 수: 추론 비용은 "유효 파라미터 수" (활성화 공간에서 방문된 방향에 해당하는 파라미터) 에 비례합니다. GPT-J 6B 의 경우 최적 작동 지점에서 유효 개수가 6B 에서 약 375M 로 감소합니다 (16 배 감소).
• 부정적 대조군: $k=128$에서의 정적 투영 (잔차 없음) 은 퍼플렉시티를 52% 증가시키고 생성 붕괴를 초래하여 게이트의 필요성을 검증했습니다.

5. 중요성 및 주장

본 논문은 GSI 가 총 파라미터 수에서 유효 파라미터 수로 비용 모델을 근본적으로 전환시키는 트랜스포머 추론을 위한 손실 없는 가속화를 제공한다고 주장합니다.

• 재학습 불필요: 이 방법은 파인튜닝이나 아키텍처 변경 없이 사전 학습된 모델에서 작동합니다.
• 어텐션 근사 없음: 어텐션 메커니즘을 변경하지 않으므로 FlashAttention 이나 ADA 와 같은 다른 가속화 기술과 호환됩니다.
• 구조적 통찰: 결과는 가중치 - 활성화 상호작용의 낮은 유효 랭크가 훈련된 트랜스포머의 구조적 속성임을 확인시켜 주며, 서로 다른 아키텍처와 모델 크기에 걸쳐 유효함을 보여줍니다.
• 실용적 영향: 선형 레이어에 대한 메모리 대역폭 요구 사항을 $d/k$배 (예: 16 배) 줄임으로써, GSI 는 배치 크기 1 추론의 주요 병목 현상을 해결하여 훨씬 빠른 디코딩 시간을 가능하게 합니다 (예: 어텐션 가속화와 결합 시 1.3ms 순전파를 0.22ms 로 단축).

저자들은 게이트 잔차 패스스루가 출력의 수학적 정확성을 희생하지 않고 높은 압축 비율을 가능하게 하는 결정적 요소라고 결론지으며, 이를 손실 없는 추론 가속화를 위한 실현 가능한 경로로 제시합니다.