PaperlayerView exact revision on arXiv

Source-linked AI summary

Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time

Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivastava, Ce Zhang, Yuandong Tian, Christopher Re, Beidi Chen

arXiv:2310.17157v1cs.LG

TL;DR

LLM inference is costly, and existing sparsity methods may require retraining, weaken in-context learning, or fail to provide wall-clock speedups. The paper introduces DejaVu, which predicts contextual sparsity online and exploits it with asynchronous, hardware-aware execution. DejaVu reduces OPT-175B latency by over 2× versus FasterTransformer without quality degradation.

  • Problem

    Existing LLM sparsity methods face retraining, quality, in-context learning, or modern-hardware wall-clock efficiency limitations.

  • Method

    DejaVu uses low-cost predictors to identify input-dependent attention heads and MLP parameters, with asynchronous lookahead to exploit contextual sparsity during inference.

  • Results

    Over 2× end-to-end latency reduction was achieved for OPT-175B versus FasterTransformer without quality degradation.

  • Takeaways & Limitations

    Contextual sparsity can be predicted and exploited to reduce LLM inference latency while retaining model quality and in-context learning ability.

  • Takeaways & Limitations

    The paper hypothesizes that contextual sparsity exists for any input, and also hypothesizes that attention heads perform mean-shift clustering.

Abstract

from arXiv · show

