This paper presents a novel approach to large language models (LLMs) that eliminates matrix multiplication (MatMul) operations, which are typically the most computationally expensive part of such models. By doing so, the authors aim to significantly reduce memory usage and improve computational efficiency, enabling the models to scale up to billions of parameters while maintaining performance comparable to state-of-the-art Transformers.
-
MatMul-Free Dense Layers: The core innovation lies in replacing MatMul operations in dense layers with addition operations using ternary weights. These ternary weights take values from {-1, 0, +1}, which allows matrix multiplications to be transformed into simple additions and subtractions.
-
MatMul-Free Self-Attention: Instead of using the traditional MatMul-based self-attention mechanism, the authors propose an alternative that employs element-wise Hadamard products and gated recurrent units (GRUs). This approach maintains the effectiveness of the self-attention mechanism without the computational overhead of MatMul.
-
Hardware-Efficient Implementations: The paper describes both a GPU-efficient implementation and a custom FPGA accelerator. The GPU implementation uses fused kernels to reduce memory usage and improve training speed. The FPGA implementation further optimizes ternary operations, showcasing the potential for even greater efficiency gains.
-
Scaling Laws and Performance: The study demonstrates that as model size increases, the performance gap between MatMul-free models and full-precision Transformers narrows. This suggests that the proposed approach becomes more competitive at larger scales.
In conventional dense layers, the input (X) and weight matrix (W) undergo a MatMul operation: [ Y = XW ] To avoid MatMul, the authors adopt BitLinear modules with ternary weights: [ W \in {-1, 0, +1} ] This ternary weight matrix transforms the MatMul into an accumulation of additions and subtractions.
The self-attention mechanism typically involves three MatMul operations for the query (Q), key (K), and value (V) matrices. The paper proposes a MatMul-free variant using a modified GRU (Gated Recurrent Unit), termed the MatMul-free Linear Gated Recurrent Unit (MLGRU), which replaces these operations with element-wise products and additions: [ \begin{align*} r_t &= \sigma(W_r x_t + U_r h_{t-1} + b_r) \ z_t &= \sigma(W_z x_t + U_z h_{t-1} + b_z) \ \hat{h}t &= \tanh(W_h x_t + r_t \circ U_h h{t-1} + b_h) \ h_t &= (1 - z_t) \circ h_{t-1} + z_t \circ \hat{h}_t \end{align*} ] Where ( r_t ) and ( z_t ) are the reset and update gates, respectively, and (\circ) denotes element-wise multiplication.
GPU Implementation:
- Uses fused kernels to combine operations and reduce memory accesses.
- Ternary weights allow for optimized CUDA kernels, increasing inference speed and reducing memory usage.
FPGA Implementation:
- Custom instruction set and assembler designed for ternary operations.
- Functional units include row-wise operations, root mean square normalization, and ternary matrix multiplication.
- Efficient use of FPGA resources leads to significant power savings and speed improvements.
- The fused BitLinear layer implementation showed a 25.6% speedup in training and a 61% reduction in memory usage compared to the vanilla implementation.
- During inference, the MatMul-free models demonstrated up to 10x reduction in memory consumption and 4.57x increase in speed compared to unoptimized models.
- The MatMul-free models achieved competitive performance across various language tasks, including question answering and commonsense reasoning.
- The performance gap between the MatMul-free models and traditional Transformers decreased as the model size increased, highlighting the scalability of the proposed approach.
The paper introduces a significant shift in the design of large language models by eliminating MatMul operations, leading to more memory-efficient and computationally efficient models. This approach not only maintains competitive performance but also points to a future where lightweight, scalable models can be more easily deployed on various hardware platforms, including GPUs and FPGAs. The work serves as a call to prioritize the development of such efficient models, especially as the scale and deployment of LLMs continue to grow.