Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time (2310.17157v1)

Abstract: LLMs with hundreds of billions of parameters have sparked a new wave of exciting AI applications. However, they are computationally expensive at inference time. Sparsity is a natural approach to reduce this cost, but existing methods either require costly retraining, have to forgo LLM's in-context learning ability, or do not yield wall-clock time speedup on modern hardware. We hypothesize that contextual sparsity, which are small, input-dependent sets of attention heads and MLP parameters that yield approximately the same output as the dense model for a given input, can address these issues. We show that contextual sparsity exists, that it can be accurately predicted, and that we can exploit it to speed up LLM inference in wall-clock time without compromising LLM's quality or in-context learning ability. Based on these insights, we propose DejaVu, a system that uses a low-cost algorithm to predict contextual sparsity on the fly given inputs to each layer, along with an asynchronous and hardware-aware implementation that speeds up LLM inference. We validate that DejaVu can reduce the inference latency of OPT-175B by over 2X compared to the state-of-the-art FasterTransformer, and over 6X compared to the widely used Hugging Face implementation, without compromising model quality. The code is available at https://github.com/FMInference/DejaVu.

• The paper introduces contextual sparsity by identifying input-dependent parameter subsets that maintain output quality while reducing computation.
• It employs a two-layer network for low-cost prediction and an asynchronous look-ahead mechanism that leverages residual connections for parallel processing.
• Experimental results show up to 85% parameter deactivation and halved inference latency on models like OPT-175B without compromising performance.

Contextual Sparsity for Efficient LLM Inference

The paper "Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time" addresses the significant computational challenges posed by LLMs during inference. While LLMs have shown remarkable capabilities, their extensive parameter sizes result in high computational costs. This paper explores the concept of contextual sparsity as a viable solution to this issue.

Key Contributions and Methodology

The research introduces the concept of contextual sparsity, which involves determining small, input-dependent subsets of attention heads and MLP parameters that can replicate the outputs of the dense model for particular inputs. This approach contrasts with previous methods that often required extensive retraining or resulted in negligible speedups on contemporary hardware.

Existence and Exploitation of Contextual Sparsity:

The paper provides evidence for the existence of substantial contextual sparsity within LLMs. By conducting forward passes through the model, the authors demonstrate that up to 85% of model parameters can be turned off for specific inputs without degrading performance. This opens the door for potentially realizing a $7\times$ parameter reduction.

Prediction of Contextual Sparsity:

The research proposes a low-cost, learning-based algorithm to predict the subsets of parameters needed, using a two-layer fully connected network. The ability to foresee these subsets effectively enables dynamic sparsity, preserving the model's quality and ensuring significant speedups.

Asynchronous Sparse Prediction:

To mitigate any overhead introduced by the sparsity prediction process, the authors propose an asynchronous look-ahead mechanism. This system exploits the properties of residual connections in LLMs, noting that token embeddings change gradually across layers—a phenomenon allowing predictions to be made one layer ahead, thus parallelizing predictions with ongoing computations.

Hardware-Efficient Implementation:

The development of hardware-aware implementations, including kernel fusion and memory coalescing, capitalizes on the block-oriented nature of modern GPUs. This efficiently reduces memory I/Os, a critical bottleneck during inference, resulting in substantial end-to-end latency improvements.

Empirical Validation

The proposed system, dejavu, is validated across several benchmarks. The authors report that it can halve the inference latency of large models like OPT-175B compared to the best existing implementations, with no compromise in model performance. The framework also extends its applicability beyond single-token generation to larger batch scenarios and demonstrates compatibility with quantization techniques.

Theoretical and Practical Implications

The findings highlight the potential for contextual sparsity to substantially lower inference costs, making it a compelling approach for deploying LLMs in latency-sensitive environments. The implications of this work extend to a broader understanding of model efficiency, offering a pragmatic way to accommodate the growing size of models without proportional increases in computational resources.

Future Developments

Looking forward, the insights from this paper suggest possible avenues for further refining the prediction mechanisms, potentially incorporating them into training phases to dynamically adjust model architectures. Moreover, the compatibility with quantization indicates additional layers of optimization that could be explored.

By providing a sophisticated blend of theoretical insights and practical implementations, this research outlines a promising pathway for efficient LLM deployment, aligning computational feasibility with the burgeoning capabilities of LLMs.