Contents
1. Q-Sparse 아키텍처 및 알고리즘
1.1 Top-K 희소성
입력 텐서 \(X\)에서 절대값이 가장 큰 상위 K개의 활성화만을 선택하여 계산 효율을 향상
\[Y = (X \odot M) \cdot W^T\]\(M\)은 \(X\)의 Top-K 활성화를 나타내는 마스크 텐서, \(\odot\)는 요소별 곱셈을 의미
1.2 양자화된 Top-K 희소성
입력 텐서 \(X\)를 8비트로 양자화 후 Top-K 희소성을 적용하여 메모리 사용량 및 계산 비용 절감
\[Q(X) = \text{RoundClip}\left(\frac{X}{\gamma + \epsilon}, -128, 127\right)\]\(\gamma\)는 \(X\)의 최대 절대값, \(\epsilon\)은 0으로 나누는 것을 방지하기 위한 작은 상수
1.3 Squared ReLU
Squared ReLU 활성화 함수를 사용하여 희소성을 증가
\[\text{ReLU}^2(x) = \max(0, x)^2,\]2. 훈련 및 학습 법칙
2.1 훈련 알고리즘
STE의 정의
\[\frac{\partial Y}{\partial X} = \frac{\partial Y}{\partial (X \odot M)} \odot M.\]2.2 학습 법칙
Q-Sparse 모델의 성능은 희소 비율 \(S\)와 모델 크기 \(N\)에 따라 변함
\[L(N, S) = E + \frac{A(S)}{N^\alpha},\]\(E\)는 무한한 파라미터를 가진 모델의 성능, \(A(S)\)는 \(S\)에 의존하는 상수, \(\alpha\)는 스케일링 지수
3. 스케일링 법칙
3.1 스케일링 실험 및 발견
3.2 모델 크기 $N$의 멱법칙
고정된 희소성 비율 \(S\)를 가질 때, 스케일링 법칙은 다음과 같이 표현됨.
\[L(N, S) \equiv E + \frac{A(S)}{N^\alpha}\]3.3 희소성 비율 $S$의 지수 법칙
성능은 희소성 비율 \(S\)에 대한 지수 법칙을 따름.
\[A(S) = B + C \exp\left(\frac{\beta}{1 - S}\right)\]3.4 파라미터 피팅
L-BFGS 알고리즘을 사용하여 예측된 로그 손실과 관찰된 로그 손실 간의 Huber 손실 최소화
\[\min_{E, B, C, \beta, \alpha} \sum_{\text{Runs } i} \left(\text{Huber}_\delta \left(\log \hat{L}(N_i, S_i) - \log L_i\right)\right)\]파라미터 추정: \(E\), \(B\), \(C\), \(\alpha\), \(\beta\)는 각각 1.86, 0.01, 1.89, 0.10, 0.05로 추정됨.
3.5 희소 활성화 모델과 밀집형 모델 간의 성능 격차 감소
동일한 모델 크기 \(N\)와 동일한 희소성 비율 \(S\)을 가진 희소 활성화 모델과 밀집형 모델 간의 성능 차이 분석
\[L(N, S) - L(N, 0) = \frac{A(S)}{N^\alpha(S)} - \frac{A(0)}{N^\alpha} = \frac{A(S) - A(0)}{N^\alpha}\]3.6 인퍼런스 최적화 스케일링 법칙
활성화된 파라미터 \(N_a\)에 의존하는 형태로 스케일링 법칙 변환
\[L(N_a, S) \equiv E + A(S)\left(\frac{1 - S}{N_a}\right)^\alpha\tag{23}\]최적의 희소성 비율:
4. 실험 및 결과
[상세]
1. Q-Sparse 아키텍처 및 알고리즘
1.1 Top-K 희소성
Q-Sparse는 입력 텐서 \(X\)에서 절대값이 가장 큰 상위 K개의 활성화만을 선택하는 Top-K 희소 함수를 도입합니다. 이를 통해 비활성화된 요소에 대한 계산을 줄여, 계산 효율을 향상시킵니다. 수학적으로 표현하면 다음과 같습니다.
\[Y = (X \odot M) \cdot W^T,\]\(M\)은 마스크 텐서로, \(X\)의 Top-K 활성화를 나타냅니다. \(\odot\)는 요소별 곱셈을 의미합니다.
1.2 Quantized Top-K 희소성
입력 텐서 \(X\)를 8비트로 양자화한 후 Top-K 희소성을 적용합니다. 이는 메모리 사용량 및 계산 비용을 줄이는 데 효과적입니다. 양자화된 함수 \(Q\)는 다음과 같이 정의됩니다.
\[Q(X) = \text{RoundClip}\left(\frac{X}{\gamma + \epsilon}, -128, 127\right)\]\(\gamma\)는 \(X\)의 최대 절대값, \(\epsilon\)은 0으로 나누는 것을 방지하기 위한 작은 상수
1.3 Squared ReLU
활성화 함수로 Squared ReLU를 사용하여 활성화의 희소성을 더욱 증가시킵니다. 이는 다음과 같이 정의됩니다.
\[\text{ReLU}^2(x) = \max(0, x)^2,\]이는 일반 ReLU에 비해 더 많은 요소를 0으로 만들어, 희소성을 증가시킵니다.
2. 훈련 및 학습 법칙
2.1 훈련 알고리즘
Q-Sparse 모델의 훈련은 기존의 역전파 알고리즘을 사용하되, 희소성을 고려하여 직선 추정기 (Straight-Through Estimator, STE)를 사용하여 그래디언트를 수정합니다. 이는 희소 마스크 $$M$$을 통해 비활성화된 요소의 그래디언트가 소실되는 것을 방지합니다. 수식은 다음과 같습니다.
\[\frac{\partial Y}{\partial X} = \frac{\partial Y}{\partial (X \odot M)} \odot M.\]2.2 학습 법칙
Q-Sparse 모델의 성능은 희소 비율 \(S\)와 모델 크기 \(N\)에 따라 변화합니다. 이 관계를 설명하는 학습 법칙은 다음과 같이 표현됩니다.
\[L(N, S) = E + \frac{A(S)}{N^\alpha},\]\(E\)는 무한한 파라미터를 가진 모델의 성능, \(A(S)\)는 \(S\)에 의존하는 상수, \(\alpha\)는 스케일링 지수
3. 스케일링 법칙
최근 대규모 언어모델(LLM) 관련 연구에서는 모델 크기와 training dataset 양이 증가함에 따라 성능이 향상된다는 점을 보여주고 있습니다.(그러나 일부 논문은 고품질의 데이터로 충분하다는 입장이 있음, 그럼에도 불구하고 어느정도의 규모는 필요하다정도) HBM+22는 \(N\) 개의 파라미터를 가진 밀집형 Transformer 모델의 수렴 성능이 다음과 같은 멱법칙 스케일링을 따른다고 언급합니다.
논문의 주장
이번 연구에서는 희소 활성화된 LLM의 스케일링 법칙을 조사합니다. 희소 활성화된 LLM의 성능 또한 멱법칙 스케일링을 따른다는 것을 발견합니다.
\[L(N, S) \equiv E + \frac{A(S)}{N^\alpha}\]3.1 스케일링 실험 및 발견
희소 활성화된 LLM의 스케일링 법칙 형태를 결정하기 위해 일련의 스케일링 실험을 시작했습니다. 실험에서는 300M에서 7B까지 다양한 규모의 Q-Sparse 언어 모델을 학습했습니다. 모델은 Redpajama 데이터셋 [Com23]을 사용해 학습되었으며, 데이터 전처리를 위해 LLaMA의 Sentencepiece 토크나이저를 사용했습니다. Q-Sparse 외에도 동일한 데이터셋과 설정으로 밀집형 모델을 학습했습니다. 자세한 내용은 Appendix B에서 확인할 수 있습니다.
희소 활성화 모델과 밀집형 모델의 관찰된 손실 값은 Figure 3에 나타나 있습니다. 요약된 발견 사항은 다음과 같습니다.
이런 발견을 바탕으로 희소 활성화 모델의 성능이 모델 크기 \(N\)에 대한 멱법칙 스케일링과 희소성 비율 \(S\)에 대한 지수 법칙 스케일링의 조합을 따른다는 가설을 세웠습니다.
3.2 모델 크기 \(N\)의 멱법칙
고정된 희소성 비율 \(S\)를 가지고, 스케일링 법칙은 [KMH+20]의 스케일링 법칙을 따라야 합니다. 이는 다음과 같이 표현할 수 있습니다.
\[L(N, S) \equiv E + \frac{A(S)}{N^\alpha(S)}\]\(\alpha(S)\)는 스케일링 지수이며, 스케일링 요소 \(A(S)\)는 희소성 비율 \(S\)의 함수입니다.
주어진 모델 크기 \(N\)에 대해, 함수 \(L(N, S)\)는 희소성 비율 \(S\)에 대해 Lipschitz 연속성을 따라야 합니다. 따라서, 스케일링 지수 \(\alpha(S)\)는 비감소 함수여야 합니다. 주어진 모델 크기 \(N\)에 대해, 함수 \(L(N, S)\)는 희소성 비율 \(S\)에 대해 증가합니다. 따라서, 스케일링 지수 \(\alpha(S)\)는 비증가 함수여야 합니다. 종합적으로, 스케일링 지수 \(\alpha(S)\)는 상수여야 하며, 스케일링 함수는 다음과 같이 표현될 수 있습니다.
\[L(N, S) \equiv E + \frac{A(S)}{N^\alpha}\]3.3 희소성 비율 \(S\)의 지수 법칙
앞서 언급한 발견에 따르면, 희소 활성화 모델의 성능은 희소성 비율 \(S\)에 대한 지수 법칙 스케일링을 따릅니다. 따라서, 스케일링 요소 \(A(S)\)도 지수 법칙을 따라야 합니다. 또한, 주어진 모델 크기 \(N\)에서 스케일링 함수는 희소성 비율 \(S\)에 대해 증가합니다. 따라서, 스케일링 요소 \(A(S)\)는 비감소 함수여야 합니다. 스케일링 요소 \(A(S)\)는 다음과 같이 표현될 수 있습니다.
\[A(S) = B + C \exp\left(\frac{\beta}{1 - S}\right)\]\(B\)는 희소한 LLM의 스케일링 요소, \(C\)는 밀집형 LLM의 스케일링 요소, \(\beta\)는 희소성 비율 \(S\)에 대한 스케일링 요소 \(A(S)\)의 스케일링 지수입니다.
3.4 파라미터 피팅
희소 활성화 모델의 관찰된 손실에 스케일링 법칙의 파라미터를 맞추었습니다. L-BFGS 알고리즘 [Noc80]을 사용하여 예측된 로그 손실과 관찰된 로그 손실 간의 Huber 손실 [Hub92]을 최소화합니다.
\[\min_{E, B, C, \beta, \alpha} \sum_{\text{Runs } i} \left(\text{Huber}_\delta \left(\log \hat{L}(N_i, S_i) - \log L_i\right)\right)\]HBM+22의 방법에 따라, \(\delta\)는 \(10^{-3}\)으로 설정했습니다. 가능한 지역 최적점 주변의 초기값 그리드에서 최적의 피팅을 선택했습니다. \(E\), \(B\), \(C\), \(\alpha\), \(\beta\)는 각각 1.86, 0.01, 1.89, 0.10, 0.05로 추정되었습니다.
3.5 희소 활성화 모델과 밀집형 모델 간의 성능 격차 감소
위의 스케일링 법칙에 따르면, 동일한 모델 크기 \(N\)와 동일한 희소성 비율 \(S\)을 가진 희소 활성화 모델과 밀집형 모델의 성능을 도출할 수 있습니다. 모델 크기 \(N\)이 증가함에 따라 희소 활성화 모델과 밀집형 모델 간의 성능 격차는 감소합니다. 성능 격차는 다음과 같이 표현될 수 있습니다.
\[L(N, S) - L(N, 0) = \frac{A(S)}{N^\alpha(S)} - \frac{A(0)}{N^\alpha} = \frac{A(S) - A(0)}{N^\alpha}\]\(\alpha\)는 \(\alpha > 0\)를 만족하는 상수이므로, 모델 크기 \(N\)이 증가할수록 성능 격차는 감소합니다. 즉, 충분히 큰 모델 크기 \(N\)을 주어지면, 동일한 모델 크기를 가진 희소 활성화 모델의 성능이 밀집형 모델의 성능에 도달할 수 있습니다.
3.6 인퍼런스 최적화 스케일링 법칙
스케일링 법칙은 활성화된 파라미터 \(N_a\)에 의존하는 형태로 변환될 수 있으며, 이는 모델 인퍼런스 중의 유효 계산(FLOPs)을 반영합니다. \(L(N_a, S) \equiv E + A(S)\left(\frac{1 - S}{N_a}\right)^\alpha\tag{23}\)
$N_a$는 모델의 활성화된 파라미터 수로, $N \times (1 - S)$에 해당합니다. $A(S)$는 증가 함수이고 $(1 - S)^\alpha$
는 감소 함수이기 때문에, 희소 활성화 모델의 손실을 최소화하는 희소성 비율 $S^* > 0$이 존재합니다. 이는 희소 활성화 모델의 인퍼런스 최적화 스케일링 법칙을 이끌어냅니다.
\[L(N_a) \equiv E + A(S^*)\left(\frac{1 - S^*}{N_a}\right)^\alpha\tag{24}\]이는 희소 활성화 모델의 성능이 동일한 인퍼런스 계산 예산을 가진 밀집형 모델보다 우수함을 나타냅니다. 최적의 희소성 비율 \(S^*\)을 찾아내었으며, 이는 약 45.58%입니다. 즉, 45.58%의 희소성 비율을 가진 희소 활성화 모델이 동일한 인퍼런스 예산 \(N_a\)로 좋은 성능을 달성할 수 있습니다. 동일한 과정으로 1.58비트 Q-Sparse 모델의 인퍼런스 최적화 스케일링 법칙을 추정한 결과, 최적의 희소성 비율은 61.25%(또는 2.58\(N_a\) 파라미터)임을 발견했습니다. Figure 4는 전체 Precision와 1.58비트 가중치로 희소 활성화 모델의 인퍼런스 최적화 스케일링 곡선을 보여줍니다. 이는 동일한 성능으로, 희소 활성화 모델이 인퍼런스 중 활성화된 파라미터 수 또는 FLOPs를 크게 줄일 수 있음을 나타냅니다.
인퍼런스 최적화 스케일링 법칙은 희소성 비율 \(S\)을 조정하여 희소 활성화 모델의 성능을 최적화할 수 있음을 보여줍니다. 이는 희소 활성화 모델의 학습을 안내하고 인퍼런스 중 모델의 성능을 최적화하는 데 사용될 수 있습니다.
4. 실험 및 결과
다양한 설정하에 Q-Sparse 모델을 훈련하고, 그 성능을 기존의 밀집 모델과 비교하여 희소 모델의 효율성 및 성능을 입증합니다. 특히, 다양한 비트 수준에서의 Q-Sparse 모델의 성능을 평가하여, 최적의 희소 비율을 도출하고, 이를 통해 인퍼런스 단계에서의 계산 비용을 최소화할 수 있는 방법을 제시합니다.
[참고자료 1] Lipschitz 연속성
Lipschitz 연속성은 함수의 변화율이 일정한 상수에 의해 제한되는 개념입니다. 이 연속성 조건은 함수가 너무 급격하게 변하지 않는다는 것을 보장합니다.
정의
함수 \(f: \mathbb{R}^n \to \mathbb{R}\)가 Lipschitz 연속(Lipschitz continuous)하다고 할 때, 주어진 두 점 \(x, y \in \mathbb{R}^n\)에 대해 다음 부등식을 만족하는 상수 \(L\)가 존재하면, 이 상수를 Lipschitz 상수라고 합니다.
\[\| f(x) - f(y) \| \leq L \| x - y \|\]이 부등식은 함수 \(f\)의 변화율이 상수 \(L\)에 의해 제한된다는 것을 의미합니다. 즉, 함수의 값이 얼마나 빨리 변할 수 있는지를 나타냅니다.
간단히 말해, 함수 \(f: \mathbb{R}^n \to \mathbb{R}\)가 Lipschitz 연속성을 가진다고 하면, 어떤 상수 \(L \geq 0\)가 존재하여 모든 \(x, y \in \mathbb{R}^n\)에 대해 다음이 성립합니다. (\(L\)은 Lipschitz 상수)
\[|f(x) - f(y)| \leq L \|x - y\|\]예를 들어, 1차원에서 \(f(x) = 3x\)라면, 이 함수는 Lipschitz 연속성이 있으며, Lipschitz 상수 \(L = 3\)으로 함수의 변화율이 최대 3배까지 제한됨을 의미합니다.
Lipschitz 상수의 존재와 계산
함수가 Lipschitz 연속성을 가지는지 여부를 확인하는 방법은 여러가지가 있습니다. 대표적인 방법 중 하나는 함수의 미분을 사용하는 것입니다. 만약 함수 \(f\)가 정의역 \(\mathbb{R}^n\)에서 미분 가능하고, 그 미분이 유계라면, \(f\)는 Lipschitz 연속입니다.
구체적으로, \(f\)의 미분 \(f'(x)\)의 노름이 모든 \(x\)에 대해 유계라면 다음과 같고,
\[\| f'(x) \| \leq M\]이 때, 상수 \(M\)이 함수 \(f\)의 Lipschitz 상수가 됩니다.
예를 들어, 함수 \(f(x) = x^2\)를 고려한다면, 이 함수는 전체 실수 영역에서 Lipschitz 연속은 아니지만, 특정 구간에서는 Lipschitz 연속일 수 있습니다. \(x\)가 구간 \([a, b]\)에 있을 때, \(f(x) = x^2\)의 Lipschitz 상수를 찾을 수 있게됩니다.
함수 \(f(x) = x^2\)의 미분은 \(f'(x) = 2x\)입니다. 이 미분의 절댓값은 \(2\\|x\\|\)이며, 구간 \([a, b]\)에서 최대값은 \(2 \max(\\|a\\|, \\|b\\|)\)이므로, 이 구간에서의 Lipschitz 상수는 \(L = 2 \max(\\|a\\|, \\|b\\|)\)입니다.
소결
Lipschitz 상수는 함수의 변화율을 제한하는 중요한 개념으로, 함수의 연속성 및 매끄러움을 분석하는 데 유용합니다. 함수가 주어진 구간에서 Lipschitz 연속성을 가지는지 확인하고, 상수를 계산하는 과정은 미분을 통해 비교적 쉽게 수행할 수 있습니다.
[참고자료 2] 멱법칙
멱법칙은 특정 현상에서 한 변수의 변화가 다른 변수의 멱(power)으로 비례하는 관계를 나타내며, 주로 다양한 자연현상과 사회현상에서 관찰됩니다.
정의
두 변수 \(x\)와 \(y\)가 다음 관계를 만족하면, \(y\)는 \(x\)에 대해 멱법칙을 따른다고 합니다.
\[y = k x^n\]\(k\)와 \(n\)은 상수
물리학에서 멱법칙은 뉴턴의 중력 법칙에서 확인할 수 있는데, 뉴턴의 중력 법칙에서 두 물체 사이의 중력은 거리의 제곱에 반비례합니다.
\[F = \frac{G m_1 m_2}{r^2}\]멱법칙은 대규모 데이터에서 일정 패턴을 찾는 데 사용되며, 주로 프랙털 구조, 지수적 성장을 나타내는 현상, 자연현상에서 자주 발견됩니다.
[참고자료 3] L-BFGS 알고리즘
전체 헤시안 행렬을 사용하지 않고 근사치를 구해 뉴턴 방법과 달리 헤시안 행렬(이차 미분)을 직접 계산하지 않고, 이전 반복에서의 기울기 정보만을 사용하여 헤시안의 근사를 업데이트 함.
L-BFGS(Limited-memory Broyden-Fletcher-Goldfarb-Shanno) 알고리즘은 제한된 메모리 조건에서 준뉴턴 방법을 사용하는 최적화 알고리즘으로, 대규모 최적화 문제를 효율적으로 해결하기 위해 개발되어 사용됩니다. 이 알고리즘은 전체 헤시안 행렬을 사용하지 않고 근사치를 구성하여 메모리 사용량을 크게 줄입니다. 따라서, 고차원의 최적화 문제에서 효과적으로 작동합니다.
배경
1980년대 초에 개발된 이 방법은 큰 규모의 최적화 문제에서 기존의 뉴턴 기반 방법들이 요구하는 계산 비용과 메모리 요구량을 감소시켰습니다.
L-BFGS는 뉴턴 방법과 달리 헤시안 행렬(이차 미분)을 직접 계산하지 않고, 이전 반복에서의 기울기 정보만을 사용하여 헤시안의 근사를 업데이트합니다.
L-BFGS 알고리즘의 주요 수식은 다음과 같습니다.
\[s_k = x_{k+1} - x_k\] \[y_k = \nabla f(x_{k+1}) - \nabla f(x_k)\] \[\rho_k = \frac{1}{y_k^T s_k}\] \[H_{k+1} = \left( I - \rho_k s_k y_k^T \right) H_k \left( I - \rho_k y_k s_k^T \right) + \rho_k s_k s_k^T\]다시 한 번 정리하면,
\[B_{k+1} = B_k + \frac{y_k y_k^T}{y_k^T s_k} - \frac{B_k s_k s_k^T B_k}{s_k^T B_k s_k}\]\(B_k\)는 k번째 반복에서의 근사 헤시안 행렬, \(s_k\)는 이전 반복과 현재 반복의 파라미터 차이 \(x_{k+1} - x_k\), \(y_k\)는 그래디언트 차이 \(\nabla f(x_{k+1}) - \nabla f(x_k)\) 로, 이전의 그래디언트와 파라미터 변화를 저장하여 헤시안의 근사를 업데이트하며 계산량을 크게 줄입니다.
[참고자료 4] Huber 손실
Huber 손실은 회귀 문제에서 예측값과 실제값의 차이가 클 때는 제곱 오차를, 작을 때는 절대값 오차를 사용하는 손실 함수입니다. 이는 이상치에 강건(robust)하면서도, MSE와 MAE의 장점을 모두 취하는 방법으로 알려져 있습니다.
배경
Huber 손실은 1964년 Peter J. Huber에 의해 제안되었으며, 이 손실 함수는 제곱 오차와 절대 오차의 결합으로, 오차가 크면 제곱 오차의 민감도를 줄이고, 오차가 작으면 MAE의 정확성을 유지합니다.
Huber 손실은 다음과 같은 수식으로 정의됩니다.
\[L_\delta (a) = \begin{cases} \frac{1}{2} a^2 & \text{if } |a| \leq \delta \\ \delta (|a| - \frac{1}{2} \delta) & \text{otherwise} \end{cases}\]\(a\)는 예측값과 실제값 사이의 차이, \(\delta\)는 오차가 제곱 오차에서 절대 오차로 전환되는 임계값
Large language models (LLMs) have achieved remarkable performance on a wide range of natural language processing (NLP) tasks. However, the deployment of LLMs in real-world applications is challenging due to their high computational cost and memory footprint, especially during the inference stage. To address this challenge, recent works [MWM+24, WMD+23, SXZ+24, XGZC23, LKM23] have focused on improving the efficiency of LLMs with various approaches, including quantization [MWM+24, WMD+23, FAHA23], pruning [XGZC23], distillation [GDWH23], better decoding [LKM23], and so on. One promising approach is to use sparsity to reduce the number of activated parameters in LLMs.
Sparsity contributes two factors to the efficiency of LLMs. First, sparsity can reduce the amount of computation of the matrix multiplication as zero elements are not computed. Second, sparsity can reduce the amount of input/output (I/O) that transfers the parameters between the memory and the computation units. The I/O transfer serves as the major bottleneck in the inference stage of LLMs.
One common approach to sparsity in LLMs is to use weight sparsity, which prunes the model weights to save the computation. However, unstructured weight sparsity is difficult to parallelize in GPU devices, while structured weight sparsity has a large impact to the accuracy of the model.
Another approach is to use activation sparsity, which reduces the number of activated elements in the activation tensors. Activation sparsity can be achieved by using the mixture-of-experts (MoE) mechanism [LLX+21, FZS21], modifying the activation function [MAM+23, SXZ+24], or predicting the position to be sparsed [LWD+23]. However, these approaches do not enable full sparsity of activations in LLMs, which can limit the efficiency gains during the inference stage. Moreover, compared to the dense models, the scaling laws for the sparsely-activated LLMs have not been well studied.
To explore the full potential of sparsity in LLMs, we introduce Q-Sparse, a simple yet effective approach to enable full sparsity of activations in LLMs. The major modification on LLMs is in the linear projection (i.e., matrix multiplication). As shown in Figure 1, for each linear projection, it has a top-K sparsification function that selects the top-K activations in the input tensor. For the backprogation, we use the straight through estimator to compute the gradients of the activations. We also introduce a squared ReLU function for the feed-forward layers to further improve the sparsity of the activations. Q-Sparse can be used with both full-precision and quantized LLMs. To study the scaling law of sparsely-activated LLMs, we conduct a series of scaling experiments and derive an inference-optimal scaling law for sparsely-activated LLMs. We summarize the findings from the scaling experiments and the implications of the scaling law as below:
We also conduct experiments to evaluate the effectiveness of Q-Sparse in different settings, includ- ing training-from-scratch, continue-training of off-the-shelf LLMs, and finetuning. We show that Q-Sparse can achieve results comparable to those of baseline LLMs with the same training cost while being much more efficient at inference time.
The Q-Sparse architecture is based on the Transformer architecture [VSP+17, TLI+23] with modifications to enable sparsity in the activations.
Top-K Sparsity
The Transformer architecture uses nn.Linear to perform the projection in both attention and feed-forward layers, which can be written as:
where \(X \in \mathbb{R}^{N \times D}\) is the input tensor, \(W \in \mathbb{R}^{M \times D}\) is the weight tensor, and \(Y \in \mathbb{R}^{N \times M}\) is the output tensor. The nn.Linear operation is equivalent to the matrix multiplication operation.
We introduce a top-K sparsity function on top of the matrix multiplication operation. The top-K sparsity function is defined as:
\[Y = (X \odot M) \cdot W^T\] \[M = \text{Topk}(|X|)\]where \(M \in \mathbb{R}^{N \times D}\) is the mask tensor that indicates the top-K activations in the input tensor \(X\) in terms of the absolute values, \(\odot\) is the element-wise multiplication operation, and \(\text{Topk}\) is the function that selects the top-K elements in the tensors.
To reduce the interval around zero, we re-scale the tensor by its \(L2\) norm after performing the top-K sparsity function.
Quantized Top-K Sparsity
Recent works [WMD+23] have shown that quantization can be used to reduce the memory footprint and computational cost of LLMs without the loss of performance. We introduce a quantized version of the top-K sparsity function. The quantized top-K sparsity function is defined as:
\[Y = (Q(X) \odot M) \cdot W^T\]where $Q(\cdot)$ is the quantization function that quantizes the input tensor $X$ to an 8-bit representation:
\[Q(X) = \text{RoundClip}\left(\frac{X}{\gamma + \epsilon}, -128, 127\right)\] \[\gamma = \max(|X|)\] \[\text{RoundClip}(X, a, b) = \min(\max(\text{round}(X), a), b)\]where \(\epsilon\) is a small constant to avoid division by zero, and \(\gamma\) is the maximum absolute value in the input tensor \(X\).
Q-Sparse can be used with both full-precision and quantized LLMs. Specifically, the quantized version of Q-Sparse is compatible with 1-bit LLMs, such as BitNet b1.58 [WMD+23]. When using Q-Sparse with 1-bit LLMs, the quantization function is performed on the weight tensor \(W\):
\[Y = (Q(X) \odot M) \cdot Q_w(W)^T\]where $Q_w(\cdot)$ is the quantization function that quantizes the weight tensor $W$ to a 1.58-bit representation:
\[Q_w(W) = \text{RoundClip}\left(\frac{W}{\alpha + \epsilon}, -1, 1\right)\]where \(\alpha\) is the mean absolute value in the weight tensor \(W\):
\[\alpha = \text{mean}(|W|)\]Squared ReLU
To further improve the sparsity of the activations, we use the squared ReLU function [SML+21] for the feed-forward layers. The squared ReLU function is defined as \(\text{ReLU}(X)^2\).
Following the LLaMA architecture, we use the gated linear unit (GLU) for the feed-forward layers. The squared ReLU function is applied with the GLU function into a ReLU2GLU function. The ReLU2GLU function is defined as:
\[\text{ReLU2GLU}(X) = X W_{\text{up}} \odot \text{ReLU2}(X W_{\text{gate}})\]Most of the existing works [MAM+23] on training sparsely-activated models use the vanilla back-propagation algorithm to compute the gradient through the sparsity function:
\[\frac{\partial Y}{\partial X} = \frac{\partial Y}{\partial (X \odot M)} \odot M\]where \(M\) is the mask tensor that indicates the top-K activations in the input tensor \(X\), and \(\odot\) is the element-wise multiplication operation.
The vanilla back-propagation algorithm has a limitation. It zero-outs the gradients of the non-activated elements, which can lead to the vanishing gradient problem, especially when the sparsity ratio is high. In this work, we propose to use the straight-through estimator [BLC13] to back-propagate the gradients through the sparsity function. In this way, the gradients are passed through the sparsity function without being zeroed-out. The straight-through estimator is defined as:
\[\frac{\partial Y}{\partial X} = \frac{\partial Y}{\partial (X \odot M)}\]We visualize the average \(l2\) norm of each projection’s gradient across different layers for dense model, Q-Sparse with and without STE. We adopt top-K as 50% for Q-Sparse. Without STE, the gradient is much smaller at the bottom layers, while STE can preserve the magnitude of the gradients. As shown in Figure 2, STE estimator significantly eases the issue of gradient vanishing, especially at the bottom of the layers. We present more visualizations for each component in Appendix A.
Figure 2: The average magnitude of each projection’s gradient of dense baseline, Q-Sparse with and without STE across different layers. The visualization is conducted with 300M model size on a subset of the valid set of C4 [RSR+19]. It shows that the gradient vanishes without STE.
Q-Sparse can be used in different settings, including training-from-scratch, continue-training, and finetuning. In the continue-train and finetuning settings, we use the same architecture and training procedure as in the training-from-scratch setting. The only difference is that we initialize the model with the pre-trained weights and continue training with the sparsity function enabled.
For the pre-trained models that do not have the squared ReLU function in the feed-forward layers, we apply the top-K sparsity function after the activation function (e.g., SiLU) in the feed-forward layers. It can improve the sparsity of the activations without changing the model architecture.
Recent work on large language models has shown that the performance of LLMs scales with the model size and the amount of training data. [HBM+22] argues that the converged performance of a dense Transformer model with \(N\) parameters follows a power-law scaling law, which can be written as:
\[L(N) \equiv E + \frac{A}{N^\alpha}\]where \(L(N)\) is the performance of the model with \(N\) parameters, \(E\) is the performance of the model with infinite parameters, \(A\) is a constant, and \(\alpha\) is the scaling exponent. Note that the number of training tokens are fixed in this setting, which is part of the constant \(E\).
In this work, we investigate the scaling law of sparsely-activated LLMs. We find that the performance of sparsely-activated LLMs also follows a power-law scaling law, which can be written as:
\[L(N, S) \equiv E + \frac{A(S)}{N^\alpha}\] \[A(S) = B + C \exp\left(\frac{\beta}{1 - S}\right)\]where \(L(N, S)\) is the performance of the sparsely-activated model with \(N\) parameters and a sparsity ratio of \(S\), and \(\alpha\) and \(\beta\) are the scaling exponents.
In the following part, we will introduce how we derive the scaling law and the corresponding findings.
To determine the form of the scaling law of sparse-activated LLMs, we begin with a series of scaling experiments. In the experiments, we train a series of language models with Q-Sparse of various scales, ranging from 300M to 7B. The models are trained on the Redpajama dataset [Com23]. We use the Sentencepiece tokenizer from LLaMA to preprocess data. Besides Q-Sparse, we also train the dense baselines with the same datasets and settings. More details can be found in the Appendix B.
The observed losses of the sparsely-activated models and the dense baselines are shown in Figure 3. We summarize the findings as below:
According to these findings, our main hypothesis is that the performance of the sparsely-activated models follows a combination of a power-law scaling law with regards to the model size \(N\) and an exponential-law scaling law with regards to the sparsity ratio \(S\).
Figure 3: The scaling curves of the sparsely-activated models regarding the model size given a fixed sparsity ratio S (Left), and regarding the sparsity ratio given a fixed model size N
With a fixed sparsity ratio \(S\), the scaling law should follow [KMH+20]’s scaling law, which can be written as:
\[L(N, S) \equiv E + \frac{A(S)}{N^\alpha(S)}\]where \(\alpha(S)\) is the scaling exponent, and the scaling factor \(A(S)\) is a function of the sparsity ratio \(S\). Given any model size \(N\), the function \(L(N, S)\) should follow the Lipschitz continuity with regards to the sparsity ratio \(S\). Therefore, the scaling exponent \(\alpha(S)\) should be a non-decreasing function. Given any model size \(N\), the function \(L(N, S)\) is increasing with the sparsity ratio \(S\), so \(\alpha(S)\) should be a non-increasing function. Above all, the scaling exponent \(\alpha(S)\) should be a constant, and the scaling function can be written as:
\[L(N, S) \equiv E + \frac{A(S)}{N^\alpha}\]According to the above finding, the performance of the sparsely-activated models follows an exponential-law scaling law with regards to the sparsity ratio \(S\). Therefore, the scaling factor \(A(S)\) should also follow an exponential law. Besides, given any model size \(N\), the scaling function is increasing with the sparsity ratio \(S\). Therefore, the scaling factor \(A(S)\) should be a non-decreasing function. The scaling factor \(A(S)\) can be written as:
\[A(S) = B + C \exp\left(\frac{\beta}{1 - S}\right)\]where \(B\) is the scaling factor for extremely sparse LLMs, \(C\) is the scaling factor for dense LLMs, and \(\beta\) is the scaling exponent of the scaling factor \(A(S)\) with regards to the sparsity ratio \(S\).
We fit the parameters of the scaling law to the observed losses of the sparsely-activated models. We use the L-BFGS algorithm [Noc80] to minimize the Huber loss [Hub92] between the predicted and observed log loss.
\[\min_{E, B, C, \beta, \alpha} \sum_{\text{Runs } i} \left(\text{Huber}_\delta \left(\log \hat{L}(N_i, S_i) - \log L_i\right)\right)\]Following [HBM+22], \(\delta\) is set as \(10^{-3}\). We select the best fit from a grid of initializations around possible local optima. \(E\), \(B\), \(C\), \(\alpha\), and \(\beta\) are estimated as 1.86, 0.01, 1.89, 0.10, and 0.05, respectively.
Given the above scaling law, we can derive the performance of the sparsely-activated models and the dense baselines with the same model size \(N\) and the same sparsity ratio \(S\). The performance gap between the sparsely-activated models and the dense baselines decreases as the model size \(N\) scales. The performance gap can be written as:
\[L(N, S) - L(N, 0) = \frac{A(S)}{N^\alpha(S)} - \frac{A(0)}{N^\alpha} = \frac{A(S) - A(0)}{N^\alpha}\]Since \(\alpha\) is a constant that satisfies \(\alpha > 0\), the performance gap decreases as the model size \(N\) scales. It means that given a large enough model size \(N\), the performance of the sparsely-activated models can eventually match the performance of the dense baselines with the same model size.
The scaling law can also be transformed into a form that is dependent on the activated parameters \(N_a\), which reflects the effective compute (i.e., FLOPs) of the model during inference: \(L(N_a, S) \equiv E + A(S)\left(\frac{1 - S}{N_a}\right)^\alpha\tag{23}\)
where $N_a$ is the number of activated parameters in the model, which is equal to $N \times (1 - S)$. Since $A(S)$ is an increasing function and $(1 - S)^\alpha$ is a decreasing function, there exists a sparsity ratio $S^* > 0$ that minimizes the loss of the sparsely-activated models. This leads to the inference-optimal scaling law of the sparsely-activated models:
\[L(N_a) \equiv E + A(S^*)\left(\frac{1 - S^*}{N_a}\right)^\alpha\tag{24}\]It shows that the performance of the sparsely-activated models is better than the dense baselines with the same inference compute budget. We further solve the optimal sparsity ratio \(S^*\), finding that \(S^* \approx 45.58\%\). It means that a sparsely-activated model with a sparsity ratio of 45.58% (or 1.84\(N_a\) parameters
) can achieve the best performance with the same inference budget \(N_a\). We follow the same process to estimate the inference-optimal scaling law for 1.58-bit Q-Sparse models. We find that the optimal sparsity ratio is 61.25% (or 2.58\(N_a\) parameters). Figure 4 shows the inference-optimal scaling curves of the sparsely-activated models with full-precision and 1.58-bit weight. It shows that with the same performance, the sparsely-activated models can achieve a significant reduction in the number of activated parameters or FLOPs during inference.
The inference-optimal scaling law shows that the performance of the sparsely-activated models can be optimized by adjusting the sparsity ratio \(S\). It can be used to guide the training of the sparsely-activated models and to optimize the performance of the models during inference.
We conduct experiments to evaluate the effectiveness of Q-Sparse in different settings, including training-from-scratch, continue-training of off-the-shelf LLMs, and finetuning.
Figure 7: The training loss curves (Left) and the overall sparsity ratio (Right) of different sparsity functions. All models are trained with 300M size and 50B tokens.
Setting We train a series of language models with Q-Sparse in both full-precision and 1.58 bits. The models are trained with 50B tokens on the Redpajama dataset [Com23]. We compare Q-Sparse with the dense baselines with the same datasets and settings.
Results The observed losses of the sparsely-activated models and the dense baselines are shown in Figure 5. It shows that Q-Sparse with a 40% sparsity ratio can match the performance of the dense baselines with the same model size and training tokens.
BitNet b1.58 + Q-Sparse We further evaluate the effectiveness of Q-Sparse on 1-bit LLMs. We train a series of BitNet b1.58 models with Q-Sparse of various scales. We plot the training loss curves of both Q-Sparse and the BitNet b1.58 baseline. Figure 6 shows that the performance of the sparsely-activated BitNet b1.58 models is better than the dense baselines with the same inference compute budget. It demonstrates that Q-Sparse is compatible with 1-bit LLMs and their synergy can be used to optimize the performance of the models during inference.
Ablation Study of top-K Sparisty and STE To evaluate the effect of the top-K sparsity function, we compare the performance of the sparsely-activated models with the top-K sparsity function and the ReLU sparsity function. Moreover, we study the effect of the STE by comparing the models with and without STE. Figure 7 illustrates the results. It shows that either removing STE or replacing it with the ReLU function significantly hurts the performance. Besides, the sparsity ratio of the models with the ReLU function decreases as the training processes. In contrast, the sparsity ratio remains unchanged with the top-K sparsity function. As shown in Figure 8, we break down the contribution of the sparsity ratio from different components, finding that the decreasing sparsity is mainly from the QKV projection, the gating projection, and the up projection of the feed-forward layers. This proves the superiority of top-K over the ReLU function.
Setting We continue-train the Mistral 7B model [BBC+23] for 40B tokens on the FineWeb-Edu dataset [LBAvWW24]. We use the Sentencepiece tokenizer from Mistral to preprocess data. We use a batch size of 4M tokens and a learning rate of 5e-5. We use the Adam optimizer with a weight decay of 0.01. More training details can be found in Appendix B.
Results For a fair comparison, we continue-train the Mistral 7B model with the same recipe as the dense baseline. We compare Q-Sparse with the ReLUfication [MAM+23] and dReLU Sparsification [SXZ+24] methods, which sparsify the model by changing the activation function. Following the origin paper [MAM+23], we adopt a two-stage training strategy that first replaces the non-ReLU activation and then adds the ReLU functions. For the dReLU Sparsification method, we implement the dReLU sparsification method following the origin paper [SXZ+24]. We evaluate these models on a range of language tasks, including ARC-Challenge [YBS19], HellaSwag [ZHB+19], Winogrande [SBBC20], MMLU [HBB+21], and TruthfulQA [LHE22]. Results are shown in Table 1. It shows that Q-Sparse achieves comparable performance to the dense baseline while being much more efficient at inference time. Moreover, Q-Sparse outperforms the ReLUfication and dReLU Sparsification methods in terms of performance and sparsity ratio.
To break down the sparsity of each component in the model, we present the sparsity ratio of the query, key, value, output, up, down, and gate tensors in Table 2. It shows that Q-Sparse achieves a higher sparsity ratio than the ReLUfication and dReLU Sparsification methods. The sparsity ratio of the query, key, value, output, up, and down tensors is higher than 40%, and the sparsity ratio of the gate tensor is higher than 60%. It demonstrates that Q-Sparse can achieve full sparsity of activations in LLMs.
Setting We finetune the base model of Mistral 7B [JSM+23] and Qwen1.5 7B [BBC+23] on the Open-Orca dataset [LGP+23] for both the dense baselines and Q-Sparse. The batch size is set to 128. The learning rates are selected from {3e-6, 5e-6, 7e-6}. All models are trained with 1 epoch for a fair comparison. The hyper-parameters are detailed in Appendix B. We conduct the evaluation for these models on a range of language tasks, including ARC-Challenge [YBS19], HellaSwag [ZHB+19], Winogrande [SBBC20], MMLU [HBB+21], and TruthfulQA [LHE22].
Results The results are shown in Table 3. It shows that Q-Sparse with 3.6B activated parameters achieves significantly better performance than the Qwen1.5 4B dense model. Moreover, Q-Sparse with around 4B activated parameters achieves comparable performance to the Mistral 7B model and the Qwen1.5 7B model. It demonstrates that Q-Sparse can be used to finetune a dense pretrained model to a much more efficient sparse model with almost no loss of accuracy.
Figure 8: The sparsity ratio of each model’s component of different sparsity functions.
Scaling BitNet b1.58 + Q-Sparse + YOCO
We have shown promising results of combining 1-bit LLMs (i.e., BitNet b1.58) and fully sparse activations (i.e., Q-Sparse). We are working on scaling up the training in terms of both model size and training tokens. Furthermore, we will incorporate YOCO [SDZ+24] to address the issue of KV cache for LLM inference. The integration of BitNet, Q-Sparse, and YOCO provides a comprehensive approach to optimizing all data types in LLM inference and deployment, which includes systematic optimization of model weights, activations, and KV cache.
Q-Sparse + MoE
Mixture-of-Experts has been the most widely used method to achieve sparse activations in LLMs. Q-Sparse is orthogonal and can be seamlessly integrated with MoE.
Q-Sparse in Batch Mode
The current Q-Sparse implementation is not friendly to batch training and inference. We are working on making Q-Sparse compatible with batch mode with innovations from both modeling and system implementation.