Tensor Product Attention Is All You Need

发表时间: 2025-01 · arXiv:2501.06425

原文: https://arxiv.org/abs/2501.06425

作者/机构：Yifan Zhang (清华大学/普林斯顿大学), Yifeng Liu (UCLA), Huizhuo Yuan (UCLA), Zhen Qin, Yang Yuan (清华大学/上海期智研究院), Quanquan Gu (UCLA), Andrew Chi-Chih Yao (清华大学/上海期智研究院)

速读

一句话结论 本文提出了一种名为张量积注意力（TPA）的新型注意力机制，通过上下文相关的动态张量分解大幅压缩键值缓存，在保持或提升大语言模型语言建模与下游任务性能的同时，将推理阶段的显存开销降低了一个数量级。

要解决什么问题 随着大语言模型处理的上下文长度不断增加，自回归推理过程中的键值缓存膨胀成为了严重的显存与内存带宽瓶颈。在标准的多头注意力机制中，模型需要为每一个生成的 token 完整缓存其对应的键张量和值张量。若注意力头数为 $h$，每头维度为 $d_h$，序列长度为 $T$，标准机制推理时的缓存总开销高达 $2 T h d_h$。当模型需要处理极长文本或在服务器端应对大批量并发请求时，这种与序列长度和头数呈线性正相关的全量张量存储会导致极高的显存占用。这不仅直接限制了单张显卡能够支持的最大上下文窗口长度，其庞大的内存读写需求也严重拖慢了解码生成的吞吐量与绝对速度。现有的静态权重分解无法动态适应输入，而传统的键值共享机制在极端压缩下往往会带来不可忽视的精度损失。因此，如何在不牺牲模型表达能力与生成质量的前提下，从底层计算机制上绕开全量张量实体化的存储限制，实现极高压缩率的动态缓存方案，是当前长序列大模型推理亟待解决的核心卡点。

怎么做的 张量积注意力机制的核心思路是放弃通过单一线性映射直接生成完整的头张量，转而利用上下文相关的张量分解技术，将查询、键和值的激活值动态拆解为低秩的潜在因子。它能绕开缓存膨胀卡点，是因为推理阶段不再存储庞大的完整键值张量，仅缓存降维后的因子组件，在数学层面上将每个 token 的存储成本从 $2 h d_h$ 降至 $(R_K + R_V)(h + d_h)$。当秩 $R_K$ 和 $R_V$ 设定为 1 或 2 时，即可实现数量级上的缓存缩减。

该方法由几个关键设计构成。首先是潜在因子映射，对于每一个输入 token 的隐藏状态，模型通过可学习的权重矩阵将其分别投影到头维度空间和特征维度空间，生成相互解耦的因子向量。其次是张量积重构，以键张量为例，其定义性公式被改写为秩为 $R_K$ 的张量积（外积）之和： $$ \mathbf{K}_{t}=\frac{1}{R_{K}} \sum_{r=1}^{R_{K}} \mathbf{a}_{r}^{K}\left(\mathbf{x}_{t}\right) \otimes \mathbf{b}_{r}^{K}\left(\mathbf{x}_{t}\right) $$ 其中 $\mathbf{a}_{r}^{K}$ 负责捕获头维度的分布，$\mathbf{b}_{r}^{K}$ 负责捕获特征维度的信息。这种设计不仅大幅减少了参数，还具有极强的理论包容性，标准的多头注意力、多查询注意力以及分组查询注意力，本质上都是该机制在特定秩和固定因子约束下的非上下文特例。

第三个关键部件是预旋转位置编码集成。该机制原生兼容旋转位置嵌入（RoPE），并提出了一种预旋转策略：直接将位置旋转操作作用于特征维度的因子上。这意味着每个键在进入缓存前就已经完成了位置信息的注入，推理阶段无需再进行显式的旋转计算，进一步加速了解码。最后，为了将理论压缩转化为实际的硬件加速，作者开发了 FlashTPA 解码算法。这是一种专为该机制优化的自回归推理算法，它利用因子化组件直接进行爱因斯坦求和约定计算，在数据流中融合了 Softmax 和值聚合步骤，完全避免了在显存中还原出完整的查询、键、值张量，从而极大地缓解了内存带宽压力。基于这一系列设计，作者构建了 T6 模型架构，将标准模型中的自注意力层直接替换为张量积注意力块，并配合 SwiGLU 前馈网络和 RMSNorm 归一化层进行端到端训练。

