One Pass Streaming Algorithm for Super Long Token Attention Approximation in Sublinear Space

📄 arXiv: 2311.14652v2 📥 PDF

作者: Raghav Addanki, Chenyang Li, Zhao Song, Chiwun Yang

分类: cs.LG, cs.CL, stat.ML

发布日期: 2023-11-24 (更新: 2024-02-05)

💡 一句话要点

提出一种单遍流式算法以解决超长令牌注意力近似问题

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

关键词: 流式算法 超长令牌 注意力机制 内存效率 大型语言模型 数据流处理 机器学习

📋 核心要点

  1. 现有方法在处理超长上下文时，时间和空间复杂度均为$O(n^2)$，导致内存使用过高。
  2. 本文提出一种单遍流式算法，利用亚线性空间存储草图矩阵，从而提高内存效率。
  3. 实验结果表明，随着令牌长度的增加，算法的误差保证减小，内存使用几乎保持不变。

📝 摘要（中文）

注意力计算的时间复杂度和空间复杂度均为$O(n^2)$，这使得在需要长上下文的流式应用中部署大型语言模型（LLMs）面临巨大的计算资源挑战。本文关注当上下文长度$n$远大于128K时的内存效率问题。我们提出了一种新算法，仅需一次数据读取，采用亚线性空间$o(n)$存储三个草图矩阵，避免了对精确$K, V$的存储需求。该算法在处理超长令牌时表现出色，随着令牌长度的增加，误差保证减小，而内存使用几乎保持不变，显示出在流式应用中高效处理LLMs的潜力。

🔬 方法详解

问题定义：本文旨在解决在流式应用中处理超长上下文时，现有注意力计算方法的高时间和空间复杂度问题。现有方法在处理长上下文时，内存需求过高，限制了其应用。

核心思路：我们提出的算法通过一次性读取数据，利用亚线性空间存储草图矩阵，避免了对完整$K, V$矩阵的存储需求，从而提高了内存效率。

技术框架：该算法的整体架构包括三个主要模块：数据流读取、草图矩阵构建和注意力计算。首先，算法从数据流中读取信息，然后构建草图矩阵，最后通过草图矩阵进行注意力计算。

关键创新：本研究的关键创新在于提出了一种仅需一次数据读取的流式算法，显著降低了内存使用，同时保持了计算精度。这与传统方法的高内存需求形成鲜明对比。

关键设计：算法中关键的参数设置包括草图矩阵的维度选择和误差控制策略。通过合理设计这些参数，确保了在处理超长令牌时的高效性和准确性。

🖼️ 关键图片

fig_0
img_1

📊 实验亮点

实验结果显示，提出的算法在处理超长令牌时，内存使用保持在亚线性水平，同时误差保证随着令牌长度的增加而减小。与传统方法相比，内存需求显著降低，展示了在流式应用中的优越性能。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、实时数据分析和流媒体服务等。通过提高大型语言模型在流式应用中的内存效率，该算法能够支持更长上下文的处理，推动智能助手、自动翻译和内容生成等领域的发展，具有重要的实际价值和未来影响。

📄 摘要（原文）

Attention computation takes both the time complexity of $O(n^2)$ and the space complexity of $O(n^2)$ simultaneously, which makes deploying Large Language Models (LLMs) in streaming applications that involve long contexts requiring substantial computational resources. In recent OpenAI DevDay (Nov 6, 2023), OpenAI released a new model that is able to support a 128K-long document, in our paper, we focus on the memory-efficient issue when context length $n$ is much greater than 128K ($n \gg 2^d$). Considering a single-layer self-attention with Query, Key, and Value matrices $Q, K, V \in \mathbb{R}^{n \times d}$, the polynomial method approximates the attention output $T \in \mathbb{R}^{n \times d}$. It accomplishes this by constructing $U_1, U_2 \in \mathbb{R}^{n \times t}$ to expedite attention ${\sf Attn}(Q, K, V)$ computation within $n^{1+o(1)}$ time executions. Despite this, computing the approximated attention matrix $U_1U_2^\top \in \mathbb{R}^{n \times n}$ still necessitates $O(n^2)$ space, leading to significant memory usage. In response to these challenges, we introduce a new algorithm that only reads one pass of the data in a streaming fashion. This method employs sublinear space $o(n)$ to store three sketch matrices, alleviating the need for exact $K, V$ storage. Notably, our algorithm exhibits exceptional memory-efficient performance with super-long tokens. As the token length $n$ increases, our error guarantee diminishes while the memory usage remains nearly constant. This unique attribute underscores the potential of our technique in efficiently handling LLMs in streaming applications.