HiRE: High Recall Approximate Top-$k$ Estimation for Efficient LLM Inference

📄 arXiv: 2402.09360v1 📥 PDF

作者: Yashas Samaga B L, Varun Yerram, Chong You, Srinadh Bhojanapalli, Sanjiv Kumar, Prateek Jain, Praneeth Netrapalli

分类: cs.LG, cs.AI

发布日期: 2024-02-14

💡 一句话要点

提出HiRE以解决LLM推理中的内存瓶颈问题

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

关键词: 大型语言模型 推理优化 稀疏性利用 高召回估计 多设备计算

📋 核心要点

  1. 现有的自回归解码方法在LLM推理中面临内存瓶颈，导致模型参数传输效率低下。
  2. HiRE通过压缩方案和DA-TOP-$k$操作符，能够高效预测top-$k$行/列，从而减少内存传输和计算延迟。
  3. 在实验中，HiRE在保持准确性的同时，将推理延迟提升了1.47倍，展示了其在大规模模型上的有效性。

📝 摘要（中文）

自回归解码的生成大型语言模型（LLMs）在加速器（GPU/TPU）上常常受到内存限制，主要时间花费在将模型参数从高带宽内存（HBM）转移到缓存上。近期研究表明，通过适当训练模型以在前馈（FFN）层上操作top-$k$行/列（$k ext{约等于} 0.05$），LLMs可以在保持质量的同时显著减少模型参数的传输，从而降低延迟。然而，利用这种稀疏性来提高延迟的潜力受到限制，因为识别top行/列是数据依赖的，通常需要使用完整矩阵操作。为了解决这些问题，我们提出了HiRE（高召回近似top-$k$估计），它包括两个新颖的组件：一种压缩方案用于廉价预测高召回的top-$k$行/列，随后对预测的子集进行完整计算，以及DA-TOP-$k$：一种高效的多设备近似top-$k$操作符。我们展示了在一个十亿参数模型上，HiRE在softmax和前馈层的应用几乎达到了预训练和下游任务的准确性，并在单个TPUv5e设备上将推理延迟加速了1.47倍。

🔬 方法详解

问题定义：论文要解决的问题是大型语言模型在推理过程中由于内存传输造成的延迟，现有方法依赖于完整矩阵操作，限制了性能提升的潜力。

核心思路：论文的核心思路是通过HiRE框架，利用压缩技术和高召回的top-$k$估计，减少对完整矩阵的依赖，从而提高推理效率。

技术框架：HiRE的整体架构包括两个主要模块：首先是压缩方案，用于快速预测top-$k$行/列；其次是DA-TOP-$k$操作符，负责在多设备上高效执行近似top-$k$计算。

关键创新：HiRE的关键创新在于其压缩方案和DA-TOP-$k$操作符的结合，显著提高了对稀疏性利用的效率，与传统方法相比，减少了对完整矩阵操作的依赖。

关键设计：在设计中，HiRE采用了特定的参数设置以优化压缩效果，并在损失函数中引入了稀疏性约束，以确保在保持准确性的同时提高推理速度。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，在一个十亿参数的模型上，HiRE在softmax和前馈层的应用几乎达到了与预训练和下游任务相同的准确性，同时在单个TPUv5e设备上将推理延迟提升了1.47倍，展示了其显著的性能优势。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和大规模文本生成等场景，能够显著提升大型语言模型的推理效率，降低计算资源消耗，具有重要的实际价值和未来影响。

📄 摘要（原文）

Autoregressive decoding with generative Large Language Models (LLMs) on accelerators (GPUs/TPUs) is often memory-bound where most of the time is spent on transferring model parameters from high bandwidth memory (HBM) to cache. On the other hand, recent works show that LLMs can maintain quality with significant sparsity/redundancy in the feedforward (FFN) layers by appropriately training the model to operate on a top-$k$ fraction of rows/columns (where $k \approx 0.05$), there by suggesting a way to reduce the transfer of model parameters, and hence latency. However, exploiting this sparsity for improving latency is hindered by the fact that identifying top rows/columns is data-dependent and is usually performed using full matrix operations, severely limiting potential gains. To address these issues, we introduce HiRE (High Recall Approximate Top-k Estimation). HiRE comprises of two novel components: (i) a compression scheme to cheaply predict top-$k$ rows/columns with high recall, followed by full computation restricted to the predicted subset, and (ii) DA-TOP-$k$: an efficient multi-device approximate top-$k$ operator. We demonstrate that on a one billion parameter model, HiRE applied to both the softmax as well as feedforward layers, achieves almost matching pretraining and downstream accuracy, and speeds up inference latency by $1.47\times$ on a single TPUv5e device.