CATS: Contextually-Aware Thresholding for Sparsity in Large Language Models

📄 arXiv: 2404.08763v4 📥 PDF

作者: Donghyun Lee, Je-Yong Lee, Genghan Zhang, Mo Tiwari, Azalia Mirhoseini

分类: cs.LG, cs.CL

发布日期: 2024-04-12 (更新: 2024-11-03)

💡 一句话要点

提出CATS框架以降低大语言模型的推理成本

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 激活稀疏化 上下文感知 大语言模型 推理成本 非线性激活函数 GPU优化 模型性能 微调

📋 核心要点

  1. 现有方法在提高激活稀疏性时，往往导致下游任务性能显著下降，难以平衡计算成本与模型性能。
  2. 本文提出的CATS框架通过引入新的非线性激活函数，实现了激活的上下文感知稀疏化，降低推理成本。
  3. 实验结果表明，CATS模型在不微调的情况下，能保持与基础模型相近的性能，并在推理延迟上有显著改善。

📝 摘要（中文）

大型语言模型（LLMs）在人工智能应用中取得了显著进展，但由于其巨大的推理成本，部署仍然面临挑战。近期研究通过提高激活稀疏性来减轻LLMs的计算成本，但在下游任务上性能显著下降。本文提出了一种新的框架——上下文感知阈值稀疏化（CATS），旨在稀疏化基础LLMs的激活并降低推理成本。CATS实现简单、易于实施且效果显著。我们展示了CATS可应用于多种基础模型，包括Mistral-7B和Llama2-7B，并在下游任务性能上优于现有稀疏化技术。CATS模型在不进行微调的情况下，通常能在激活稀疏性达到50%时，保持与基础模型1-2%的性能差距。此外，CATS模型在微调时收敛更快，任务性能更佳。我们还开发了自定义GPU内核，以高效实现CATS，显著提升了推理延迟。

🔬 方法详解

问题定义：本文旨在解决大型语言模型在推理过程中面临的高计算成本问题。现有方法在提高激活稀疏性时，往往导致下游任务性能显著下降，难以实现性能与效率的平衡。

核心思路：CATS框架的核心思想是引入一种新的非线性激活函数，使得激活稀疏化能够根据上下文信息进行调整，从而在保持模型性能的同时降低推理成本。

技术框架：CATS的整体架构包括激活稀疏化模块和自定义GPU内核。激活稀疏化模块负责根据上下文信息动态调整激活值，而自定义GPU内核则优化了稀疏激活的计算效率。

关键创新：CATS的主要创新在于其上下文感知的激活稀疏化方法，这与传统的稀疏化技术不同，后者通常是静态的，无法根据输入变化进行调整。

关键设计：CATS采用了新的非线性激活函数，并在激活稀疏性达到50%时，仍能保持与基础模型1-2%的性能差距。此外，CATS的自定义GPU内核实现了约15%的推理延迟改善，提升了模型的实际应用效率。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，CATS模型在不进行微调的情况下，能够在激活稀疏性达到50%时，保持与基础模型1-2%的性能差距。此外，CATS的自定义GPU内核实现了约15%的推理延迟改善，显著提升了Llama-7B和Mistral-7B的推理效率。

🎯 应用场景

CATS框架具有广泛的应用潜力，尤其在需要高效推理的大型语言模型部署场景中，如实时对话系统、智能客服和内容生成等领域。其显著降低的推理延迟和保持的任务性能，使得在资源受限的环境中部署大型模型成为可能，推动了AI技术的普及与应用。

📄 摘要（原文）

Large Language Models (LLMs) have dramatically advanced AI applications, yet their deployment remains challenging due to their immense inference costs. Recent studies ameliorate the computational costs of LLMs by increasing their activation sparsity but suffer from significant performance degradation on downstream tasks. In this work, we introduce a new framework for sparsifying the activations of base LLMs and reducing inference costs, dubbed Contextually Aware Thresholding for Sparsity (CATS). CATS is relatively simple, easy to implement, and highly effective. At the heart of our framework is a new non-linear activation function. We demonstrate that CATS can be applied to various base models, including Mistral-7B and Llama2-7B, and outperforms existing sparsification techniques in downstream task performance. More precisely, CATS-based models often achieve downstream task performance within 1-2% of their base models without any fine-tuning and even at activation sparsity levels of 50%. Furthermore, CATS-based models converge faster and display better task performance than competing techniques when fine-tuning is applied. Finally, we develop a custom GPU kernel for efficient implementation of CATS that translates the activation of sparsity of CATS to real wall-clock time speedups. Our custom kernel implementation of CATS results in a ~15% improvement in wall-clock inference latency of token generation on both Llama-7B and Mistral-7B.