效果如何 实验基于 FineWeb-Edu 数据集，在 500 亿 tokens 上进行了单轮预训练。作者训练了 1.24 亿、3.53 亿、7.73 亿和 15 亿参数四个规模的模型，硬件环境使用了 4 到 8 张 A100 显卡。对比基线涵盖了当前主流的注意力路线：代表传统全量缓存路线的标准多头注意力（MHA）、代表键值共享压缩路线的多查询注意力（MQA）与分组查询注意力（GQA），以及代表潜在空间压缩路线的多头潜在注意力（MLA）。为保证公平，所有基线的注意力头数均经过调整以对齐参数总量。

量化结果表明，该方法在收敛速度和最终验证困惑度上均优于或持平于所有基线。在下游任务（包含 ARC、BoolQ、HellaSwag、MMLU 的零样本与少样本测试）中，3.53 亿参数的张量积注意力模型取得了 51.41% 的平均准确率，击败了 MHA 的 50.11% 和 MLA 的 50.13%。在 7.73 亿和 15 亿参数规模下，仅对键和值进行因子化的 TPA-KVonly 变体展现出最佳的平衡，分别以 53.52% 和 55.03% 的得分保持领先。在长序列解码速度测试中，当上下文长度超过 16K 时，FlashTPA 的解码耗时增长曲线显著平缓于 MQA 和 GQA，其绝对速度超越了 MLA，并大幅领先于 MHA。在局限性方面，作者承认当前的 FlashTPA 解码算法是基于 Triton 实现的，要彻底榨干其在极长序列下的硬件性能潜力，仍需在未来开发高度优化的 CUDA 内核。此外，附录中探索的三阶张量分解等高阶变体目前仅在小模型上验证了可行性，其在更大规模场景下的表现仍有待验证。

A1 主要贡献

随着大语言模型（LLM）处理上下文长度的增加，推理过程中的键值（KV）缓存成为严重的内存瓶颈。本文提出了张量积注意力（Tensor Product Attention, TPA），这是一种新型注意力机制，其核心创新点和贡献如下：

  1. 提出 TPA 机制：通过上下文相关的张量分解（Tensor Decompositions）来压缩查询（Q）、键（K）和值（V）的激活值。与静态权重分解（如 LoRA）不同，TPA 动态分解激活值，构建低秩且具上下文信息的表示。
  2. KV Cache 显著缩减：TPA 在推理阶段大幅减少了 KV Cache 的大小，通常比标准多头注意力（MHA）降低一个数量级，同时在预训练验证损失（困惑度）和下游任务表现上持平或更优。
  3. 高度兼容性：TPA 原生支持旋转位置嵌入（RoPE）及任何位置编码，也支持输出门控和 KV 移位。它可以作为 LLaMA、Qwen 和 Gemma 等现代 LLM 架构中 MHA 层的掉入式（Drop-in）替换方案。
  4. 提出 T6 模型架构：基于 TPA 构建了 Tensor ProducT ATTenTion Transformer (T6) 序列建模架构。
  5. 开发 FlashTPA Decoding 算法：这是一种针对 TPA 优化的自回归推理算法，利用因子化组件直接进行爱因斯坦求和（einsum）计算，避免显式还原全量张量，在长序列推理速度上超越了优化的 MHA、MQA、GQA 和 MLA。

Fig 1 Tensor ProducT ATTenTion Transformer (T6) 中的张量积注意力 (TPA)
Fig 1 展示了在 T6 每一层中，输入隐藏状态 $[x_t]$ 如何通过线性层产生 Q、K、V 的潜在因子矩阵，应用 RoPE 后通过张量积形成各头张量，最后进行缩放点积注意力计算。

A3 背景知识/关键观察/设计原则

