一种新型注意力机制 _ 张量积注意力（Tensor Product Attention, TPA）
论文提出了一种名为 Tensor Product Attention (TPA) 的新型注意力机制，旨在解决扩展语言模型输入序列长度时面临的内存开销问题。
TPA 通过将查询、键和值分解为潜在因子矩阵
A 和 B，并应用张量积和 RoPE 旋转，最终通过缩放点积注意力和线性投影得到输出的过程。
这种方法在保持模型性能的同时，显著减少了内存开销，提高了模型的可扩展性。
输入与线性变换
-
输入：图中最下方的
Linear层接收输入的隐藏状态（hidden state），并将其分别转换为查询（Query）、键（Key）和值（Value）的潜在因子矩阵A和B。 -
RoPE 应用：对
B_Q和B_K应用旋转位置编码（RoPE），以增强模型对位置信息的捕捉能力。RoPE 是一种通过旋转查询和键的向量来编码相对位置信息的方法。 -
张量积计算: - 查询、键和值的多头向量通过
A和B的张量积得到：
- 查询：A_Q 和 B_Q 的张量积。 - 键：A_K 和 B_K 的张量积。 - 值：A_V 和 B_V 的张量积。
- 张量积的结果经过缩放（Scale），缩放因子分别为 1/R_Q、1/R_K 和 1/R_V，其中 R_Q、R_K 和 R_V 是张量积的秩。
缩放点积注意力（ Scaled Dot-Product Attention）:
- 计算查询和键的点积，得到注意力权重。 - 将注意力权重与值相乘，得到加权后的输出。
输出与线性投影
-
Concat：将多个头的输出拼接在一起。
-
Linear：对拼接后的结果进行线性投影，得到最终的输出。
1. 问题背景
语言模型在处理长输入序列时，通常需要较大的键值（KV）缓存来存储中间结果，这会导致推理过程中内存开销显著增加。例如，标准的多头注意力（MHA）机制在处理长序列时，内存消耗会随着序列长度线性增长，限制了模型能够处理的最大上下文窗口。
2. 核心思想
TPA 通过张量分解技术，将查询（Q）、键（K）和值（V）表示为上下文相关的低秩张量，从而显著减少 KV 缓存的内存占用。具体来说：
1. 张量分解：TPA 将 Q、K 和 V 分解为多个低秩张量的和，每个张量由两个因子函数的外积组成。
2. 动态适应性：与静态的低秩分解方法（如 LoRA）不同，TPA 的分解是动态的，依赖于输入的上下文信息，能够更好地捕捉序列中的语义。
3. 内存优化：通过只存储分解后的低秩因子，TPA 可以将 KV 缓存的内存占用降低一个数量级。
2.1 传统 MHA 的问题
假设我们有一个语言模型，需要处理长度为 1024 的输入序列。在传统的多头注意力（MHA）机制中，每个头的维度为 64，总共有 8 个头。那么：
u 查询（Q）、键（K）和值（V）的维度分别为 1024×64。
u 每个头的 KV 缓存大小为 1024×64，总缓存大小为 1024×64×8。
这意味着在推理时，我们需要存储大量的 KV 缓存，内存开销非常大。如果序列长度增加到 2048 或更长，内存需求会进一步增长，这在硬件资源受限的情况下是不可持续的。
2.2 张量积注意力（TPA）核心思想
TPA 的核心在于通过张量分解减少 KV 缓存的内存占用，同时保持模型的性能。以下是关键点的详细解释和举例说明：
（1） 张量分解
TPA 将 Q、K 和 V 分解为低秩张量的和。假设我们将每个头的 Q、K 和 V 分解为两个低秩因子的外积，例如：
u 原始的 Q、K 和 V 的维度为 1024×64。
u 通过张量分解，可以将它们表示为两个低秩因子的乘积，例如 1024×8 和 8×64。
这种分解方式可以显著减少存储需求：
u 原始的存储需求为 1024×64=65,536。
u 分解后的存储需求为 1024×8+8×64=8,192+512=8,704。
通过这种分解，TPA 可以将 KV 缓存的内存占用减少到原来的约 1/7。
（2）动态上下文分解
TPA 的分解是动态的，依赖于输入的上下文信息。这意味着每个低秩因子的值会根据输入序列的变化而变化，从而更好地捕捉语义信息。例如：
l 在处理句子“猫在桌子上睡觉”时，TPA 会根据“猫”和“睡觉”的语义关系动态调整低秩因子的值。
l 在处理句子“狗在院子里奔跑”时，TPA 会根据“狗”和“奔跑”的语义关系动态调整低秩因子的值。
这种动态调整使得 TPA 能够灵活地捕捉不同上下文中的语义信息，而不会像静态分解方法（如 LoRA）那样固定不变。
（3）与 RoPE 的无缝集成
l TPA 与旋转位置编码（RoPE）兼容，可以直接应用于支持 RoPE 的模型架构。RoPE 是一种能够处理任意长度序列的位置编码方法，通过旋转向量来表示位置信息。例如：
l 在处理长序列时，RoPE 可以通过旋转向量来区分不同位置的词，而不会像传统位置编码那样在长序列中失效。
l TPA 与 RoPE 的结合使得模型能够更好地处理长序列，同时减少内存开销。
3. T6对比（ MLA 、 MHA、MQA、GQA ）
基于 TPA，论文提出了 T6，一种用于序列建模的新模型架构。T6 在语言建模任务上的表现优于现有的 Transformer 基线模型，包括 MHA、MQA、GQA 和 MLA。
TPA-KVonly 和 TPA：训练损\验证损失下降速度最快，且最终损失值最低，表明 TPA 及其变体在验证集上也表现最佳。
TPA-KVonly 和 TPA：验证困惑度下降速度最快，且最终困惑度最低，表明 TPA 及其变体在大规模模型上也表现最佳。
（1） 性能提升
在语言建模任务中，T6 的困惑度（Perplexity）明显低于传统 Transformer 模型。例如：
u 在一个包含 100 万句子的数据集上，传统 MHA 的训练后期稳定在约 15.0 左右 ，而 T6 的困惑度为 14.0 左右 。
u 这表明 T6 能够更好地预测下一个词，从而提升模型的性能。
（2） 内存效率
TPA 的内存优化使得 T6 能够在固定资源约束下处理更长的序列。例如：
u 在相同的硬件资源下，传统 MHA 可以处理的最大序列长度为 1024，而 TPA 可以处理的序列长度为 2048。
u 这使得 T6 能够更好地捕捉长距离依赖关系，从而提升模型的性能。
4. 实验结果
论文通过广泛的实验验证了 TPA 和 T6 的有效性：
u 在大规模预训练任务中，T6 在训练和验证损失上均优于传统注意力机制。
u 在多个标准基准测试（如 ARC、BoolQ、HellaSwag 等）中，T6 的零样本和少样本性能显著优于现有方法。
（1）零样本评估
在零样本评估中，T6 能够更好地理解未见过的句子。这表明 T6 能够更好地理解问题和上下文之间的关系。
（2） 少样本评估
在少样本评估中，T6 能够更快地适应新任务。这表明 T6 能够更好地利用有限的样本进行学习。
5 . 代码开源
TPA 和 T6 的代码已经在 GitHub 上开源，地址为：github.com/tensorgi/T6。
6 . 总结
TPA 通过张量分解显著减少了 KV 缓存的内存占用，同时通过动态上下文分解和与 RoPE 的无缝集成，保持了模型的性能。基于 TPA 的 T6 模型在语言建模任务中表现优异，同时在内存效率上具有显著优势。
通过这些优化，TPA 和 T6 为解决现代语言模型中的内存和可扩展性问题提供了新的思路和解决方案。
