1. 서론
딥러닝에서의 행렬 곱셈(MatMul)은 주로 GPU가 최적화되어 있는 동작으로, Dense Layer에서는 벡터-행렬 곱셈(Vector-Matrix Multiplication, VMM), 컨볼루션은 블록-희소 VMM으로, 그리고 자기-주의 메커니즘은 행렬-행렬 곱셈(Matrix-Matrix Multiplication, MMM)으로 구현됩니다. 이런 MatMul 작업의 효율화는 CUDA 및 cuBLAS와 같은 선형 대수 라이브러리를 통해 병렬화 및 가속화가 가능합니다. 이는 AlexNet의 ILSVRC2012 대회 승리와 딥러닝의 부상에 결정적인 역할을 했습니다. 그러나, MatMul 작업은 광범위한 데이터 이동 및 동기화를 필요로 하며, 이는 계산 처리량 및 효율성을 크게 저해할 수 있습니다. 예를 들어, BitNet 모델에서는 양자화된 주의 메커니즘의 행렬들이 성능 저하와 모델 수렴 실패를 초래했습니다($\text{Fig. 1}$ 참조).
이 연구에서는 Dense Layer에서의 덧셈 작업 및 자기-주의 같은 기능에 대해 원소별 하다마드 곱(Hadamard Product)만을 사용함으로써 모든 MatMul 작업을 제거하는 첫 번째 확장 가능한 MatMul-free 언어 모델을 개발했습니다. 특히, 삼진 가중치를 사용함으로써 Dense Layer에서 MatMul을 제거했고, 자기-주의를 제거하기 위해 GRU를 최적화하여, SOTA Transformers와 경쟁하면서 모든 MatMul 작업을 제거할 수 있었습니다.
2. 관련 작업
언어 모델에 대한 이진화, 삼진화 및 저Precision 양자화는 BERT를 삼진 BERT로 축소하고, GLUE 벤치마크에서 평균 41% 정확도를 달성한 것으로 시작되었습니다. 최근 연구에서는 점진적인 양자화 접근 방식을 소개하여 모델을 32비트에서 2비트, 그리고 최종적으로 이진 모델 파라미터로 점진적으로 양자화했습니다. 이와 병행하여, 저Precision 언어 생성 모델이 발전하면서 3억 파라미터를 가진 이진 및 삼진 모델인 BitNet은 Llama와 같은 언어 모델과 경쟁력 있는 성능을 유지했습니다.
MatMul-free Transformers는 주로 SNNs의 영역에서 사용되었습니다. Spikformer는 Transformer 아키텍처를 SNN과 통합한 첫 사례를 이끌었고, 이후의 연구는 비전 작업에서 성공을 입증했습니다. 언어 이해 분야에서는 SpikeBERT가 BERT에 SNN을 적용하여 감성 분석을 수행하는 데 지식 증류 기술을 사용했습니다. 언어 생성에서는 spiking RWKV 아키텍처를 사용하여 216M 파라미터 생성 모델인 SpikeGPT를 훈련했습니다. 이런 모델들은 크기가 제한적이며, 이는 이진 활성화로 확장하는 데의 챌린지를 반영합니다. BNNs도 이 영역에서 상당한 진전을 보였습니다. BinaryViT와 BiViT는 시각적 작업에 이진 비전 Transformers를 성공적으로 적용했습니다. 더 나아가, Kosson 등은 곱셈, 나눗셈, 비선형성을 조각별 선형 근사로 대체하여 성능을 유지하면서 곱셈 없는 훈련을 달성했습니다.
3. 방법
이 섹션에서는 제안된 MatMul-free 언어 모델(LM)의 구성 요소를 자세히 설명합니다. 먼저 삼진 가중치를 사용하는 MatMul-free Dense Layer(BitLinear 층)에 대해 설명하고, 토큰 믹서와 채널 믹서를 포함한 MatMul-free LM 아키텍처에 대한 추가적인 세부사항을 제공합니다.
3.1 삼진 가중치를 이용한 MatMul-free Dense Layer
표준 Dense Layer에서는 입력 \(x \in \mathbb{R}^{1 \times d}\)와 가중치 행렬 \(W \in \mathbb{R}^{d \times m}\) 간의 MatMul이 다음과 같이 표현됩니다.
\[y = xW = \sum_{j=1}^d x_j W_{ij} \quad \text{for } i = 1, 2, \ldots, m\]\(y \in \mathbb{R}^{1 \times m}\)는 출력입니다. 표준 MatMul 기반 Dense Layer을 사용하지 않기 위해, BitNet을 채택하여 MatMul을 포함한 Dense Layer을 삼진 가중치를 사용하는 BitLinear 모듈로 대체하였습니다. 삼진 가중치를 사용할 때, 가중치 행렬 \(W\)의 원소는 {-1, 0, +1} 집합으로 제한됩니다. \(\tilde{W}\)를 삼진 가중치 행렬이라고 하면, 삼진 가중치를 사용한 MatMul은 다음과 같이 표현될 수 있습니다.
\[\tilde{Y} = x \odot \tilde{W} = \sum_{j=1}^d x_j \tilde{W}_{ij}, \quad \tilde{W}_{ij} \in \{-1, 0, +1\} \quad \text{for } i = 1, 2, \ldots, m\]\(\tilde{Y} \in \mathbb{R}^{1 \times m}\)는 출력이며, \(\odot\)은 삼진 MatMul을 나타냅니다. 이는 누적으로 단순화될 수 있습니다. 삼진 가중치 \(\tilde{W}_{ij}\)는 {-1, 0, +1} 값만을 가질 수 있으므로, MatMul에서의 곱셈 연산은 간단한 덧셈 또는 뺄셈 연산으로 대체될 수 있습니다.
\[x_j \tilde{W}_{ij} = \begin{cases} x_j & \text{if } \tilde{W}_{ij} = 1 \\ 0 & \text{if } \tilde{W}_{ij} = 0 \\ -x_j & \text{if } \tilde{W}_{ij} = -1 \end{cases}\]따라서, 삼진 MatMul은 다음과 같이 쓸 수 있습니다.
\[\tilde{Y}_i = \sum_{j: \tilde{W}_{ij} =1} x_j - \sum_{j: \tilde{W}_{ij} =-1} x_j \quad \text{for } i = 1, 2, \ldots, m\]3.2 하드웨어 효율적인 Fused BitLinear 층
BitNet은 삼진 층을 안정화하기 위해 BitLinear 입력 전에 추가적인 RMSNorm이 필요함을 보여주었습니다. 그러나 기존의 BitNet 구현은 효율적이지 않습니다. 현대의 GPU는 대용량의 글로벌 고대역폭 메모리(HBM)와 더 작고 빠른 공유 메모리(SRAM)를 특징으로 하며, BitNet의 구현은 많은 입출력 작업을 도입했습니다. 이전 활성화를 SRAM으로 읽어 RMSNorm을 실행하고, 양자화를 위해 다시 쓰고, 양자화를 위해 다시 읽고, 저장하고, 선형 작업을 위해 한 번 더 읽습니다. 이런 비효율성을 해결하기 위해, 활성화를 한 번만 읽고 SRAM에서 RMSNorm과 양자화를 융합된 연산으로 적용합니다.
Algorithm 1은 SRAM의 최적 활용을 통해 HBM의 입출력 비용을 줄이면서 BitLinear 층의 하드웨어 효율성을 향상시키는 접근 방식을 제시합니다. 이 모델에서는 삼진 가중치와 원소별 연산의 양 때문에 활성화가 가중치보다 더 큰 메모리 FOOTPRINT를 가지므로, 최적화 노력은 활성화에 초점을 맞춥니다.
3.3 MatMul-free 언어 모델 아키텍처
이 섹션에서는 토큰 믹서와 채널 믹서를 포함한 MatMul-free 언어 모델의 구조에 대해 설명합니다.
3.3.1 MatMul-free 토큰 믹서
현대 언어 모델에서 가장 일반적인 토큰 믹서는 자기-주의(self-attention)로, 세 행렬 Q, K, V 간의 행렬 곱셈에 의존합니다. 이 연산들을 덧셈으로 전환하기 위해, 적어도 두 개의 행렬을 이진화하거나 삼진화합니다. 모든 Dense Layer 가중치가 삼진화된다고 가정하면, Q와 K를 양자화하여 삼진 주의 맵을 생성함으로써 자기-주의에서 곱셈을 제거합니다. 그러나 실험에서 보듯이, 이런 방식으로 훈련된 모델은 수렴하지 못합니다. 이는 활성화에 중요한 이상치들이 포함되어 있고, 이를 효과적으로 양자화하기 어렵기 때문일 수 있습니다. 이런 챌린지를 해결하기 위해, 행렬 곱셈에 의존하지 않고 토큰을 믹싱하는 대안적 방법을 모색합니다.
삼진 RNN을 사용함으로써 원소별 연산과 누적을 결합하여 MatMul-free 토큰 믹서를 구성할 수 있습니다. 다양한 RNN 아키텍처 중에서, GRU는 그 간결성과 효율성으로 인해 주목받으며, LSTM과 유사한 성능을 더 적은 게이트와 간단한 구조로 달성합니다. 따라서 GRU를 MatMul-free 토큰 믹서를 구축하기 위한 기반으로 선택합니다. 표준 GRU를 재검토한 후, MLGRU를 어떻게 파생하는지 단계별로 설명합니다.
MatMul-free 선형 게이트 순환 유닛
우선, hidden state 관련 가중치 \(W_{cc}\), \(W_{hr}\), \(W_{hf}\) 및 hidden state 간의 활성화(\(\tanh\))를 제거합니다. 이 수정은 모델을 MatMul-free로 만들 뿐만 아니라 변형된 RNN들, 예를 들어 Linear Recurrent Unit, Hawk, RWKV-4 같은 최근의 RNN 모델에서 볼 수 있는 특징입니다. 이어서, ht와 ot 사이에 데이터 의존적 출력 게이트를 추가합니다.
\[g_t = \sigma(x_t \odot W_g + b_g) \in \mathbb{R}^{1 \times d}, \quad t = g_t \odot h_t \in \mathbb{R}^{1 \times d}, \quad o'_t = o'_t \odot W_o + b_o \in \mathbb{R}^{1 \times d}.\]후보 hidden state의 계산을 간소화하기 위해, 이를 입력의 단순 선형 변환으로 유지합니다. 마지막으로, 남은 모든 가중치 행렬을 삼진 가중치 행렬로 대체하여 MatMul 작업을 완전히 제거합니다. 결과적으로 생성된 MLGRU 아키텍처는 다음과 같이 공식화될 수 있습니다.
\[f_t = \sigma(x_t \odot W_f + b_f) \in \mathbb{R}^{1 \times d}, \quad c_t = \tau(x_t \odot W_c + b_c) \in \mathbb{R}^{1 \times d}, \quad h_t = f_t \odot h_{t-1} + (1 - f_t) \odot c_t \in \mathbb{R}^{1 \times d}, \quad g_t = \sigma(x_t \odot W_g + b_g) \in \mathbb{R}^{1 \times d}, \quad t = g_t \odot h_t \in \mathbb{R}^{1 \times d}, \quad o'_t = o'_t \odot W_o + b_o \in \mathbb{R}^{1 \times d}.\]3.3.2 MatMul-free 채널 믹서
MatMul-free 채널 믹싱을 위해, GLU를 사용하며, 이는 많은 현대 LLM에서 널리 채택되고 있습니다. GLU는 세 주요 단계로 구성됩니다.
이를 통해 비싼 MatMul을 필요로 하는 대신, 정보를 채널 간에 효율적으로 혼합할 수 있습니다.
4. 실험
4.1 MatMul-free LM의 확장 법칙
신경망 확장 법칙은 훈련 세트 크기와 모델 크기의 함수로서 모델 오류가 감소한다고 가정합니다. 이런 예측은 큰 모델을 훈련하는 데 드는 비용이 증가함에 따라 중요해집니다. 먼저 작은 모델로 확장성을 테스트하는 것이 일반적인 최선의 방법입니다. GPT-4 기술 보고서에 따르면 최종 모델 크기의 1/10,000에 해당하는 예측 모델이 여전히 전체 크기 모델 성능을 정확하게 예측할 수 있습니다.
MatMul-free LM과 Transformer++ 모델의 370M, 1.3B, 2.7B 파라미터 모델에 대한 확장 법칙을 평가했습니다. Transformer++의 모든 가중치와 활성화는 BF16이고, MatMul-free LM의 BitLinear 층은 삼진 파라미터를 사용하며, BF16 활성화를 사용합니다. MatMul-free LM의 평균 연산은 Transformer++보다 계산적으로 저렴합니다.
MatMul-free LM의 확장 예측은 Transformer++보다 더 가파른 하강을 보여줍니다. 이는 MatMul-free LM이 추가적인 계산 자원을 더 효율적으로 활용하여 성능을 개선할 수 있음을 시사합니다. 이에 따라 MatMul-free LM의 확장 곡선은 약 1023 FLOPs에서 Transformer++의 확장 곡선과 교차할 것으로 예측됩니다. 이 계산 규모는 Llama-3 8B 모델과 Llama-2 70B 모델 훈련에 필요한 FLOPs와 대략적으로 동일합니다. 이 결과는 MatMul-free LM이 효율성뿐만 아니라 스케일업 시 손실 측면에서도 우수한 성능을 낼 수 있음을 시사합니다.
4.2 downstream 작업
BitNet 벤치마킹과 일치하게, 여러 언어 작업에서 이 모델들의 0-shot 성능을 평가했습니다. MatMul-free LM 모델은 모든 작업에서 Transformer++ 베이스라인과 경쟁력 있는 성능을 보여줍니다. 특히, 2.7B MatMul-free LM 모델은 ARC-Challenge와 OpenbookQA에서 Transformer++ 모델을 능가하며, 다른 작업에서는 비슷한 성능을 유지합니다. 모델 크기가 증가함에 따라, MatMul-free LM과 Transformer++ 사이의 성능 격차는 확장 법칙에 일치하게 좁혀집니다.
4.3 학습률 분석
학습률은 언어 모델 훈련에서 중요한 하이퍼파라미터이며, 삼진/이진 가중치 체제에서는 학습률에 더 민감합니다. 370M 모델을 사용하여 1.5e-3에서 3e-2 범위 내에서 최적의 학습률을 찾기 위한 탐색을 수행했습니다. 이 결과는 학습률이 1.5e-3에서 1e-2로 증가함에 따라 최종 훈련 손실이 단조롭게 감소하는 것을 보여줍니다. 학습률이 2e-2를 초과할 때만 모델이 불안정해집니다.
4.4 훈련 효율성 비교
제안된 Fused BitLinear과 Vanilla BitLinear 구현을 훈련 시간과 메모리 사용 측면에서 평가했습니다. 실험 결과, Fused 연산자는 더 큰 배치 크기에서 더 빠른 훈련 속도와 감소된 메모리 소비의 이점을 보여줍니다.
4.5 인퍼런스 효율성 비교
MatMul-free LM은 Transformer++에 비해 모든 모델 크기에서 일관되게 더 낮은 메모리 사용과 지연 시간을 보여줍니다. 가장 큰 모델 크기인 13B 파라미터에서 MatMul-free LM은 4.19 GB의 GPU 메모리만 사용하며, 695.48 ms의 지연 시간을 달성합니다. 반면 Transformer++는 48.50 GB의 메모리를 필요로 하고 3183.10 ms의 지연 시간을 보여줍니다. 이 결과는 MatMul-free LM이 큰 규모의 언어 모델링 작업에서 유망한 접근 방식임을 강조합니다.
5. FPGA 구현 및 결과
5.1 구현
삼진 연산을 더 잘 활용할 수 있는 맞춤 하드웨어에서 MatMul-free LM의 전력 사용 및 효과를 테스트하기 위해 SystemVerilog에서 FPGA 가속기를 만들었습니다. 이 디자인에는 네 개의 기능 단위가 포함되어 있으며, 각각은 간단한 아웃오브오더 실행을 허용합니다.
5.2 결과
MatMul-free 토큰 생성 코어의 RTL 구현은 Intel FPGA Devcloud의 D5005 Stratix 10 프로그래머블 가속 카드(PAC)에서 배포되었습니다. 이 코어는 43ms 동안 블록의 순방향 패스를 완료하며, 60MHz의 클록 속도를 달성합니다. 이 구현은 낮은 동적 전력을 보여주며, 코어가 비활성 상태일 때 측정된 전력과 구별하기 어렵습니다. 각 코어는 DDR4 인터페이스와 MMIO 브리지에 대한 액세스가 필요합니다. 이 구현에서는 주로 제공된 쉘 논리에 자원이 할당되며, 프로그래밍 가능한 논리 자원의 0.4%만이 DDR4 인터페이스/MMIO에 대한 코어 인터커넥트 및 중재 논리에 할당됩니다.
[참고자료 1] Swish Activation Function과 GLU(Gated Linear Unit)
Swish는 딥러닝에서 사용되는 활성화 함수로, Google Brain 팀이 “Searching for Activation Functions”라는 논문에서 처음 제안했습니다. 이 함수는 ReLU(Rectified Linear Unit)의 대안으로 고안되었으며, 일부 경우에서 ReLU보다 우수한 성능을 보이는 것으로 알려져 있습니다.
Swish 함수는 입력값에 자기 자신을 곱한 뒤, 시그모이드 함수를 통해 정규화하는 구조를 가집니다. 이는 입력 값의 범위를 자연스럽게 조절하며, 작은 입력 값에 대해서도 그래디언트가 사라지지 않는 이점을 제공합니다.
Swish 함수는 다음과 같은 수식으로 정의됩니다.
\[f(x) = x \cdot \sigma(\beta x)\]\(x\)는 입력, \(\sigma\)는 시그모이드 함수, \(\beta\)는 학습 가능한 파라미터 또는 고정된 상수(e.g., 1)입니다. 시그모이드 함수 \(\sigma(x)\)는 다음과 같습니다.
\[\sigma(x) = \frac{1}{1 + e^{-x}}\]예를 들어,
입력 \(x\)가 1일 때, \(\beta\)를 1로 설정하면, Swish 함수의 출력은 다음과 같습니다.
\[f(1) = 1 \cdot \frac{1}{1 + e^{-1}} \approx 0.731\]다른 활성화 함수와의 비교
Swish 외에도 여러 활성화 함수들이 널리 사용되며, 각 함수의 특징과 비교를 표로 정리하면 다음과 같습니다.
| 활성화 함수 | 수식 | 특징 |
|---|---|---|
| ReLU | \(f(x) = \max(0, x)\) | 음수 입력에 대해 0을 출력, 그래디언트 소실 문제를 해결하지만 죽은 뉴런 문제 발생 |
| Leaky ReLU | \(f(x) = \max(\alpha x, x)\) | 음수 입력에 대해 작은 기울기 \(\alpha\)를 부여하여 죽은 뉴런 문제 완화 |
| ELU | \(f(x) = \begin{cases} x & \text{if } x > 0 \\ \alpha(e^x - 1) & \text{if } x \leq 0 \end{cases}\) | 음수 값에 대해 포화를 방지하며, 일정 부분 그래디언트 소실 문제 해결 |
| Sigmoid | \(f(x) = \frac{1}{1 + e^{-x}}\) | 출력 범위가 [0,1], 그래디언트 소실 문제 발생 |
| Tanh | \(f(x) = \frac{e^x - e^{-x}}{e^x + e^{-x}}\) | 출력 범위가 [-1,1], 중심이 0이라는 점에서 학습 효율성 |
GLU (Gated Linear Unit)와의 관련성
GLU는 채널 믹싱을 위해 사용되며, 특히 LLM(Large Language Models)에서 활용됩니다. GLU는 입력 \(x_t\)를 두 개의 가중치 행렬 \(W_g\), \(W_u\)를 통해 변환하고, 그 결과를 게이팅하여 출력을 조절합니다. Swish 함수를 활용하여 비선형성을 추가함으로써, 모델의 표현력을 향상시킬 수 있습니다.
GLU의 수행 단계는 다음과 같습니다.
GLU는 복잡한 연산 없이도 효과적인 표현 학습을 가능하게 하는 구조적 접근 방식을 제공합니다. 이를 통해 대규모 언어모델에서 높은 성능을 달성할 수 있습니다.
[참고자료 2] Dead Neuron Problem
죽은 뉴런 문제(Dead Neuron Problem)는 특히 ReLU 활성화 함수에서 흔히 볼 수 있는 현상으로, 네트워크의 특정 뉴런이 항상 0을 출력하게 되는 상태를 말합니다. 이런 뉴런은 네트워크의 학습 과정에 참여하지 못하며, 그 결과로 네트워크의 전체적인 학습 능력이 저하될 수 있습니다.
원인 분석
ReLU 활성화 함수는 음수 입력에 대해 0을 출력합니다. 따라서, 입력값과 가중치의 조합이 음수가 되는 경우, 해당 뉴런의 출력은 계속 0이 됩니다. 만약 학습 과정에서 뉴런의 가중치가 잘못 조정되어 계속 음수 입력만을 받게 된다면, 그 뉴런은 학습에서 제외되는 죽은 뉴런이 됩니다.
죽은 뉴런 시뮬레이션(예시)
네트워크 설정
간단한 신경망을 가정하면 네트워크는 단일 입력 \(x\)를 받아 단일 출력 \(y\)를 내는 구조로, 가중치 \(w\)와 바이어스 \(b\)를 포함하고, 활성화 함수로 ReLU를 사용합니다. 수식으로 표현하면 다음과 같습니다.
\[y = \text{ReLU}(wx + b)\]가중치 \(w = -1\), 바이어스 \(b = 0\)로 설정하고, 입력 \(x\)가 항상 양수인 경우(e.g., \(x = 1\))를 가정하면,
이 조건에서 출력 \(y\)는 다음과 같이 계산됩니다.
\[y = \text{ReLU}(-1 \cdot 1 + 0) = \text{ReLU}(-1) = 0\]상기 식에서 \(y\)의 값은 항상 0이 됩니다. 이는 입력이 양수일 때, 가중치와 바이어스의 조합이 음수를 만들어내고, ReLU가 이를 0으로 변환하기 때문입니다.
문제 해결 방안
Leaky ReLU 도입
죽은 뉴런 문제를 해결하기 위해, Leaky ReLU와 같은 변형된 활성화 함수를 사용할 수 있습니다. Leaky ReLU는 음수 입력에 대해서도 아주 작은 양의 기울기(\(\alpha\))를 부여하여 출력이 완전히 0이 되지 않도록 합니다. 수식은 다음과 같습니다.
\[y = \text{LeakyReLU}(wx + b) = \text{LeakyReLU}(-1) = -0.01 \quad (\text{if } \alpha = 0.01)\]이 경우, 음수 입력에서도 출력이 \(-0.01\)과 같이 작은 값이 되어, 뉴런이 완전히 죽지 않고 gradient가 일부 전달될 수 있으므로 뉴런이 계속해서 학습 과정에 포함되고, 전체 네트워크의 학습 능력이 향상될 수 있습니다.