符号系统与张量定义。本文使用粗体大写字母表示矩阵，粗体小写表示向量。令 $[d_{model}]$ 为嵌入维度，$[h]$ 为头数，$[d_h]$ 为每头维度。$[x_t \in \mathbb{R}^d]$ 表示第 $[t]$ 个 token 的输入。$[Q, K, V \in \mathbb{R}^{T \times h \times d_h}]$ 分别表示查询、键、值张量。向量 $[a, b]$ 的张量积（外积）定义为 $[a \otimes b = C]$，其中 $[C_{ij} = a_i b_j]$。矩阵 $[C]$ 的向量化记为 $[vec(C)]$。

缩放点积注意力与多头注意力 (MHA)。缩放点积注意力通过比较 $[Q]$ 和 $[K]$ 来确定关注重点，并对 $[V]$ 进行加权组合：

$$\operatorname{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V})=\operatorname{Softmax}\left(\frac{\mathbf{QK}^{\top}}{\sqrt{d_{k}}}\right) \mathbf{V},$$


MHA（【60, Attention is all you need+2017+NIPS+https://arxiv.org/abs/1706.03762】）通过将模型表示划分为多个头来扩展此机制。对于每 个 token 嵌入 $[x_t]$，计算每个头如下：

$$\mathbf{Q}_{t,i} = (\boldsymbol{W}_i^Q)^\top \mathbf{x}_t \in \mathbb{R}^{d_h}, \mathbf{K}_{t,i} = (\boldsymbol{W}_i^K)^\top \mathbf{x}_t \in \mathbb{R}^{d_h}, \mathbf{V}_{t,i} = (\boldsymbol{W}_i^V)^\top \mathbf{x}_t \in \mathbb{R}^{d_h},$$
$$\mathbf{head}_i = \text{Attention}(\mathbf{Q}_i, \mathbf{K}_i, \mathbf{V}_i),$$
其中 $[W_i^Q, W_i^K, W_i^V \in \mathbb{R}^{d_{model} \times d_h}]$ 是可学习的投影矩阵。最后连接所有头并映射回原始维度：
$$\operatorname{MHA}(\mathbf{X})=\operatorname{Concat}\left(\mathbf{head}_{1}, \ldots, \mathbf{head}_{h}\right) \boldsymbol{W}^{O} .$$

A2 方法细节

查询、键和值的张量因子化。TPA 不再通过单一线性映射形成头张量，而是将每个 $[Q_t, K_t, V_t]$ 因子化为秩分别为 $[R_q, R_k, R_v]$ 的（上下文相关）张量积之和。具体而言，对于每个 token $[t]$：

$$\mathbf{Q}_{t}=\frac{1}{R_{Q}} \sum_{r=1}^{R_{Q}} \mathbf{a}_{r}^{Q}\left(\mathbf{x}_{t}\right) \otimes \mathbf{b}_{r}^{Q}\left(\mathbf{x}_{t}\right), \quad \mathbf{K}_{t}=\frac{1}{R_{K}} \sum_{r=1}^{R_{K}} \mathbf{a}_{r}^{K}\left(\mathbf{x}_{t}\right) \otimes \mathbf{b}_{r}^{K}\left(\mathbf{x}_{t}\right),$$

$$\mathbf{V}_{t}=\frac{1}{R_{V}} \sum_{r=1}^{R_{V}} \mathbf{a}_{r}^{V}\left(\mathbf{x}_{t}\right) \otimes \mathbf{b}_{r}^{V}\left(\mathbf{x}_{t}\right),$$
其中 $[a_r^Q(x_t), a_r^K(x_t), a_r^V(x_t) \in \mathbb{R}^h]$，$[b_r^Q(x_t), b_r^K(x_t), b_r^V(x_t) \in \mathbb{R}^{d_h}]$。这意味着查询、键、值在头维度（Head dimension）和特征维度（Feature dimension）上被解耦分解。

潜在因子映射与秩索引合并。张量积中的每个因子都取决于 token 的隐藏状态 $[x_t]$。以查询为例：

$$\mathbf{a}_r^Q(\mathbf{x}_t) = \boldsymbol{W}_r^{a^Q} \mathbf{x}_t \in \mathbb{R}^h, \quad \mathbf{b}_r^Q(\mathbf{x}_t) = \boldsymbol{W}_r^{b^Q} \mathbf{x}_t \in \mathbb{R}^{d_h},$$