Large language models (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.

1 Introduction

DejaVu addresses the cost and limitations of inference-time sparsity by exploiting input-dependent contextual sparsity and predicting it efficiently during execution. Its asynchronous, hardware-aware system preserves model quality while substantially reducing OPT-175B latency.

  • Motivation: LLMs are expensive at inference time, while existing sparsity methods face retraining, quality, in-context learning, or hardware-speedup limitations.The paper seeks less computation and memory without sacrificing pretrained LLM capabilities.
  • Contextual Sparsity: Contextual sparsity selects small, input-dependent sets of attention heads and MLP parameters that approximately preserve the dense model’s output.This approach avoids modifying pretrained models during inference.
  • Contextual Sparsity: 85% structured sparsity can yield essentially the same output while potentially reducing parameters 7× for each input.The sparsity is structured across attention heads and MLP parameters and is intended to maintain accuracy.
  • Prediction: Accurate sparsity prediction requires contextual token embeddings and can be formulated using similarity between layer parameters and the previous layer’s output.Purely dynamic information is insufficient for accurate prediction.
  • System: DejaVu uses low-cost learning-based predictors, asynchronous lookahead, and hardware-aware sparse matrix multiplication to exploit contextual sparsity during inference.Predictors select relevant heads or MLP parameters for subsequent computation while reducing sequential overhead.
  • Results: Over 2× end-to-end latency reduction was achieved for OPT-175B versus FasterTransformer without quality degradation, with further gains versus Hugging Face at small batch sizes.The comparison is against FasterTransformer and the widely used Hugging Face implementation.

2 Related Work and Problem Formulation

The paper motivates contextual sparsity from LLM inference bottlenecks, formalizes sparse attention and MLP computation, and frames prediction as selecting structured subsets that approximate dense outputs. Token generation, parameter movement, and inter-GPU communication define the efficiency constraints.

  • 2.2 LLM Inference Latency Breakdown: LLM inference has prompting and token-generation phases, with token generation often dominating latency because of parameter-loading I/O.The cited setting uses tensor model parallelism and emphasizes latency-sensitive generation.
  • 2.2 LLM Inference Latency Breakdown: Inter-GPU communication accounts for around 15% of token-generation latency, limiting the maximum speedup from skipping transformer computation to around 6×.This boundary is specific to the tensor-parallel regime described.
  • 2.3 Problem Formulation: The problem formulation seeks attention-head and MLP-neuron subsets that minimize sparse-versus-full computation error under a compute budget.The sparse subsets are denoted S_A for attention heads and S_M for MLP neurons.
  • 2.3 Problem Formulation: A sparsified MLP computes only a small set of neurons, and sparsity in the first linear layer also sparsifies the second through sparse activations.The MLP uses two linear layers and an activation function such as ReLU or GeLU.
  • 2.3 Problem Formulation: Sparsified attention retains a small set of heads whose combined output approximately matches full attention for the current input.The formulation represents token embeddings, the current MHA input, projection matrices, and the selected head set.

3 Pre-trained LLMs are Contextually Sparse

Experiments verify that pre-trained LLMs contain substantial input-dependent sparsity in attention heads and MLP neurons. Attention behaves like token clustering, while residual connections help explain slowly changing embeddings and enable similarity-based prediction.

  • 3.1 Contextual Sparsity Hypothesis: Two forward passes identify input-specific heads and neurons with large output norms, then evaluate the model using only those recorded parameters.The verification covers OPT-175B, OPT-66B, and OPT-30B on downstream datasets including OpenBookQA and Wiki-Text.
  • 3.1 Contextual Sparsity Hypothesis: Up to 80% of attention heads and 95% of MLP neurons can be sparse on average, yielding about 85% total sparsity and potentially 7× speedup.These are structured sparsity patterns that vary across input examples.
  • 3.2 Token Clustering in Attention Layers: Attention heads differ in token selectivity: heavy-hitter heads concentrate on particular tokens, whereas uniform heads can be omitted without affecting the example prediction.The example associates heavy attention with tokens such as “like” and “shipping.”
  • 3.2 Token Clustering in Attention Layers: The authors hypothesize that each attention head performs a mean-shift clustering step in a projection space, with different heads learning different spaces.This accounts for token clustering and head-specific attention patterns.
  • 3.3 Slowly Changing Embeddings across Layers: These observations motivate similarity-based sparsity prediction for DejaVu.The analysis connects the natural occurrence of contextual sparsity with the system’s prediction design.
  • 3.3 Slowly Changing Embeddings across Layers: Consecutive-layer embeddings remain highly similar: OPT models exceed 95% similarity, with OPT-175B around 0.99 from the second layer onward.Residual connections explain this behavior because the input norm is significantly larger than the transformation norm.
  • 3.3 Slowly Changing Embeddings across Layers: High contextual sparsity may contribute to small transformation norms because many MLP outputs and attention-head outputs have small norms.The paper connects this sparsity-related norm behavior to slowly changing embeddings.

4 DEJAVU

DejaVu predicts input-dependent sparsity for MLPs and attention heads, then uses slowly changing embeddings, asynchronous execution, and hardware-aware kernels to reduce inference latency. Its design addresses prediction accuracy and overhead, which otherwise can outweigh the savings from sparse computation.

  • 4.1 Contextual Sparsity Prediction in MLP Blocks: DejaVu frames MLP sparsity prediction as approximate maximum inner-product search between input embeddings and neuron parameters.The MLP’s first-layer parameters form the dataset and the input embedding forms the query.
  • 4.1 Contextual Sparsity Prediction in MLP Blocks: A neural-network classifier replaces slower GPU near-neighbor searches because HNSW takes over 10 ms and FAISS over 4 ms, versus 0.2 ms for OPT-175B MLP computation.The design targets prediction overhead that would otherwise exceed the sparse computation savings.
  • 4.2 Contextual Sparsity Prediction in Attention Blocks: Attention-head prediction uses the same classifier architecture as MLP prediction, treating each head as a class and selecting a set of heads for computation.After the first few layers, prediction can use the current token embedding and similarity to head parameters.
  • 4.3 Reducing Overhead with Asynchronous Execution: Look-ahead prediction and residual connections allow sparse decisions to be made across layers while relaxing otherwise sequential computation.The informal guarantee states that sufficiently small embedding changes preserve the quality of the MaxIP decision up to a changed approximation factor.
  • 4.3 Reducing Overhead with Asynchronous Execution: Prediction overhead can outweigh savings because attention and MLP computation must otherwise wait for sparse-predictor decisions.DejaVu parallelizes prediction with block computation to reduce this dependency.
  • 4.4 Hardware-efficient Implementation: Kernel fusion and memory coalescing make the implementation hardware-efficient, achieving up to 2× end-to-end speedup over FasterTransformer.The implementation accounts for GPU memory-I/O bottlenecks and block-oriented memory access.

5 Empirical Evaluation

DejaVu maintains model quality and in-context learning while exploiting contextual sparsity for substantial inference speedups across models, blocks, and batch sizes.

  • 5.1 End-to-End Results: 75% sparsity causes no average accuracy drop across zero-shot tasks, with a similar trend in five-shot evaluation.The five-shot result verifies preservation of in-context learning ability.
  • 5.1 End-to-End Results: At around 75% sparsity, DEJAVU speeds up OPT-175B token generation by 1.8-2× versus FasterTransformers and 4.8-6× versus Hugging Face.The comparison uses batch size 1.
  • 5.2 Ablation Results: Union contextual sparsity does not grow linearly with batch size, suggesting that similar inputs could be batched to obtain higher sparsity.The union operation is used to realize a fast sparse GEMM.
  • 5.2 Ablation Results: 85% MLP sparsity introduces no accuracy loss across zero-shot tasks and language modeling when the Attention block remains dense.The MLP sparse predictor achieves over 99% validation accuracy in shallow layers and around 93% in ending layers.
  • 5.2 Ablation Results: Attention sparsity introduces no accuracy loss at around 50% sparsity when the MLP block remains dense.Validation accuracy is around 93% in middle layers and near 99% in shallow and deep layers.
  • 5.2 Ablation Results: DEJAVU reports no accuracy loss at 50% sparsity on OPT-66B and at attention sparsity 50% plus MLP sparsity 30% on BLOOM.The lower BLOOM MLP sparsity is attributed to a difference in activation function.
  • 5.2 Ablation Results: Contextual prediction preserves accuracy, whereas non-contextual prediction using only the original input embedding causes accuracy losses even at 50% sparsity.This supports using the activation at every layer as the predictor input.
  • 5.2 Ablation Results: DEJAVU combined with W4A16 quantization almost always achieves better accuracy than either DEJAVU or quantization alone.The DEJAVU-OPT-175B model uses 75% sparsity in this comparison.

6 Conclusion

The conclusion presents DEJAVU as a practical approach to efficient LLM inference that retains model quality and in-context learning while reducing latency. It emphasizes lightweight prediction, asynchronous lookahead, and hardware-efficient sparsity as the basis for the reported speedup.

  • 6 Conclusion: DEJAVU targets efficient LLM inference so pretrained models’ in-context learning abilities can be used in more application domains.The conclusion frames inference efficiency as the central practical goal.
  • 6 Conclusion: DEJAVU uses lightweight learning-based prediction, asynchronous lookahead predictors, and hardware-efficient sparsity to speed up inference in wall-clock time.The system predicts contextual sparsity on the fly rather than modifying the pretrained model.
  • 6 Conclusion: DEJAVU reduces OPT-175B inference latency by over 2× versus FasterTransformer without model quality drops.The conclusion characterizes the empirical results as encouraging.
  • A Related Work: The paper situates DEJAVU among specialized LLM inference systems and identifies a lack of careful algorithm-and-system co-design for hardware efficiency.Related systems include Faster Transformer, Orca, LightSeq, PaLM inference, TurboTransformers, and DeepSpeed-Inference.
  • A Related Work: Near-neighbor search is a prior technique used in recommendation, question answering, and natural language processing, providing context for DEJAVU’s prediction approach.The related-work discussion also covers quantization, pruning, distillation, and residual connections.

B Additional Observation on Slowly Changing Observation

Residual connections make representations change slowly across transformer layers, helping explain why contextual sparsity can be predicted from preceding-layer outputs. This behavior also affects how sparsity can be exploited for larger batches.

  • Slowly Changing Observation: Cosine similarity between activations across adjacent layers is higher for larger OPT models.Figure 9 compares layer l with layer l+1 across models.
  • Slowly Changing Observation: Residual connections compute X+F(X), keeping X close to the layer output because ∥X∥ is generally much larger than ∥F(X)∥.The first layer is an exception, where ∥F(X)∥ is larger and cosine similarity is lower.
  • Slowly Changing Observation: Layer normalization scales ∥X∥ to a consistent magnitude across layers, with example scales of 85 for OPT-30B and 110 for OPT-175B.The overall norm trend is similar across models of different sizes.
  • Larger-Batch Observation: Union Contextual Sparsity measures the fraction of MLP neurons or Attention heads unused by every input in a batch.It is calculated as 1.0 minus the union of activated units divided by the total units; the union enables fast sparse GEMM.
  • Larger-Batch Observation: Activated neurons and heads do not grow linearly with batch size, suggesting a power-law rather than uniform parameter-access distribution.Larger batches can also cause out-of-memory issues for long sequences because of GPU memory, model size, and KV-cache storage.

C.2 Near Neighbor classifier

DejaVu formulates sparsity prediction as near-neighbor search and uses a neural classifier to reduce prediction cost on GPUs. The appendix contrasts this choice with HNSW and notes additional sparsification possibilities.

  • C.2 Near Neighbor Classifier: Any near-neighbor search method under the inner product metric can predict a sparsity pattern in the DEJAVU framework.Training the predictor reduces on-the-fly prediction cost rather than training the language model.
  • C.2 Near Neighbor Classifier: HNSW produced no perplexity drop at 90% sparsity during exploration, but its 10 ms CPU prediction time exceeded MLP computation time.The overhead arose from high-dimensional embeddings and HNSW’s reliance on the CPU.
  • C.2 Near Neighbor Classifier: DejaVu chooses a neural network classifier to exploit fast GPU matrix multiplication and lower training cost.The classifier serves as the near-neighbor search method for predicting sparsity patterns.
  • C.2 Near Neighbor Classifier: Union contextual sparsity is used to support fast sparse GEMM when combining inputs in a batch.The associated figures examine union contextual sparsity for larger batch sizes.
  • C.2 Near Neighbor Classifier: Skipping or parallelizing entire transformer blocks may avoid catastrophic test-time accuracy drops, according to the cited additional observation.This observation concerns depth-wise sparsification rather than the classifier itself.

C.3 Future Possibility: Skipping Layer

The paper explores reducing inference cost by exploiting slowly changing activations across Transformer blocks, enabling parallelization, reordering, or skipping layers while preserving accuracy.

  • Motivation: Adjacent-layer activations often have cosine similarity above 0.99, suggesting that consecutive blocks can reuse similar representations.This observation was reported across seven OPT model sizes using C4 validation activations.
  • Results: Preliminary results show that parallelizing sequentially trained blocks does not significantly hurt downstream-task performance.The paper reports this finding for OPT-175B and Bloom in Table C.3.
  • Method: Parallelizing attention and MLP blocks across two or four Transformer layers replaces sequential dependencies with shared-input computation.The supplied equations define two-block and four-block parallelization using the same input activation for the relevant sub-blocks.
  • Results: Randomly skipping 25% of layers does not lead to catastrophic quality, indicating a possible direction for model compression and optimization.The authors connect this result to relatively consistent activation patterns across blocks.

D Implementation Details

DejaVu implements sparse attention and MLP computation through predicted indices, fused sparse kernels, and memory layouts designed for efficient GPU access.

  • Sparse computation: Figure 12 depicts predicted head and neuron indices selecting sparse attention and MLP computations for an input.The example selects attention indices 0,3 and MLP indices 0,2.
  • Kernel optimizations: Kernel fusion combines sparse indexing with matrix multiplication to avoid extra memory reads and writes.Separate indexing can incur three times the memory I/O, whereas fusion performs selection and multiplication together.
  • Kernel optimizations: Column-major storage for selected-column access improves memory coalescing without adding generation-time cost.The required transposition is performed once when loading the model.

E Benchmarking Sparse MLP and Sparse Attention

Hardware-aware sparse MLP and attention implementations deliver wall-clock speedups over PyTorch baselines and remain faster than dense computation across substantial densities.

  • MLP: Up to 4.5× faster than the baseline implementation in PyTorch, the sparse MLP implementation remains faster than dense MLP for density up to 0.8.The benchmark uses OPT-175B on 8xA100s.
  • Benchmark setup: The implementation validates wall-clock speedup against both dense computation and standard PyTorch implementations.The supplied passage identifies this as the purpose of the hardware-aware sparse MLP and attention evaluation.
  • Benchmark results: 4-5× faster than the baseline implementation in Pytorch, sparse MLP and attention remain faster than dense versions for density up to 0.8.The comparison covers both sparse MLP and sparse attention implementations.
  • Attention: Up to 5× faster than the baseline implementation in PyTorch, the sparse attention implementation remains faster than dense MLP for density up to 0.8.The benchmark uses OPT-175B on 8xA100s.

G.1 Soft-Max Functions

This section defines softmax variants, ReLU, subspace embeddings, and supporting random-matrix tools used in the paper’s theoretical analysis.

  • Softmax functions: The standard softmax is defined as an ℓ1-based function, while the paper also considers an ℓ2 generalization.Both functions map vectors in Rs to vectors in Rs.
  • Softmax functions: For the ℓ2 softmax construction, the resulting function approximately preserves vector norms on a rank-k subspace with probability 1−δ.The guarantee is bounded by factors of (1−ϵ)τ and (1+ϵ)τ under the stated dimension condition.
  • ReLU functions: The paper defines ReLU as ϕ(z)=max{z,0} and uses it in norm-preservation guarantees for transformed vectors.The stated bound compares the transformed output norm with the norm of ϕ(Kx).
  • Subspace embeddings: An ℓ2 subspace embedding preserves norms simultaneously for vectors in a matrix column space up to a multiplicative 1±ϵ factor with probability 1−δ.The section introduces standard sketching matrices and supporting net and concentration arguments for these guarantees.

J Nearest Neighbor Search Data Structure

This section defines approximate nearest-neighbor and MaxIP search, then reduces direction search to projected approximate MaxIP through transformations to unit-sphere ANN.

  • Approximate nearest-neighbor search: Approximate ANN returns a vector within c·r when the query has a dataset neighbor within distance r.
  • Approximate nearest-neighbor search: LSH solves unit-sphere ANN with query time O(d·nρ), preprocessing O(dn1+o(1)), and space O(n1+o(1)+dn).
  • MaxIP formulation: Approximate MaxIP retrieves z whose inner product is at least c times the dataset maximum when that maximum exceeds τ.
  • MaxIP formulation: Projected MaxIP applies separate query and data transforms before comparing inner products in a lower-dimensional space.
  • Reduction to MaxIP: The proposed transformations map direction search in convex optimization to MaxIP on the unit sphere while preserving the relevant argmax or approximation relation.
  • Reduction to MaxIP: Convexity ensures the transformed direction-search instance has a non-negative maximum inner product, satisfying the positivity condition required by LSH-based MaxIP.

J.4 Data Structures

This section gives LSH-based projected MaxIP data structures and develops an explanation of how self-attention and training induce contextual clustering in embeddings.

  • J.4 Data Structures: LSH solves unit-sphere approximate MaxIP with probability at least 0.97 using O(d·nρ) query time, O(dn1+o(1)) preprocessing, and O(n1+o(1)+dn) space.
  • J.4 Data Structures: Projected approximate MaxIP adds transform costs Tϕ and Tψ to the query and preprocessing complexities, respectively.
  • Self-attention and clustering: Self-attention resembles mean-shift clustering because attention aggregates value projections according to similarity between queries and keys.
  • Embedding capacity: A fixed low-dimensional embedding generally cannot satisfy all m2 pairwise distance constraints with only md embedding parameters.
  • Embedding capacity: Self-attention addresses this limitation by trading model size for additional computation and grouping embeddings through a multi-layer structure.
  • Training-induced clustering: When x⊺y > −0.4576, the negative gradient pushes normalized embeddings toward one another, producing static embedding clustering during training.
  • Training-induced clustering: For multiple embeddings, when pS > 1/2, the gradient has a positive component toward the weighted mean ¯y, moving x toward that mean.
  • Self-attention and generative models: Equal weighting of embeddings around xi cancels their contributions, leaving xk aligned with xi.
Loading 2310.17157v1…