其中 $[W_r^{aQ} \in \mathbb{R}^{h \times d_{model}}]$ 和 $[W_r^{bQ} \in \mathbb{R}^{d_h \times d_{model}}]$ 是可学习的权重。为了实现效率，秩索引常被合并到单一输出维度：

$$\mathbf{a}^{Q}(\mathbf{x}_{t})=\boldsymbol{W}^{a^{Q}} \mathbf{x}_{t} \in \mathbb{R}^{R_{q} \cdot h}, \mathbf{b}^{Q}(\mathbf{x}_{t})=\boldsymbol{W}^{b^{Q}} \mathbf{x}_{t} \in \mathbb{R}^{R_{q} \cdot d_{h}}$$
随后重塑为 $[A_Q(x_t) \in \mathbb{R}^{R_Q \times h}]$ 和 $[B_Q(x_t) \in \mathbb{R}^{R_Q \times d_h}]$。则第 $[t]$ 个 token 的查询张量可表示为：
$$\mathbf{Q}_{t}=\frac{1}{R_{Q}} \mathbf{A}_{Q}\left(\mathbf{x}_{t}\right)^{\top} \mathbf{B}_{Q}\left(\mathbf{x}_{t}\right) \in \mathbb{R}^{h \times d_{h}}.$$
该操作等同于 $[Q_t = \frac{1}{R_Q} \sum_{r=1}^{R_Q} a_r^Q(x_t) (b_r^Q(x_t))^\top]$。

多头注意力计算与参数初始化。一旦完成因子化，多头注意力将按照标准 Transformer 进行。对于每个头 $[i]$：

$$\mathbf{head}_i = \text{Softmax} \left( \frac{1}{\sqrt{d_h}} \mathbf{Q}_i (\mathbf{K}_i)^\top \right) \mathbf{V}_i,$$


连接所有头后，通过输出权重 $[W^O]$ 投影回原始嵌入维度：

$$\text{TPA}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{Concat}(\mathbf{head}_1, \dots, \mathbf{head}_h) \boldsymbol{W}^O.$$
因子权重矩阵采用 Xavier 初始化（【15, Understanding the difficulty of training deep feedforward neural networks+2010+AISTATS】）。

RoPE 兼容性与加速优化。在标准流程中，RoPE 应用于重构后的张量。TPA 提供了一种优化策略：直接将 RoPE 集成到因子化中。
1. 因子预旋转：可以预先旋转 token 维度的因子，例如：

$$\widetilde{\mathbf{B}}_{K}(\mathbf{x}_{t}) \leftarrow \mathrm{RoPE}_{t}(\mathbf{B}_{K}(\mathbf{x}_{t})),$$


2. 预旋转表示：得到预旋转后的键表示：

$$\widetilde{\mathbf{K}}_t = \frac{1}{R_K} \sum_{r=1}^{R_K} \mathbf{a}_{(r)}^K(\mathbf{x}_t) \otimes \mathrm{RoPE}_t(\mathbf{b}_{(r)}^K(\mathbf{x}_t)) = \frac{1}{R_K} \mathbf{A}_K(\mathbf{x}_t)^\top \mathrm{RoPE}_t(\mathbf{B}_K(\mathbf{x}_t)).$$
这样每个键 $[K_t]$ 在缓存前就已经完成了旋转，推理阶段无需再显式计算旋转，从而加速解码。Theorem 3.1 证明了 TPA 因子化与 RoPE 是数学兼容的，保持了相对平移属性（详见附录 D.1 证明）。

KV 缓存与内存减少分析。标准注意力需要缓存 $[K_t, V_t \in \mathbb{R}^{h \times d_h}]$，总开销为 $[2 T h d_h]$。
- TPA 因子化 KV 缓存：TPA 仅存储因子组件 $[A_K(x_t), \tilde{B}_K(x_t), A_V(x_t), B_V(x_t)]$。
- 开销计算：每个 token 的存储成本为 $[(R_K + R_V)(h + d_h)]$。
- 压缩率：与标准缓存相比，比例为 $[\frac{(R_K + R_V)(h + d_h)}{2 h d_h}]$。对于典型的 $[d_h = 64]$ 或 $[128]$，当秩 $[R_K, R_V]$ 设为 1 或 2 时，KV Cache 可获得量级上的缩减。Table 1 详细对比了各机制的参数量与 KV 缓存大小。

Table 1 不同注意力机制对比
Table 1 不同注意力机制对比

MHA, MQA, GQA 作为非上下文 TPA 的表达。作者指出 MHA、MQA 和 GQA 实际上是 TPA 的特殊非上下文变体，即通过对 TPA 因子施加约束：
- MHA：等同于秩 $[R=h]$ 且头维度因子为固定正交的标准基向量 $[e_i]$ 的 TPA。
- MQA：等同于键值秩 $[R_K, R_V = 1]$ 且头维度因子为全 1 向量 $[1_h]$ 的 TPA。
- GQA：等同于秩为组数 $[G]$ 的 TPA，头维度因子为掩码向量。

T6 模型架构细节。T6 采用了 LLaMA（【58, Llama: Open and efficient foundation language models+2023+arXiv】）的整体架构，但将自注意力块替换为 TPA 块。FFN 部分采用 SwiGLU 层（【47, Glu variants improve transformer+2020+arXiv】）：

$$\begin{aligned} \begin{aligned} \mathbf{x} & \leftarrow \mathbf{x}+\operatorname{TPA}(\operatorname{RMSNorm}(\mathbf{x})), \\ \mathbf{x} & \leftarrow \mathbf{x}+\operatorname{SwiGLU-FFN}(\operatorname{RMSNorm}(\mathbf{x})). \end{aligned} \end{aligned}$$


在每个子层前应用层归一化（如 RMSNorm）。

A5 FlashTPA 解码算法

算法核心逻辑。为了高效生成 token，FlashTPA Decoding 利用因子化表示，通过一系列爱因斯坦求和（einsum）计算注意力输出，完全避免了完整 $Q, K, V$ 张量的实体化。
- 数据流过程：如图 2 所示。先计算头共享的特征空间点积 $[P]$，将其与头特定的 $[A_Q, A_K]$ 混合获得 logits，应用 Softmax，最后通过 $[A_V, B_V]$ 聚合值。
- 内存与带宽优势：随着序列长度 $[M]$ 增长，KV Cache 的因子化存储大幅降低了内存带宽压力。
- 实现细节：提供了基于 Triton 的内核实现（Algorithm 3），该实现融合了 Softmax 和值聚合步骤，具有良好的数值稳定性（在线 log-sum-exp）。

Fig 2 FlashTPA Decoding 数据流图
Fig 2 展示了单查询 (N=1) 与 M 个缓存项交互的过程，矩形代表张量，圆圈代表 einsum 缩并或元素级乘法。

A4 实验环境

  • 数据集：使用 FineWeb-Edu 100B 数据集进行预训练（1000 亿 token 训练，1 亿 token 验证）。
  • 模型规模：训练了四种规模：Small (124M)、Medium (353M)、Large (773M) 和 XL (1.5B)。所有模型均在 50B token 上训练 1 个 epoch。
  • 基准对比：对比了标准 LLaMA (MHA) 及其变体 MQA、GQA、MLA。为公平起见，调整头数以确保各注意力机制的参数量一致。
  • 软件配置：基于 nanoGPT 代码库实现。使用 AdamW 优化器，学习率采用余弦退火策略，设置了 2000 步预热。
  • 硬件配置：
    • Small: 4x A100 GPU。
    • Medium/Large/XL: 8x A100 GPU。
    • 具体的 batch size 和梯度累积步数见附录 Table 9。

A4 实验结果

语言建模性能。

  • 收敛速度与损失：在所有规模（Medium, Large, XL）中，TPA（红线）和 TPA-KVonly（粉线）的收敛速度均优于或持平于基准。如图 4 所示，TPA 在训练各阶段的验证损失均明显低于 MHA 和 MLA。
  • 验证困惑度：TPA-based 方法在预训练结束时达到了最低的困惑度（附录 Fig 9）。
    Fig 4 验证损失曲线对比

下游任务评测。
- 在 ARC, BoolQ, HellaSwag, MMLU 等主流基准上测试了 0-shot 和 2-shot 表现。
- Medium (353M)：TPA 平均准确率为 51.41%，高于 MHA (50.11%) 和 MLA (50.13%)。
- Large (773M)：TPA-KVonly 获得最高平均得分 53.52%。
- XL (1.5B)：TPA-KVonly 继续保持领先（55.03%）。
- 详细结果见 Table 2 和 Table 3。

Table 2 Medium 模型下游任务结果
Table 2 Medium 模型下游任务结果

FlashTPA 解码速度。
- 可扩展性：在长序列生成上表现卓越。随着序列长度增加，FlashTPA 的解码时间增长速度慢于 MQA/GQA。
- 性能对比：在序列长度超过 $[2^{14} \sim 2^{15}]$ 时，基于 Triton 的 FlashTPA 速度超越了 MLA，并显著快于 MHA（Fig 5）。
Fig 5 不同 batch size 下的解码时间对比

A5 结论

本文提出的 Tensor Product Attention (TPA) 通过上下文相关的低秩分解，有效地解决了 LLM 推理中的 KV Cache 膨胀问题。实验证明 TPA 在大幅压缩存储的同时，不仅没有损失模型性能，反而提升了语言建模的质量。其与 RoPE 的天然兼容性以及高效的 FlashTPA 解码算法，使其成为长文本建模极具竞争力的方案。
未来工作：作者指出目前的实现基于 Triton，未来开发高度优化的 CUDA 内核有望进一步挖掘其性能。此外，TPA 在其他模态（如图像、多模态）的泛化能力也值得深入研究。

A6 附录/补充细节

高阶张量积注意力 (Higher-Order TPA)。附录 C 探讨了将特征向量 $[b_r(x_t)]$ 进一步分解为 $[b_r(x_t) \otimes c_r(x_t)]$ 的可能性（3 阶分解）。这可以增强模型的表达能力，通过额外的上下文调制因子来调整特征。作者证明了高阶 TPA 同样兼容 RoPE（Theorem C.1），并在 Small 模型上初步验证了其可行性。

非上下文因子变体。
- Non-contextual B：固定 token 维度因子，仅让头维度因子 $[a(x_t)]$ 随输入变化。实验显示其表现略逊于标准 TPA。
- TPA-KVonly：仅对 K 和 V 进行因子化，Q 使用标准投影。这是一种在压缩率和复杂度之间的平衡方案。
- 共享 B 因子：让 K 和 V 共享 token 维度基向量 $[b_r]$，可进一步减少参数量和缓存开销。

复杂度分析 (Appendix A)。详细对比了显式还原（Naive）与特殊路径（Specialized）的 FLOPs。推导得出当秩 $[R]$ 较小时，TPA 的计算复杂度优势随头数 $[H]$ 或特征维度 $[D]$ 的增加而扩大。在超大规模模型（如 1~2T MoE）中，TPA 在投影成本和注意力成本上均具有显著优势。

汇总方法细节中的引用

编号 参考文献名 发表年份 会议/期刊 引用段落及原文描述
[60] Attention is all you need 2017 NIPS 背景章节：“Scaled dot-product attention [60] determines how to focus on different parts...”
[52] Roformer: Enhanced transformer... 2024 Neurocomputing 方法章节：“A key advantage of TPA is its native compatibility with rotary positional embeddings (RoPE) [52]...”
[15] Understanding the difficulty of training... 2010 AISTATS 方法章节：“We use Xavier initialization [15] for the factor weight matrices...”
[58] Llama: Open and efficient... 2023 arXiv 架构章节：“T6 utilizes the overall architecture of LLaMA [58] while changing the self-attention block...”
[47] Glu variants improve transformer 2020 arXiv 架构章节：“The feed-forward network (FFN) adopts a SwiGLU layer, as in [47, 58].”
[32] Deepseek-v2: A strong, economical... 2024 arXiv 引言/背景：“Multi-Head Latent Attention (MLA) in Deepseek-V2 [32] caches compressed KV representations...”
[57] Triton: An intermediate language... 2019 ACM 解码算法：“Our current FlashTPA implementation utilizes Triton [57].”
[24] NanoGPT 2022 GitHub 实验章节：“All experiments... are implemented based on the nanoGPT codebase [24].”