1. 引言

在现代人工智能领域,多头自注意力机制(Multi-Headed Self Attention, MHSA)可以说是最重要的架构范式之一。它是Transformer模型的核心组件,而Transformer又是当前最先进的大型语言模型的基础架构。本文将深入浅出地解析多头自注意力机制的工作原理,通过手动计算的方式,让读者对其内部运作有一个直观而全面的理解。

2. 多头自注意力机制的背景

在深入MHSA之前,我们先简要回顾一下自然语言处理(NLP)领域的相关发展历程。早期的NLP模型主要依赖于循环神经网络(RNN)和长短期记忆网络(LSTM)等序列模型。这些模型虽然能够处理序列数据,但在处理长序列时存在长期依赖问题。

2017年,Google提出了Transformer模型,其核心就是多头自注意力机制。MHSA能够并行处理输入序列,捕捉序列中的长距离依赖关系,大大提高了模型的性能和效率。自此,MHSA成为了NLP领域的主流技术,被广泛应用于各种大型语言模型中。

3. 多头自注意力机制的工作原理

让我们通过一个具体的例子,step by step地计算多头自注意力机制的输出。我们将遵循以下步骤:

1. 定义输入
2. 定义可学习参数
3. 计算查询(Query)、键(Key)和值(Value)
4. 划分多个注意力头
5. 计算Z矩阵
6. 掩码操作
7. 计算注意力矩阵
8. 计算注意力头的输出
9. 拼接多个注意力头的输出

3.1 定义输入

MHSA可以应用于各种类型的数据,但通常情况下,输入是一个向量序列。在自然语言处理中,这通常是词嵌入(word embedding)与位置编码(positional encoding)的组合。

假设我们有一个简单的输入序列,包含3个词,每个词用4维向量表示:

Input = [
    [1, 2, 3, 4],
    [5, 6, 7, 8],
    [9, 10, 11, 12]
]

这个4×3的矩阵代表了我们的输入序列。

3.2 定义可学习参数

MHSA主要学习三个权重矩阵,用于构造”查询”(Query)、”键”(Key)和”值”(Value)。在本例中,我们假设模型已经学习到了以下权重矩阵:

W_Q = [
    [0.1, 0.2],
    [0.3, 0.4],
    [0.5, 0.6],
    [0.7, 0.8]
]

W_K = [
    [0.1, 0.2],
    [0.3, 0.4],
    [0.5, 0.6],
    [0.7, 0.8]
]

W_V = [
    [0.1, 0.2],
    [0.3, 0.4],
    [0.5, 0.6],
    [0.7, 0.8]
]

这些4×2的矩阵代表了模型的可学习参数。

3.3 计算查询、键和值

接下来,我们将输入与权重矩阵相乘,得到查询、键和值:

Query = Input * W_Q
Key = Input * W_K
Value = Input * W_V

让我们计算Query:

Query = [
    [1*0.1 + 2*0.3 + 3*0.5 + 4*0.7, 1*0.2 + 2*0.4 + 3*0.6 + 4*0.8],
    [5*0.1 + 6*0.3 + 7*0.5 + 8*0.7, 5*0.2 + 6*0.4 + 7*0.6 + 8*0.8],
    [9*0.1 + 10*0.3 + 11*0.5 + 12*0.7, 9*0.2 + 10*0.4 + 11*0.6 + 12*0.8]
]

Query = [
    [5.0, 6.0],
    [13.0, 15.0],
    [21.0, 24.0]
]

同理可得Key和Value:

Key = [
    [5.0, 6.0],
    [13.0, 15.0],
    [21.0, 24.0]
]

Value = [
    [5.0, 6.0],
    [13.0, 15.0],
    [21.0, 24.0]
]

3.4 划分多个注意力头

多头自注意力机制的”多头”体现在这一步。我们将Query、Key和Value划分为多个子矩阵,每个子矩阵对应一个注意力头。在本例中,我们使用两个注意力头:

Query_1 = [
    [5.0],
    [13.0],
    [21.0]
]

Query_2 = [
    [6.0],
    [15.0],
    [24.0]
]

Key_1 = [
    [5.0],
    [13.0],
    [21.0]
]

Key_2 = [
    [6.0],
    [15.0],
    [24.0]
]

Value_1 = [
    [5.0],
    [13.0],
    [21.0]
]

Value_2 = [
    [6.0],
    [15.0],
    [24.0]
]

这样,我们就得到了两组Query、Key和Value,分别用于两个注意力头的计算。

3.5 计算Z矩阵

接下来,我们需要计算Z矩阵,这是构造注意力矩阵的中间步骤。Z矩阵由Query和Key的矩阵乘法得到。我们以第一个注意力头为例:

Z_1 = Query_1 * Key_1^T

Z_1 = [
    [5.0],    [5.0, 13.0, 21.0]
    [13.0], *
    [21.0]
]

Z_1 = [
    [5.0*5.0, 5.0*13.0, 5.0*21.0],
    [13.0*5.0, 13.0*13.0, 13.0*21.0],
    [21.0*5.0, 21.0*13.0, 21.0*21.0]
]

Z_1 = [
    [25, 65, 105],
    [65, 169, 273],
    [105, 273, 441]
]

为了防止Z矩阵的值随着序列长度的增加而过大,我们通常会将Z矩阵除以序列长度的平方根。在本例中,序列长度为3,所以我们将Z_1除以:

Z_1 = [
    [14.43, 37.53, 60.62],
    [37.53, 97.58, 157.62],
    [60.62, 157.62, 254.62]
]

同理可得Z_2。

3.6 掩码操作

在某些应用场景中,如语言模型预测下一个词时,我们需要进行掩码操作,以确保模型在预测时不会”看到”未来的信息。这通常通过将Z矩阵中的某些位置设置为负无穷来实现。在本例中,我们假设不需要掩码操作。

3.7 计算注意力矩阵

注意力矩阵是通过对Z矩阵的每一行进行softmax运算得到的。softmax函数的定义如下:

让我们以Z_1的第一行为例计算softmax:

row = [14.43, 37.53, 60.62]
e_row = [1,850,752, 20,215,372,293, 1,800,537,936,918]
sum_e_row = 1,822,604,060,963

softmax(row) = [
    1,850,752 / 1,822,604,060,963,
    20,215,372,293 / 1,822,604,060,963,
    1,800,537,936,918 / 1,822,604,060,963
]

softmax(row) ≈ [0.001, 0.011, 0.988]

对Z_1的每一行都进行这样的计算,我们就得到了注意力矩阵Attention_1:

Attention_1 ≈ [
    [0.001, 0.011, 0.988],
    [0.000, 0.000, 1.000],
    [0.000, 0.000, 1.000]
]

同理可得Attention_2。

3.8 计算注意力头的输出

得到注意力矩阵后,我们将其与Value相乘,得到每个注意力头的输出:

Output_1 = Attention_1 * Value_1

Output_1 ≈ [
    [0.001*5.0 + 0.011*13.0 + 0.988*21.0],
    [0.000*5.0 + 0.000*13.0 + 1.000*21.0],
    [0.000*5.0 + 0.000*13.0 + 1.000*21.0]
]

Output_1 ≈ [
    [20.86],
    [21.00],
    [21.00]
]

同理可得Output_2。

3.9 拼接多个注意力头的输出

最后,我们将所有注意力头的输出拼接起来,得到多头自注意力机制的最终输出:

Final_Output = [Output_1 | Output_2]

Final_Output ≈ [
    [20.86, 24.00],
    [21.00, 24.00],
    [21.00, 24.00]
]

这个3×2的矩阵就是多头自注意力机制的输出结果。

4. 多头自注意力机制的优势

通过上述计算过程,我们可以看出多头自注意力机制具有以下优势:

1. 并行计算: MHSA可以并行处理输入序列中的所有元素,大大提高了计算效率。
2. 捕捉多种关系: 通过使用多个注意力头,模型可以同时关注输入序列中的不同特征和关系。
3. 长距离依赖: MHSA可以有效捕捉序列中的长距离依赖关系,克服了RNN等传统模型的局限性。
4. 灵活性: MHSA可以应用于各种类型的序列数据,不仅限于自然语言处理。
5. 可解释性: 注意力权重可以提供模型决策过程的一定解释性,有助于理解模型的工作原理。

5. 多头自注意力机制的应用

MHSA在自然语言处理领域有广泛的应用,包括但不限于:

1. 机器翻译: Transformer模型在机器翻译任务中取得了突破性的成果。
2. 文本生成: GPT系列模型使用了基于MHSA的架构,能够生成高质量的文本。
3. 问答系统: BERT等模型在问答任务中表现出色,为智能问答系统提供了强大的支持。
4. 文本分类: MHSA可以有效捕捉文本的语义特征,提高分类准确率。
5. 语音识别: 在语音识别任务中,MHSA也展现出了优秀的性能。
6. 图像处理: 虽然最初设计用于NLP任务,但MHSA也被成功应用于计算机视觉领域。

6. 结论

多头自注意力机制是现代人工智能,特别是自然语言处理领域的核心技术之一。通过本文的详细计算过程,我们深入了解了MHSA的工作原理。尽管实际应用中的计算规模要大得多,但基本原理是相同的。

理解MHSA的工作原理对于深入学习和应用先进的AI技术至关重要。随着技术的不断发展,我们可以期待MHSA在更多领域发挥重要作用,推动人工智能技术的进步。

参考文献

1. Warfield, D. (2024). Multi-Headed Self Attention — By Hand. Intuitively and Exhaustively Explained. https://iaee.substack.com/p/multi-headed-self-attention-by-hand
2. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., … & Polosukhin, I. (2017). Attention is all you need. Advances in neural information processing systems, 30.

随着人工智能技术的快速发展，大语言模型(Large Language Models, LLMs)在自然语言处理领域扮演着越来越重要的角色。然而，这些模型通常规模庞大、计算密集，给部署和推理带来了巨大挑战。为了解决这一问题，英特尔推出了OpenVINO GenAI Flavor，这是一个专门针对生成式AI模型优化的推理引擎。本文将详细介绍如何使用OpenVINO GenAI Flavor来高效运行LLMs，帮助开发者充分发挥硬件性能，实现快速、高效的模型推理。

OpenVINO GenAI Flavor概述

OpenVINO GenAI Flavor是OpenVINO工具套件的一个专门版本，旨在优化生成式AI模型的推理性能。它集成了多项先进技术，如动态形状支持、稀疏计算和高效内存管理等，特别适合处理LLMs这类大规模、复杂的模型。

主要特点

1. 专为LLMs优化：针对Transformer架构和生成式任务进行了特殊优化。
2. 动态形状支持：能够处理变长输入序列，无需固定批处理大小。
3. 高效内存管理：通过智能缓存和内存复用技术，显著减少内存占用。
4. 稀疏计算加速：利用模型的稀疏性，提高计算效率。
5. 多硬件支持：可在CPU、GPU等多种硬件平台上运行，充分利用硬件特性。

安装和设置

要开始使用OpenVINO GenAI Flavor，首先需要安装必要的软件包。您可以通过pip命令轻松完成安装：

pip install openvino openvino-genai

这将安装最新的OpenVINO开发版本以及GenAI Flavor专用组件。

模型准备

在使用OpenVINO GenAI Flavor之前，需要将LLM转换为OpenVINO的中间表示（IR）格式。这一步骤可以通过OpenVINO的模型转换工具完成。以下是转换过程的基本步骤：

1. 导出原始模型：从训练框架（如PyTorch或TensorFlow）导出模型。
2. 转换为ONNX：将模型转换为ONNX格式，这是一个通用的深度学习模型表示格式。
3. ONNX到IR转换：使用OpenVINO的Model Optimizer工具将ONNX模型转换为IR格式。

示例代码：

from openvino.runtime import Core
from transformers import AutoTokenizer, AutoModelForCausalLM

# 加载预训练模型和分词器
model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# 转换为ONNX格式
onnx_model_path = "gpt2.onnx"
dummy_input = tokenizer("Hello, how are you?", return_tensors="pt").input_ids
torch.onnx.export(model, dummy_input, onnx_model_path, opset_version=11)

# 使用OpenVINO转换为IR格式
core = Core()
ov_model = core.read_model(onnx_model_path)
compiled_model = core.compile_model(ov_model, "CPU")

使用OpenVINO GenAI Flavor进行推理

一旦模型转换完成，就可以使用OpenVINO GenAI Flavor进行高效推理。以下是一个基本的推理流程示例：

import numpy as np
from openvino.runtime import Core, Tensor
from transformers import AutoTokenizer

# 初始化OpenVINO Core和模型
core = Core()
model = core.read_model("path/to/your/model.xml")
compiled_model = core.compile_model(model, "CPU")

# 准备输入数据
tokenizer = AutoTokenizer.from_pretrained("gpt2")
input_text = "OpenVINO is"
input_ids = tokenizer.encode(input_text, return_tensors="np")

# 创建推理请求
infer_request = compiled_model.create_infer_request()

# 设置输入并执行推理
infer_request.set_input_tensor(Tensor(input_ids))
infer_request.infer()

# 获取输出
output = infer_request.get_output_tensor().data

# 解码输出
generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
print(generated_text)

这个示例展示了如何使用OpenVINO GenAI Flavor加载模型、处理输入、执行推理并获取输出。

高级功能和优化技巧

OpenVINO GenAI Flavor提供了多种高级功能和优化技巧，可以进一步提升LLMs的推理性能：

1. 动态形状处理

GenAI Flavor支持动态输入形状，这对于处理变长序列非常有用：

# 设置动态形状
model.reshape({0: [-1, -1]})  # 第一维为批次大小，第二维为序列长度

2. KV缓存优化

对于自回归生成任务，GenAI Flavor实现了高效的KV缓存机制：

# 启用KV缓存
compiled_model = core.compile_model(model, "CPU", config={"PERFORMANCE_HINT": "LATENCY"})

3. 稀疏计算加速

GenAI Flavor可以自动检测和利用模型中的稀疏性，无需额外配置即可获得加速：

# 稀疏计算会自动应用，无需特殊设置

4. 批处理推理

对于需要处理多个输入的场景，可以使用批处理来提高吞吐量：

# 准备批处理输入
batch_inputs = tokenizer(["Hello", "How are you", "OpenVINO is great"], 
                         padding=True, return_tensors="np")

# 执行批处理推理
results = compiled_model(batch_inputs.input_ids)

5. 异步推理

利用异步推理可以更好地利用硬件资源，提高整体效率：

# 创建异步推理请求
infer_request = compiled_model.create_infer_request()

# 启动异步推理
infer_request.start_async()

# 等待结果
infer_request.wait()

# 获取结果
output = infer_request.get_output_tensor().data

性能优化和调优

为了获得最佳性能，可以考虑以下优化策略：

1. 选择合适的硬件：根据模型大小和推理需求，选择CPU、GPU或专用AI加速器。
2. 量化：对模型进行INT8量化可以显著减少内存占用和推理时间。
3. 模型剪枝：移除不必要的模型参数，减小模型大小。
4. 缓存优化：合理设置缓存大小，平衡内存使用和性能。
5. 并行推理：在多核系统上，利用多线程并行处理多个推理请求。

示例代码：

# 使用INT8量化
quantized_model = core.quantize_model(model, "CPU", {"STAT_TYPE": "DYNAMIC"})

# 设置线程数
core.set_property("CPU", {"INFERENCE_NUM_THREADS": 4})

# 启用缓存
compiled_model = core.compile_model(quantized_model, "CPU", 
                                    config={"CACHE_DIR": "./model_cache"})

最佳实践和注意事项

1. 内存管理：对于大型LLMs，合理管理内存至关重要。使用流式处理或分段处理来减少内存占用。
2. 输入预处理：确保输入数据格式正确，并考虑将预处理步骤集成到模型中以提高效率。
3. 错误处理：实现健壮的错误处理机制，以应对可能的推理失败或异常情况。
4. 模型更新：定期更新模型和OpenVINO版本，以获得最新的性能优化和功能支持。
5. 性能监控：使用OpenVINO提供的性能分析工具来识别瓶颈并进行针对性优化。

结论

OpenVINO GenAI Flavor为运行大语言模型提供了强大而灵活的解决方案。通过利用其专门的优化技术和高级功能，开发者可以显著提升LLMs的推理性能，使这些复杂的模型能够在各种硬件平台上高效运行。随着生成式AI技术的不断发展，OpenVINO GenAI Flavor将继续演进，为开发者提供更多工具和能力，以应对未来的挑战和机遇。

参考文献

1. Intel Corporation. (2024). Run LLMs with OpenVINO GenAI Flavor — OpenVINO™ documentation. https://docs.openvino.ai/2024/learn-openvino/llm_inference_guide/genai-guide.html

大型语言模型 (LLM) 正在改变世界，从生成文本到翻译语言，再到编写代码，LLM 的应用范围越来越广。然而，LLM 的推理速度一直是制约其应用的关键因素。为了解决这个问题，研究人员一直在探索各种方法来加速 LLM 推理。

本文将介绍一篇名为 “FlashDecoding++: Faster Large Language Model Inference on GPUs” 的论文，该论文提出了一种新的 LLM 推理加速技术，可以在 GPU 上显著提高推理速度。

LLM 推理加速的挑战

加速 LLM 推理面临着三大挑战：

1. 同步部分 Softmax 更新： Softmax 操作需要对每个部分 Softmax 结果进行同步更新，这导致了 LLM 中注意力计算的约 20% 的开销。
2. 扁平 GEMM 的计算利用率低下： LLM 推理中执行 GEMM 的矩阵形状是扁平的，导致计算利用率低下，在之前的设计中，填充零后会导致超过 50% 的性能损失。
3. 静态数据流导致的性能损失： LLM 中的内核性能取决于不同的输入数据特征、硬件配置等。单一且静态的数据流会导致 LLM 推理中不同形状的 GEMM 出现 50.25% 的性能损失。

FlashDecoding++ 的解决方案

FlashDecoding++ 针对上述挑战提出了以下解决方案：

1. 异步 Softmax 与统一最大值： FlashDecoding++ 引入了一种统一最大值技术，用于不同的部分 Softmax 计算，从而避免同步。
2. 双缓冲的扁平 GEMM 优化： FlashDecoding++ 指出不同形状的扁平 GEMM 面临着不同的瓶颈。然后，引入了双缓冲等技术。
3. 硬件资源自适应的启发式数据流： FlashDecoding++ 使用不同的硬件资源，考虑输入动态，启发式地优化数据流。

性能提升

FlashDecoding++ 的优化策略使其在 NVIDIA 和 AMD GPU 上都取得了显著的性能提升，与 Hugging Face 实现相比，分别实现了高达 4.86 倍和 2.18 倍的加速。与主流 LLM 上最先进的 LLM 推理引擎相比，FlashDecoding++ 的平均加速比为 1.37 倍。

总结

FlashDecoding++ 提出了一套全面的 LLM 推理加速解决方案，有效地解决了 LLM 推理中的三大挑战。其在主流 LLM 和硬件平台上的出色表现，为 LLM 的广泛应用提供了强有力的支持。

参考文献

• [2311.01282] FlashDecoding++: Faster Large Language Model Inference on GPUs
• Stanford CRFM
• GitHub – opengear-project/GEAR: GEAR: An Efficient KV Cache Compression Recipefor Near-Lossless Generative Inference of LLM
• DistServe/distserve at main · LLMServe/DistServe · GitHub

注： 本文仅对 FlashDecoding++ 论文进行了简要介绍，更多细节请参考原文。

---

在GPU上推理大规模语言模型（LLM）的性能至关重要，而FlashDecoding++是一款针对LLM推理的快速引擎，通过解决同步部分softmax更新、未充分利用的扁平化GEMM计算和静态数据流等挑战，实现了显着的推理加速效果。

解决同步部分softmax更新的挑战：
FlashDecoding++引入了异步softmax和统一最大值的技术，避免了在计算部分softmax结果时需要同步更新的问题。每个部分softmax结果可以独立进行处理，无需进行同步操作，从而减少了计算中的开销。

解决未充分利用的扁平化GEMM计算的挑战：
FlashDecoding++通过双缓冲技术对扁平化GEMM计算进行了优化，隐藏了内存访问延迟，提高了计算利用率。它在共享内存中分配了两个独立的缓冲区，一个缓冲区用于进行GEMM计算，而另一个缓冲区则用于加载下一个GEMM操作所需的数据。通过这种方式，计算和内存访问可以同时进行，实现了计算与存储的重叠。

解决静态数据流的挑战：
FlashDecoding++采用了启发式数据流和硬件资源适应的方法。它根据输入动态和硬件配置，在不同的线性工作负载下动态优化数据流，选择最佳的实现方式。通过根据不同工作负载的特点进行灵活调整，FlashDecoding++实现了最佳的推理性能。

性能评估：
FlashDecoding++在多个硬件平台上进行了性能评估，包括NVIDIA和AMD的GPU。与Hugging Face、vLLM、DeepSpeed、TensorRT-LLM、OpenPPL和FlashDecoding等LLM推理引擎进行了比较。结果表明，FlashDecoding++相对于这些基线引擎实现了显着的加速效果，提供了高达4.86倍的推理速度提升。

---

近年来，大型语言模型（LLM）在各个领域都取得了显著的进步，其强大的文本理解和生成能力令人惊叹。那么，LLM 是否能胜任模拟文本世界的任务呢？换句话说，能否利用 LLM 来构建虚拟环境，并准确预测行动对世界状态的影响，从而省去繁琐的人工编码？

本文将探讨 LLM 作为文本世界模拟器的潜力，并通过一个名为 BYTESIZED32-State-Prediction 的全新基准数据集进行评估。

模拟世界的挑战

模拟世界对于理解和研究世界至关重要，但传统上，构建一个复杂的模拟环境需要大量的人工投入，耗费大量时间和精力。LLM 的出现为我们提供了一种新的思路，即利用其庞大的预训练数据集，直接将其作为模拟器使用。

然而，LLM 真的能胜任模拟器的角色吗？为了回答这个问题，研究人员将目光投向了文本游戏领域。文本游戏以自然语言描述环境和动态变化，长期以来被用于决策过程、信息提取和人工智能推理等研究领域。

两种利用 LLM 进行世界建模的方法

研究人员提出两种利用 LLM 进行世界建模和模拟的方法：

1. 神经符号方法: 利用 LLM 生成符号表示的代码，以便进行形式化规划或推理。例如，REASONING VIA PLANNING (RAP) 方法利用 LLM 的先验知识构建世界模型，然后使用专门的规划算法来决定代理策略。
2. 直接模拟方法: 利用 LLM 直接生成文本描述，构建虚拟环境，并根据用户输入的行动进行模拟。例如，AI-DUNGEON 项目使用 LLM 生成文本描述，构建一个纯粹由语言模型驱动的游戏世界。

本文重点关注第二种方法，即直接模拟方法，并首次对 LLM 直接模拟虚拟环境的能力进行了量化分析。

BYTESIZED32-State-Prediction 基准数据集

为了评估 LLM 作为文本世界模拟器的能力，研究人员构建了一个名为 BYTESIZED32-State-Prediction (BYTESIZED32-SP) 的全新基准数据集。该数据集包含 76,369 个文本游戏状态转换，每个转换都由一个七元组 (S, A, T, O, R, C, D) 表示，分别对应状态空间、动作空间、转换函数、观察函数、奖励函数、上下文信息和完成指示函数。

该数据集从 BYTESIZED32 数据集派生而来，BYTESIZED32 数据集包含 32 个由人类编写的文本游戏，每个游戏模拟不同的科学或常识推理概念。研究人员通过修改每个 BYTESIZED32 游戏，使其能够在每个时间步输出游戏状态 (st, rt, dt) 和中间状态 sactt+1，并以 JSON 对象的形式存储。

LLM-Sim 任务

研究人员定义了一个名为 LLM-Sim 的预测任务，用于评估 LLM 作为可靠模拟器的能力。LLM-Sim 任务的目标是实现一个函数 F : C × S × A → S × R × {0, 1}，该函数将给定的上下文信息、状态和行动 (c, st, at) 映射到后续状态、奖励和游戏完成状态 (st+1, rt+1, dt+1)。

为了更好地理解 LLM 模拟不同类型状态转换的能力，研究人员将模拟函数 F 分解为三个步骤：

1. 动作驱动转换模拟器 Fact: 预测给定上下文信息、状态和行动 (c, st, at) 后，状态的直接变化 sactt+1。
2. 环境驱动转换模拟器 Fenv: 预测给定上下文信息和动作驱动转换后的状态 (c, sactt+1) 后，环境因素引起的额外状态变化 st+1。
3. 游戏进度模拟器 FR: 预测给定上下文信息、状态和行动 (c, st+1, at) 后，游戏的奖励 rt+1 和完成状态 dt+1。

研究人员分别评估了 LLM 模拟 Fact、Fenv 和 FR 的能力，以及模拟完整 F (即包含所有转换) 的能力。

实验结果

研究人员使用 GPT-4 对 BYTESIZED32-SP 数据集进行了评估，结果表明：

• 预测动作驱动转换比预测环境驱动转换更容易: GPT-4 在模拟动态动作驱动转换方面的最佳准确率为 77.1%，而在模拟动态环境驱动转换方面的最佳准确率仅为 49.7%。
• 预测静态转换比预测动态转换更容易: 模拟静态转换比模拟动态转换更容易，因为静态转换只需要判断是否发生状态变化，而动态转换还需要模拟环境因素的影响。
• 预测完整游戏状态对于动态状态更容易，而预测状态差异对于静态状态更容易: 预测动态状态的差异可以显著提高模拟静态转换的性能，但会降低模拟动态转换的性能。
• 游戏规则很重要，LLM 能够生成足够好的游戏规则: 当上下文信息中没有提供游戏规则时，GPT-4 在所有三个模拟任务上的性能都会下降。然而，研究人员没有发现人类专家编写的游戏规则和 LLM 生成的游戏规则之间存在明显的性能差异。
• GPT-4 能够在大多数情况下预测游戏进度: 当上下文信息中包含游戏规则时，GPT-4 能够在 92.1% 的测试用例中正确预测游戏进度。
• 人类在 LLM-Sim 任务上比 GPT-4 表现更好: 研究人员进行了一项初步的人类研究，结果表明，人类在模拟 Fact 方面的准确率为 80%，而 GPT-4 的准确率为 50%。
• GPT-4 在需要算术、常识或科学知识时更容易出错: 研究人员发现，GPT-4 在模拟需要算术、常识或科学知识的动态转换时更容易出错。

结论

本文提出了 BYTESIZED32-State-Prediction 基准数据集，用于评估 LLM 作为模拟器的能力。研究人员使用 GPT-4 对该数据集进行了评估，结果表明，LLM 尚未能够可靠地充当文本世界模拟器。

进一步的错误分析表明，虽然 LLM 在模拟用户行动的结果方面表现较好，但难以处理环境驱动转换以及需要算术、常识或科学知识的转换。

局限性和伦理问题

本文的研究存在一些局限性和伦理问题：

• 模型选择: 研究人员只评估了 GPT-3.5 和 GPT-4 两种模型，可能存在其他模型在该任务上表现更好。
• 状态表示: 研究人员使用了两种状态表示形式，即完整状态空间和状态差异，可能存在其他更有效的表示形式。
• 领域局限: 本文的研究主要集中在常识和基础科学推理领域，可能无法推广到其他领域，例如物理或医学模拟。
• 伦理风险: LLM 作为文本世界模拟器可能会生成虚假或误导性信息，在某些应用场景中可能存在伦理风险。

未来展望

尽管 LLM 作为文本世界模拟器仍面临挑战，但该领域的研究具有重要的意义，未来需要进一步探索以下方向：

• 提高 LLM 的推理能力: 增强 LLM 对算术、常识和科学知识的理解和运用能力。
• 开发更有效的训练方法: 针对模拟任务，开发专门的训练方法，提高 LLM 的模拟精度。
• 探索新的状态表示形式: 研究更适合模拟任务的新的状态表示形式。
• 关注伦理问题: 在开发和应用 LLM 作为模拟器时，需要充分考虑伦理问题，确保其安全性和可靠性。

参考文献

• Achiam, J., et al. (2023). GPT-4. [Online; accessed 2023-03-14].
• Ammanabrolu, P., & Hausknecht, M. (2020). A text-based adventure game for interactive learning and evaluation of natural language understanding. arXiv preprint arXiv:2005.02294.
• Adhikari, A., et al. (2020). Towards a text-based game for evaluating grounded language understanding. arXiv preprint arXiv:2005.03442.
• Côté, M.-A., et al. (2018). The unreasonable effectiveness of deep learning for text-based games. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing (pp. 1830-1839).
• Fan, A., et al. (2020). Learning to play text-based games with a language model. arXiv preprint arXiv:2003.07617.
• Fakhoury, S., et al. (2023). Evaluating the factual consistency of language models. arXiv preprint arXiv:2301.07187.
• Hao, B., et al. (2023). Reasoning via planning: A language model-based approach to symbolic reasoning. arXiv preprint arXiv:2303.16960.
• Hausknecht, M., et al. (2020). A text-based adventure game for interactive learning and evaluation of natural language understanding. arXiv preprint arXiv:2005.02294.
• Jansen, P. (2022). Text-based games for grounded language understanding: A survey. arXiv preprint arXiv:2206.02437.
• Kaelbling, L. P., et al. (1998). Reinforcement learning: A survey. Journal of artificial intelligence research, 8, 237-282.
• Liu, H., et al. (2023). Code as a language: Towards a unified framework for code and natural language. arXiv preprint arXiv:2303.17581.
• Nottingham, W., et al. (2023). Towards general-purpose language models for code generation. arXiv preprint arXiv:2303.16627.
• Shridhar, K., et al. (2020). Text-based adventure games as a testbed for grounded language understanding. arXiv preprint arXiv:2003.04604.
• Tang, Y., et al. (2024). Towards a unified framework for code and natural language via large language models. arXiv preprint arXiv:2401.04156.
• Urbanek, T., et al. (2019). Learning to play text-based games with a language model. arXiv preprint arXiv:1907.03718.
• Valmeekam, V., et al. (2023). Language models are not planners. arXiv preprint arXiv:2303.16487.
• Walton, C. (2020). AI Dungeon: A text adventure powered by GPT-3. [Online; accessed 2023-03-14].
• Wang, R., et al. (2022). A text-based game for evaluating commonsense reasoning in language models. arXiv preprint arXiv:2205.14871.
• Wang, R., et al. (2023). BYTESIZED32: A benchmark for evaluating scientific reasoning in language models. arXiv preprint arXiv:2303.11612.
• Wong, A., et al. (2023). Code as a language: Towards a unified framework for code and natural language. arXiv preprint arXiv:2303.17581.

---

Q: 这篇论文试图解决什么问题？

A: 这篇论文探讨的主要问题是：当前的语言模型（Language Models，简称LLMs）是否能够作为文本基础的世界模拟器（text-based world simulators），准确地预测动作如何改变不同的世界状态，从而绕过手动编码的需要。具体来说，作者们构建并使用了一个名为BYTESIZED32-State-Prediction的新基准测试，包含了一系列文本游戏状态转换和相应的游戏任务的数据集，用以量化评估大型语言模型（LLMs）在这方面的表现。研究发现，尽管GPT-4等模型表现出色，但如果没有进一步的创新，它们仍然不能作为一个可靠的世界模拟器。这项工作不仅提供了对当前LLMs能力和弱点的新见解，还提供了一个新的基准测试，以跟踪未来模型发展的进步。

Q: 有哪些相关研究？

A: 论文中提到了多项相关研究，主要集中在使用大型语言模型（LLMs）来模拟环境和执行决策任务的领域。以下是一些具体的相关研究：

1. Côté et al., 2018: 提出了Textworld，一个用于文本游戏的学习环境。
2. Fan et al., 2020: 研究了如何生成交互式世界。
3. Urbanek et al., 2019: 研究了在幻想文本冒险游戏中的学习和行动。
4. Shridhar et al., 2020: 提出了Alfworld，一个将文本与具身环境对齐的交互式学习平台。
5. Hausknecht et al., 2020: 研究了交互式小说游戏。
6. Jansen, 2022: 对文本世界进行了系统性的调查。
7. Wang et al., 2022: 提出了Scienceworld，一个测试代理是否比五年级学生更聪明的平台。
8. Wang et al., 2023: 提出了ByteSized32，一个用于生成特定任务世界模型的文本游戏的语料库和挑战任务。

此外，论文还提到了一些使用语言模型生成代码以进行形式化规划或推理的研究，例如：

• Liu et al., 2023: 研究了如何增强大型语言模型的规划能力。
• Nottingham et al., 2023: 研究了使用语言引导的世界建模来进行具身决策。
• Tang et al., 2024: 提出了Worldcoder，一个通过编写代码和与环境交互来构建世界模型的模型。

这些研究展示了LLMs在不同领域的应用潜力，特别是在模拟环境和决策制定方面。论文通过这些相关工作，为读者提供了一个关于LLMs在文本游戏模拟器领域应用的全面背景。

Q: 论文如何解决这个问题？

A: 论文通过以下步骤解决这个问题：

1. 构建基准测试：作者创建了一个名为BYTESIZED32-State-Prediction的新基准测试，它包含了一系列文本游戏状态转换和相应的游戏任务的数据集。
2. 定义预测任务：提出了一个名为LLMas-a-Simulator (LLM-Sim) 的预测任务，用于量化评估语言模型作为可靠模拟器的能力。这个任务涉及从给定的上下文、状态和动作映射到后续状态、奖励和游戏完成状态。
3. 分解模拟器功能：将模拟器功能进一步分解为三个步骤：
   • 动作驱动的转换模拟器（Fact）：预测由动作直接引起的状态变化。
   • 环境驱动的转换模拟器（Fenv）：预测在动作影响之后，环境如何进一步改变状态。
   • 游戏进度模拟器（FR）：预测奖励和游戏完成状态。
4. 数据收集：从BYTESIZED32语料库中收集数据，该语料库包含32个人工编写的文本游戏，每个游戏都模拟不同的科学或常识推理概念。
5. 评估方法：使用提出的基准测试来评估GPT-4模型的性能，通过准确率来衡量模型对状态转换的预测能力。
6. 实验和分析：对GPT-4进行系统性分析，发现它在模拟非平凡状态转换时的准确性不超过59.9%。此外，还进行了人类研究，比较了人类注释者和GPT-4在LLM-Sim任务上的表现。
7. 结果讨论：论文讨论了GPT-4在模拟任务中的表现，指出了其在模拟环境驱动转换和需要算术、常识或科学知识转换时的不足。
8. 结论：基于实验结果，论文得出结论，尽管LLMs在某些方面表现出前景，但目前还不能作为可靠的文本世界模拟器，除非进行进一步的创新。
9. 未来工作：论文提出了对未来模型的评估和开发更高性能的语言模型模拟器的期望。

通过这些步骤，论文不仅提供了对当前LLMs能力的深入分析，还为未来的研究提供了一个基准和方向。

Q: 论文做了哪些实验？

A: 论文中进行了以下实验来评估语言模型作为文本世界模拟器的能力：

1. LLM-Sim任务评估：提出了LLMas-a-Simulator (LLM-Sim)任务，用于定量评估语言模型作为可靠模拟器的性能。这包括模拟动作驱动的转换（Fact）、环境驱动的转换（Fenv）以及游戏进度（FR）。
2. 数据集构建：创建了BYTESIZED32-State-Prediction (BYTESIZED32-SP)数据集，包含76,369个文本游戏状态转换，这些转换是从31个不同的文本游戏中收集的。
3. 数据收集方法：从BYTESIZED32语料库中修改游戏，以转储每个时间步的游戏状态，并确定有效的转换。
4. 上下文信息：为每个游戏提供了上下文信息，包括动作规则、对象规则、得分规则和示例转换。
5. 模型评估：使用GPT-4模型对LLM-Sim任务进行评估，测量其在完整状态预测和状态差异预测两种模式下的准确性。
6. 动态与静态转换的比较：分析了模型在模拟动态（状态发生变化）和静态（状态不发生变化）转换时的性能差异。
7. 人类与模型性能比较：进行了一项初步的人类研究，让人类注释者模拟Fact任务，并与GPT-4的性能进行了比较。
8. 错误分析：对GPT-4在模拟动态转换时的错误模式进行了深入分析，特别是当转换需要算术、常识或科学知识时。
9. 游戏进度预测：评估了GPT-4在预测游戏进度方面的能力，包括游戏得分和游戏结束状态。
10. 规则的重要性：研究了游戏规则对GPT-4性能的影响，并比较了人类编写的规则和LLM生成的规则。

这些实验提供了对当前语言模型在模拟文本世界方面能力的全面评估，并揭示了它们在这一任务上的局限性。

Q: 有什么可以进一步探索的点？

A: 根据论文的结论和局限性部分，以下是一些可以进一步探索的点：

1. 更广泛的语言模型评估：论文中提到，尽管GPT-3.5和GPT-4是当时表现最好的模型，但并没有详尽地评估所有大型语言模型。未来的研究可以探索更多种类的模型，以确定是否有其他模型在模拟任务上表现更好。
2. 不同的表示形式：论文中使用了JSON对象来表示状态空间。未来的工作可以探索其他表示形式，以确定是否有更有效的方式来表示和处理状态空间。
3. 特定领域的模拟：论文中的实验主要集中在常识和初级科学推理任务上。未来的研究可以探索使用LLMs作为模拟器在更特定、更复杂的领域，如物理模拟或医学模拟。
4. 改进的算法和架构：论文指出，当前的LLMs在模拟环境驱动的转换和需要算术、常识或科学知识的任务时存在困难。未来的研究可以探索改进的算法和模型架构，以提高这些领域的性能。
5. 减少模拟错误的影响：由于模拟错误会随着时间的推移而累积，研究如何减少单步模拟误差的影响，或者如何设计能够自我纠正的模拟器，可能是一个有价值的研究方向。
6. 伦理和社会影响：论文提到了LLMs在生成误导性或非事实信息方面的潜在风险。未来的研究可以更深入地探讨这些伦理问题，并开发策略来减轻这些风险。
7. 教育和实际应用：论文建议LLMs作为模拟器的应用在教育和与儿童直接互动的环境中可能不安全。研究如何安全地将这些技术应用于教育和其他实际场景，将是一个重要的研究方向。
8. 长期跟踪和评估：论文提出了一个基准测试来评估现有和未来的模型。继续维护和更新这个基准测试，以跟踪语言模型在模拟任务上的进步，将是一个持续的过程。

这些探索点可以帮助推动语言模型在模拟任务上的性能，同时确保这些技术的发展和应用是安全和负责任的。

Q: 总结一下论文的主要内容

A: 这篇论文的主要内容可以总结如下：

1. 研究问题：探讨当前语言模型（LLMs）是否能够作为文本基础的世界模拟器，预测动作如何改变世界状态，从而减少手动编码的需求。
2. 方法论：
   • 构建了一个新的基准测试，名为BYTESIZED32-State-Prediction，包含文本游戏状态转换和任务的数据集。
   • 提出了LLMas-a-Simulator (LLM-Sim)任务，用于量化评估LLMs作为模拟器的能力。
3. 实验设计：
   • 使用BYTESIZED32-SP数据集对GPT-4模型进行测试。
   • 评估了模型在完整状态预测和状态差异预测两种模式下的性能。
4. 主要发现：
   • GPT-4在模拟动作驱动的转换方面表现相对较好，但在环境驱动的转换和需要算术、常识或科学知识的任务上表现不佳。
   • 人类在LLM-Sim任务上的表现优于GPT-4。
5. 结论：
   • 尽管LLMs在某些方面有潜力，但目前还不能作为可靠的文本世界模拟器，需要进一步的创新。
6. 局限性和伦理问题：
   • 论文讨论了研究的局限性，包括评估的模型范围、表示形式的选择和任务类型的限制。
   • 提出了使用LLMs作为模拟器可能带来的伦理和社会风险，并强调了负责任使用的重要性。
7. 未来工作：
   • 论文提出了对未来模型的评估和开发更高性能的语言模型模拟器的期望。

论文通过构建新的基准测试和对现有模型的评估，为理解LLMs在模拟任务上的能力提供了新的见解，并为未来的研究指明了方向。

---

大型语言模型（LLM）在自然语言处理领域取得了巨大成功，其参数规模不断攀升，突破了千亿甚至万亿级别。然而，这些模型的推理效率却面临着严峻挑战，尤其是当需要处理长序列文本并满足严格的延迟要求时。本文将深入探讨 Transformer 模型推理的效率问题，并介绍一系列工程优化策略，旨在突破模型规模和推理效率之间的瓶颈。

推理成本的权衡

随着模型规模的增长，推理成本也随之增加。我们主要关注三个关键指标：延迟、吞吐量和模型 FLOPS 利用率 (MFU)。延迟是指完成一次推理所需的时间，可以细分为处理输入文本的时间（称为“预填充”）和生成输出文本的时间（称为“解码”）。解码延迟也可以按“每步”计算，即除以每个序列中的令牌数。吞吐量是指每秒处理或生成的令牌数。MFU 则是实际吞吐量与理论峰值吞吐量的比率，反映了硬件资源的利用效率。

大型模型通常无法完全容纳在一块加速器芯片的内存中，需要进行模型划分，将模型参数和激活张量分布在多个芯片上。这种划分虽然可以降低每个芯片的内存和计算压力，但也引入了芯片间通信的开销。

内存成本：模型参数和 KV 缓存（每个层中的注意力键和值张量）需要存储在芯片上的高带宽内存 (HBM) 中。每次前向传播（预填充或解码步骤）都需要将这些张量从 HBM 加载到计算核心，这会消耗一定的时间，称为“内存时间”。在小批量和短序列情况下，加载权重的时间占主导地位。而在大批量和长序列情况下，加载 KV 缓存的时间则会占主导地位。

计算成本：一个拥有 N 个参数的解码器模型，每个令牌需要进行 2N 次矩阵乘法运算。如果所有芯片都以峰值 FLOPS 运行，这些矩阵乘法需要一定的时间，称为“计算时间”。注意力机制的矩阵乘法通常占用的 FLOPS 较少，但在长序列情况下，KV 缓存的内存占用和带宽需求会显著增加。

优化策略：模型划分

为了高效地进行推理，我们需要对大型模型进行合理的划分。本文将介绍几种模型划分策略，并分析其在不同模型规模、序列长度和应用需求下的性能表现。

1. 前馈层划分

1D 权重固定布局：这是最简单的划分策略，将每个 E × F 权重矩阵沿 E 或 F 轴进行划分（或分片），每个权重分片在相应的芯片上与激活分片进行乘法运算，并将结果通过全聚合和/或降维散射操作进行聚合。这种策略在芯片数量较少时，内存延迟和计算延迟会随着芯片数量的增加而线性下降。然而，通信延迟基本保持不变，因为每次矩阵乘法都需要将整个激活矩阵进行聚合。当芯片数量增加时，通信成为瓶颈。

2D 权重固定布局：当芯片数量较多时，可以将每个 E × F 权重矩阵沿 E 和 F 轴进行划分，使得每个分片近似为正方形。这种策略称为 2D 权重固定布局。虽然计算成本与 1D 权重固定布局相同，但通信效率更高。通过交替地沿 E 和 F 轴进行激活聚合，可以确保每个芯片始终拥有所需的激活分片，而无需完全复制激活张量。通信时间随着芯片数量的增加而减小，因此即使在通信成为瓶颈的情况下，也可以通过增加芯片数量来降低延迟。

权重聚合布局：在权重固定布局中，每个芯片存储一个权重矩阵的分片，并负责将其与相应的激活分片进行乘法运算。每个芯片的矩阵乘法结果需要进行聚合，才能作为后续操作的输入。然而，当批量大小（和序列长度）增加时，输出激活的大小可能远大于权重的大小。在这种情况下，将激活固定在每个芯片上，并将权重在芯片之间进行传输会更经济。对于非常大的批量大小，最好将激活完全固定在连续的矩阵乘法之间，这需要将权重完全传输到所有芯片之间。我们称这种方法为 XYZ 权重聚合。对于中等批量大小，使用“混合”方法是有益的，即权重和激活都沿不同的轴进行部分传输。我们将这些方法称为 X 权重聚合和 XY 权重聚合。

2. 注意力层划分

多头注意力：多头注意力可以与前馈层类似地进行划分，将 nheads 视为 dff。然而，多头注意力在存储和加载 KV 缓存方面会产生大量的内存容量和带宽成本，这在大批量或长序列情况下可能会成为主要的性能瓶颈。

多查询注意力：多查询注意力是一种替代方案，它仍然为查询张量输出 nheads，但仅为键和值张量输出一个头，该头在所有 nheads 查询头之间共享。这将 KV 缓存张量的大小减少了 nheads 倍，从而减少了加载它们所需的内存时间。但它也去掉了原本用于并行化的一个轴，因此 KV 缓存和相关的计算需要进行不同的划分。

优化策略：为了最大程度地减少加载 KV 缓存所需的内存时间，我们将 Q、K 和 V 矩阵沿批次 B 轴进行划分。这种策略可以将每个芯片加载 KV 缓存的内存成本降低 nheads 倍，从而减少内存时间。虽然这种策略会增加额外的通信成本，但与沿头轴划分相比，它可以显著减少 KV 缓存的内存占用，从而提升推理效率。

3. 并行注意力/前馈层

PaLM 模型采用了一种并行化的 Transformer 块结构，将注意力层和前馈层并行计算，并将其结果相加得到输出。这种结构有两个主要优势：

• 减少延迟：每个层只有一个层归一化操作，而不是两个，这在小批量情况下可以降低延迟。
• 提高 FLOPS 利用率：可以将前馈层的输入矩阵与注意力层的查询投影矩阵 WQ 进行融合，将键/值投影矩阵 WK 和 WV 融合到注意力层中，并将前馈层的输出矩阵与注意力层的输出投影矩阵 WO 进行融合。这种融合可以提高 FLOPS 利用率，因为它可以更有效地运行加速器上的大型矩阵乘法。更重要的是，它还消除了每个 Transformer 层中用于 dff/nheads 并行化所需的两个全聚合操作之一，将沿该轴的通信时间减少了一半。

低级优化

除了模型划分策略，我们还采用了一系列低级优化技术，进一步提升推理效率：

• 循环集体 einsum：将通信与计算并行执行，可以部分或完全隐藏大部分降维散射和全聚合操作的通信时间。
• 集体 einsum 优化：通过开发集体 einsum 的多种变体，可以针对不同的场景进行优化，例如延迟与吞吐量、不同的环形轴数量以及与不同输入/输出集体进行融合。
• 内存布局优化：优化张量在内存中的布局，可以减少矩阵乘法过程中的填充和复制操作。
• 快速 top-k/top-p 实现：加速解码采样过程。
• 快速 log-base-2 实现：加速 Softmax 和 Swish 函数的计算。
• 增量序列处理：支持在预填充阶段对序列进行增量处理。

量化

为了降低权重存储的内存成本，我们使用 AQT 库将 16 位权重转换为 8 位整数。这可以节省权重加载所需的内存时间，尤其是在小批量情况下，并减少权重聚合布局中的通信量。

案例研究：PaLM 模型

为了验证上述优化策略的有效性，我们对 PaLM 模型家族进行了实验，包括 8B、62B 和 540B 参数模型，并使用 bfloat16 或 int8 格式的权重。

1. 前馈层划分

实验结果表明，2D 权重固定布局在解码阶段的性能优于 1D 权重固定布局，因为其在芯片数量增加时的扩展性更好。在预填充阶段，随着批量大小的增加，最优的划分布局从 2D 权重固定布局转变为权重聚合布局。权重聚合布局在小批量情况下效率较低，但在高批量情况下效率最高，可以实现高达 76% 的 MFU。

2. 注意力层划分

实验结果表明，将多查询注意力沿批次 B 轴进行划分可以显著提高推理效率，与沿头轴划分相比，它可以支持 32 倍以上的上下文长度。

3. 并行注意力/前馈层

实验结果表明，并行化 Transformer 块结构可以有效地降低延迟，尤其是在解码阶段。

4. 端到端结果

我们通过调整批量大小、芯片数量和划分策略，获得了 PaLM 模型家族在不同模型规模、序列长度和延迟要求下的 Pareto 前沿。结果表明，在高批量情况下，推理成本与模型参数数量成正比。通过降低批量大小，可以提高延迟，但会增加每个令牌的成本。

实验结果还表明，int8 权重量化可以有效地降低延迟。在低延迟目标情况下，int8 权重量化可以将成本降低一半以上。在高批量情况下，int8 和 bfloat16 权重的成本差异较小。

与 FasterTransformer 的比较

我们还将我们的实现与 FasterTransformer 基准进行了比较，结果表明，我们的实现可以实现更高的 MFU 和更低的延迟，尤其是在大批量情况下。这主要归功于我们采用的 2D 权重固定布局和 TPU v4 的高带宽互连网络。

总结

本文提出了一系列工程优化策略，可以有效地提高 Transformer 模型推理的效率，尤其是当需要处理长序列文本并满足严格的延迟要求时。这些策略可以应用于各种硬件平台，包括 GPU 和 TPU。

未来，我们希望通过进一步探索稀疏性技术、自适应计算技术等方法，进一步降低 Transformer 模型的 FLOPS 数量和通信量，从而实现更高的推理效率。

参考文献

2211.05102v1.pdf (https://arxiv.org/pdf/2211.05102)

近年来，巨型语言模型（LLM）在自然语言处理领域取得了显著进展，其参数规模已突破千亿甚至万亿大关。然而，这些模型的推理效率却面临着巨大的挑战，尤其是当需要处理长序列文本并满足严格的延迟要求时。本文将深入探讨如何通过模型分割和优化策略来提高巨型语言模型的推理效率。

推理效率的挑战

巨型语言模型的推理效率主要受以下因素影响：

• 庞大的内存占用： 训练好的模型参数和解码过程中的中间状态都需要占用大量的内存空间。
• 低并行性： 与训练过程相比，推理过程的并行性较低，因为每个token的生成都依赖于前一个token。
• 注意力机制的二次复杂度： 注意力机制的计算量与输入序列长度的平方成正比，这会随着序列长度的增长而显著增加计算成本。

分割策略：高效利用硬件资源

为了解决上述挑战，本文提出了一个基于模型分割的框架，旨在高效利用硬件资源，并根据应用需求选择最佳分割策略。

3.1 分割符号和通信机制

本文采用了一种基于TPU v4架构的3D torus拓扑结构的分割符号和通信机制。例如，符号 BLExyz 表示将一个逻辑形状为 BLE 的张量沿最后一个维度 E 分割成 X × Y × Z 个分区，其中 x，y 和 z 分别代表TPU v4的三个物理轴，每个芯片上的张量形状为 [B, L, E/(X × Y × Z)]。

3.2 前馈层分割策略

3.2.1 一维权重固定布局

最简单的分割策略是将每个 E × F 权重矩阵沿 E 或 F 轴分割成 nchips 个分区，每个分区在相应的芯片上与激活张量进行矩阵乘法，然后使用 all-gather 和 reduce-scatter 操作进行跨芯片聚合。这种策略在芯片数量较少时效率较高，但随着芯片数量的增加，通信成本会成为瓶颈。

3.2.2 二维权重固定布局

为了提高通信效率，可以将每个 E × F 权重矩阵沿 E 和 F 轴进行二维分割，使每个分区近似为正方形。这种策略被称为二维权重固定布局。它可以有效减少通信成本，因为我们可以交替地在两个轴上进行激活张量的聚合，从而避免在每个矩阵乘法过程中都进行全量复制。

3.2.3 权重收集布局

当批处理大小和序列长度较大时，激活张量的尺寸可能会超过权重张量，此时可以将激活张量固定在每个芯片上，并将权重张量在芯片之间进行传输。这种策略被称为权重收集布局。

实验结果和结论

本文对PaLM系列巨型语言模型进行了实验验证，结果表明：

• 通过合理的模型分割策略，可以有效提高推理效率，降低延迟和成本。
• 多查询注意力机制可以有效减少内存占用，从而提高批处理大小，进而提升吞吐量。
• 在64个TPU v4芯片上，PaLM 540B模型可以实现29ms/token的低延迟生成速度，以及76%的模型FLOPS利用率，同时支持2048个token的上下文长度。

总而言之，本文提出的模型分割和优化策略为高效部署巨型语言模型提供了重要的参考，并为进一步提升推理效率提供了新的思路。

参考文献

• Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. (2020). Language models are few-shot learners. arXiv preprint arXiv:2005.14165.
• Kaplan, J., McCandlish, S., Henighan, T., Brown, T., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., Amodei, D., et al. (2020). Scaling laws for neural language models. arXiv preprint arXiv:2001.08202.
• Rae, J. W., Borgeaud, S., Cai, T., Olah, C., Leike, J., Allen, L., Jeffery, S., Rosenthal, S., Ganguli, S., Molloy, I., et al. (2021). Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:2112.11400.
• Hoffmann, J., Habib, M., Lu, Y., Goyal, N., Zhang, X., Khandelwal, U., Das, A., Lee, K., Mishra, N., Gruslys, A., et al. (2022). Training language models with tens of trillions of parameters. arXiv preprint arXiv:2203.15556.
• Chowdhery, A., Bhatia, S., Mishra, N., Gruslys, A., Rajbhandari, S., Kumar, A., Leike, J., Allen, L., Rosenthal, S., Ganguli, S., et al. (2022). Scaling language models to 540 billion parameters. arXiv preprint arXiv:2203.15556.
• Smith, T., Zambaldi, V., Sukhbaatar, S., Raffel, C., Dhariwal, P., Leike, J., Allen, L., Rosenthal, S., Ganguli, S., Molloy, I., et al. (2022). Training language models with tens of trillions of parameters. arXiv preprint arXiv:2203.15556.
• Thoppilan, R., Sukhbaatar, S., He, J., Lee, K., Mishra, N., Gruslys, A., Rajbhandari, S., Kumar, A., Leike, J., Allen, L., et al. (2022). Scaling language models to 540 billion parameters. arXiv preprint arXiv:2203.15556.
• Sukhbaatar, S., Szlam, A., Weston, J., and Fergus, R. (2019). End-to-end efficient language modeling with data-parallel distributed attention. arXiv preprint arXiv:1907.04020.
• Choromanski, K., Rowland, M., So, A., Khan, M. E., Ballard, A., and Recht, B. (2020). Rethinking attention with performers. arXiv preprint arXiv:2009.13821.
• Dao, T., Guu, K., Lee, K., Tung, H. W., Pasupat, P., and Chang, M. W. (2022). Sparsity in deep learning: Overcoming the memory wall. arXiv preprint arXiv:2203.14722.
• Zheng, S., Li, Y., Yu, Y., Zhang, Z., and Liu, Z. (2022). Efficient large-scale language model inference on tpu v4 pods. arXiv preprint arXiv:2205.07354.
• Xu, B., Zhang, Z., Li, Y., Yu, Y., and Liu, Z. (2021). Efficient large-scale language model inference on tpus. arXiv preprint arXiv:2104.04420.
• Clarke, L., Glover, R., and MPI Forum (1994). MPI: A message-passing interface standard. Journal of Parallel and Distributed Computing, 22(1), 6-21.
• Rajbhandari, S., Rasheed, A., Madaan, A., Kumar, A., and Ganguli, S. (2020). Efficient large-scale language model training on tpus. arXiv preprint arXiv:2006.16668.
• Shoeybi, M., Patel, M., Goldfarb, C., Fevzi, B., Lee, J., Tran, L., and Parmar, N. (2019). Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053.

在现代的深度学习领域，如何优化模型性能成为了一个热门话题。许多人会依赖于一些曾经有效的小技巧，比如“使用就地操作！”、“将梯度设置为 None！”、“安装 PyTorch 1.10.0 而不是 1.10.1！”等。虽然这些方法有时能带来显著的性能提升，但从基本原理出发进行分析，往往能更系统地解决问题。

理解深度学习系统的三大组成部分

深度学习系统的性能可以分解为以下三个主要组件：

1. 计算（Compute）：在 GPU 上进行实际浮点运算（FLOPS）的时间。
2. 内存（Memory）：在 GPU 内部传输张量的时间。
3. 开销（Overhead）：其他所有时间，比如 Python 解释器运行和 CUDA 内核启动等。

理解自己深度学习系统的性能瓶颈所在，可以帮助我们有针对性地进行优化。例如，如果你的系统主要花费时间在内存传输上（即内存带宽受限），那么增加 GPU 的计算能力并不会有帮助。相反，如果你主要时间都在进行大量的矩阵乘法运算（即计算受限），那么减少开销的优化也无济于事。

接下来，我们将分别讨论计算、内存带宽和开销这三个组件。

计算：如何最大化 GPU 的 FLOPS

为了充分利用 GPU 的计算能力，我们需要尽量减少在其他部分花费的时间。GPU 的 FLOPS 越高，我们的计算效率就越高。然而，计算能力增长的速度远快于内存带宽的增长速度，使得实现高效计算变得更加困难。

例如，现代机器学习加速器（如 Nvidia 的 Tensor Cores）主要针对矩阵乘法进行了优化。如果你的计算主要不是矩阵乘法，那么你将无法完全利用 GPU 的计算能力。尽管如此，其他操作（如层归一化或激活函数）所需的 FLOPS 相对较少，因此对整体性能影响不大。

内存带宽：数据传输的成本

内存带宽成本是指在系统中移动数据的成本。这包括从 CPU 到 GPU、一个节点到另一个节点，甚至从 CUDA 全局内存到 CUDA 共享内存的传输成本。频繁的数据传输会占用大量时间，导致计算资源无法充分利用。

举个例子，当我们执行一个简单的操作如 torch.cos 时，我们将数据从存储单元移到计算单元，进行计算后再将数据移回存储单元。这种频繁的数据传输非常耗时，因此我们需要尽量减少这种操作。

操作融合：减少内存开销的关键

操作融合（Operator Fusion）是深度学习编译器中最重要的优化之一。它的基本思想是将多个计算操作合并为一个，从而减少数据在全局内存和计算单元之间的传输次数。例如，执行 x.cos().cos() 时，通常需要进行四次全局读取和写入，而通过操作融合，我们只需两次全局内存读取和写入。

这种优化在执行大规模计算时尤其重要，可以显著减少内存带宽成本，提高整体性能。

开销：代码执行之外的时间

开销指的是代码在执行实际计算和传输数据之外所花费的时间。现代 GPU 的计算速度非常快，而 Python 解释器相对较慢，因此在 Python 解释器中花费的时间会导致 GPU 资源闲置。

如何识别和减少开销

要识别是否受到开销限制，可以通过增加数据量来观察运行时间的变化。如果增加数据量后运行时间没有成比例增加，那么系统可能受到开销限制。使用 PyTorch 的分析工具也可以帮助识别 CPU 和 GPU 之间的配合情况，从而找到优化方向。

结论

要提升深度学习系统的性能，关键是识别模型的性能瓶颈，并有针对性地进行优化。以下是不同性能瓶颈的优化策略：

性能瓶颈	可能的解决方案
开销受限	使用 tracing，操作融合，避免使用 Python，采用真正的 JIT
内存带宽受限	操作融合
计算受限	使用 Tensor Cores，升级 GPU 硬件

理解这些基本原理和优化策略，可以帮助我们更有效地提升深度学习模型的性能。

---

参考文献：

• Horace He, Making Deep Learning Go Brrrr From First Principles

作为一名资深科技专栏作家，我接触过许多想要提升深度学习模型性能的用户。他们常常会采取一些“偏方”，比如使用“in-place operations”、将梯度设置为“None”、安装特定版本的PyTorch等等。

这些方法并非完全无效，但更像是炼金术而非科学。现代系统，特别是深度学习，其性能表现常常让人捉摸不透。然而，如果我们从基础原理出发，就能排除很多无效的方法，从而更高效地解决问题。

三大核心要素：计算、内存和开销

我们可以将深度学习系统的效率拆解为三个核心要素：

• 计算: GPU 用于实际浮点运算 (FLOPS) 的时间。
• 内存: 在 GPU 内部传输张量所花费的时间。
• 开销: 除此之外的一切时间消耗。

就像训练机器学习模型一样，了解系统的瓶颈所在，才能有的放矢地进行优化。例如，如果大部分时间都花在内存传输上（即内存带宽受限），那么提升 GPU 的 FLOPS 就毫无意义。反之，如果大部分时间都在进行大型矩阵乘法（即计算受限），那么用 C++ 重写模型逻辑以减少开销也无济于事。

计算：深度学习的引擎

理想情况下，我们希望最大化计算时间，毕竟我们花费了大量资金购买高性能 GPU，就应该充分利用其计算能力。然而，为了让矩阵乘法引擎高效运转，我们需要减少其他方面的耗时。

为什么 focus on 计算而不是内存带宽呢？ 因为我们无法在不改变实际操作的情况下减少所需的计算量，但可以通过优化来降低开销或内存成本。

雪上加霜的是，计算能力的增长速度远超内存带宽。下表展示了 CPU FLOPS 和内存带宽的翻倍时间：

指标	翻倍时间
CPU FLOPS	1 年
内存带宽	3 年

这种差距意味着，尽管 GPU 的计算能力越来越强，但如果内存带宽无法跟上，整体性能提升仍然有限。

内存带宽：数据传输的成本

内存带宽成本指的是将数据从一个地方移动到另一个地方所花费的成本。这可能包括将数据从 CPU 移动到 GPU、从一个节点移动到另一个节点，甚至从 CUDA 全局内存移动到 CUDA 共享内存。

回到工厂的比喻，GPU 的 DRAM 就好比仓库，用于存储大量数据和结果。每次执行 GPU 内核时，都需要将数据从仓库运送到工厂进行计算，然后再将结果运回仓库。

对于像 torch.cos 这样的简单操作，我们需要将数据从仓库运送到工厂，执行简单的计算，然后再将结果运回仓库。由于数据传输成本高昂，因此大部分时间都花在了数据传输上，而不是实际计算上。

操作融合：减少数据搬运的利器

为了减少内存带宽成本，我们可以采用操作融合技术。简单来说，就是将多个操作合并成一个，避免重复的数据读写。

例如，执行 x.cos().cos() 通常需要 4 次全局内存读写操作：

x1 = x.cos() # 从全局内存读取 x，写入 x1
x2 = x1.cos() # 从全局内存读取 x1，写入 x2

但通过操作融合，我们只需要 2 次全局内存读写操作：

x2 = x.cos().cos() # 从全局内存读取 x，写入 x2

操作融合是深度学习编译器中最重要的优化之一。它可以将多个操作合并到一起，从而节省内存带宽成本。

开销：Python 和框架的负担

开销是指代码执行过程中，除了张量传输和计算之外的所有时间消耗。例如，Python 解释器、PyTorch 框架、启动 CUDA 内核（但不执行）等都会产生开销。

现代 GPU 速度极快，而 Python 解释器却非常慢。在一个 A100 GPU 执行一次 FLOP 的时间内，Python 解释器只能执行几千万次加法运算。

PyTorch 等框架也存在多层调度机制，这也会增加开销。

为了减少开销，可以采用 JIT 编译、CUDA Graphs 等技术。

总结：对症下药，才能药到病除

总而言之，想要提升深度学习系统的性能，首先要了解系统的瓶颈所在。

性能瓶颈	解决方案
开销受限	JIT 编译、操作融合、避免使用 Python
内存带宽受限	操作融合
计算受限	使用 Tensor Cores、购买更强大的 GPU

当然，用户需要考虑这些问题，本身就反映了框架设计上的不足。PyTorch 的编译器和性能分析 API 并不完善，但也在不断改进。

希望本文能够帮助你更好地理解深度学习系统的性能优化，从而让你的模型运行得更快。

参考文献

He, H. (2022). Making Deep Learning Go Brrrr From First Principles. Retrieved from https://horace.io/brrr_intro.html

大型语言模型（LLM）如 ChatGPT 和 Llama 近年来备受关注，但它们的运行成本依然高昂。即使生成单个回复可能只需要 0.01 美元（在 AWS 上使用 8xA100 实例运行几秒钟），但当扩展到数十亿用户时，成本会迅速增加，因为这些用户每天可能与 LLM 进行多次交互。一些用例的成本更高，例如代码自动补全，因为每次输入新字符时都会运行。随着 LLM 应用的增加，即使生成时间略微提高效率，也会产生巨大的影响。

LLM 推理（或“解码”）是一个迭代过程：一次生成一个词元。生成包含 N 个词元的完整句子需要对模型进行 N 次前向传递。幸运的是，可以缓存先前计算的词元：这意味着单个生成步骤不依赖于上下文长度，除了一个操作：注意力机制。该操作的计算量随着上下文长度的增加而迅速增长。

LLM 的一些重要新兴用例利用了长上下文。有了更长的上下文，LLM 可以推理更长的文档，例如对文档进行摘要或回答有关文档的问题，它们可以跟踪更长的对话，甚至在编写代码之前处理整个代码库。例如，大多数 LLM 在 2022 年的上下文长度最多为 2k（GPT-3），但现在我们拥有上下文长度扩展到 32k（Llama-2-32k）甚至 100k（CodeLlama）的开源 LLM。在这种情况下，注意力机制在推理过程中占用了大量时间。

当扩展批次大小维度时，即使上下文相对较短，注意力机制也会成为瓶颈。这是因为要读取的内存量随着批次大小维度而扩展，而它仅取决于模型大小。

我们提出了一种名为 Flash-Decoding 的技术，它可以显著加快推理过程中的注意力机制，对于非常长的序列，可以使生成速度提高 8 倍。主要思想是尽可能快地并行加载键和值，然后分别重新缩放和组合结果以保持正确的注意力输出。

解码的多头注意力机制

在解码过程中，每个新生成的词元都需要关注所有先前的词元，以计算：

softmax(queries @ keys.transpose) @ values

此操作已在训练情况下使用 FlashAttention（最近的 v1 和 v2 版本）进行了优化，其中瓶颈是读取和写入中间结果（例如 Q @ K^T）的内存带宽。然而，这些优化不直接适用于推理情况，因为瓶颈不同。对于训练，FlashAttention 在批次大小和查询长度维度上并行化。在推理过程中，查询长度通常为 1：这意味着如果批次大小小于 GPU 上的流式多处理器数量（A100 为 108），则操作将仅使用 GPU 的一小部分！当使用长上下文时尤其如此，因为它需要更小的批次大小才能适应 GPU 内存。如果批次大小为 1，FlashAttention 将使用不到 GPU 的 1%！

FlashAttention 仅在查询块和批次大小上并行化，无法在解码过程中占用整个 GPU。

注意力机制也可以使用矩阵乘法原语来完成，而无需使用 FlashAttention。在这种情况下，操作会完全占用 GPU，但会启动许多内核来写入和读取中间结果，这并非最佳选择。

解码的更快注意力机制：Flash-Decoding

我们新的方法 Flash-Decoding 基于 FlashAttention，并添加了一个新的并行化维度：键/值序列长度。它结合了上述两种方法的优点。与 FlashAttention 一样，它存储到全局内存的额外数据很少，但即使批次大小很小，只要上下文长度足够长，它也能充分利用 GPU。

Flash-Decoding 也在键和值上并行化，但需要一个小的最终归约步骤。

Flash-Decoding 分三个步骤进行：

1. 首先，我们将键/值分成更小的块。
2. 我们使用 FlashAttention 并行计算查询与每个块的注意力。我们还为每行和每个块写入一个额外的标量：注意力值的 log-sum-exp。
3. 最后，我们使用 log-sum-exp 来缩放每个块的贡献，通过对所有块进行归约来计算实际输出。

所有这些都是可能的，因为注意力/softmax 可以迭代计算。在 Flash-Decoding 中，它在两个级别上使用：在块内（类似于 FlashAttention），以及在块之间进行最终归约。

实际上，步骤 (1) 不涉及任何 GPU 操作，因为键/值块是完整键/值张量的视图。然后我们有两个独立的内核分别执行 (2) 和 (3)。

CodeLlama 34B 的基准测试

为了验证这种方法，我们对 CodeLLaMa-34b 的解码吞吐量进行了基准测试。该模型与 Llama 2 具有相同的架构，更一般而言，结果应该可以推广到许多 LLM。我们测量了不同序列长度（从 512 到 64k）下的解码速度（tok/s），并比较了几种计算注意力机制的方法：

• Pytorch：使用纯 PyTorch 原语（不使用 FlashAttention）运行注意力机制。
• FlashAttention v2。
• FasterTransformer：使用 FasterTransformer 注意力内核。
• Flash-Decoding。
• 以及一个上限，计算为读取整个模型以及 KV 缓存所需的时间。

Flash-Decoding 在解码速度方面为非常长的序列带来了高达 8 倍的加速，并且比其他方法的扩展性更好。

所有方法在小提示情况下表现相似，但随着序列长度从 512 增加到 64k，扩展性都很差，除了 Flash-Decoding。在这种情况下（批次大小为 1），使用 Flash-Decoding，扩展序列长度对生成速度几乎没有影响。

组件级微基准测试

我们还在 A100 上对不同序列长度和批次大小的缩放多头注意力机制进行了微基准测试，输入为 f16。我们将批次大小设置为 1，并使用 16 个维度为 128 的查询头，用于 2 个键/值头（分组查询注意力），这与在 4 个 GPU 上运行 CodeLLaMa-34b 时使用的维度相匹配。

设置	算法	运行时间（us）
B=256, seqlen=256	PyTorch Eager	3058.6
B=256, seqlen=256	Flash-Attention v2.0.9	390.5
B=256, seqlen=256	Flash-Decoding	63.4
B=128, seqlen=512	PyTorch Eager	3151.4
B=128, seqlen=512	Flash-Attention v2.0.9	366.3
B=128, seqlen=512	Flash-Decoding	67.7
B=64, seqlen=1024	PyTorch Eager	3160.4
B=64, seqlen=1024	Flash-Attention v2.0.9	364.8
B=64, seqlen=1024	Flash-Decoding	77.7
B=32, seqlen=2048	PyTorch Eager	3158.3
B=32, seqlen=2048	Flash-Attention v2.0.9	352
B=32, seqlen=2048	Flash-Decoding	58.5
B=16, seqlen=4096	PyTorch Eager	3157
B=16, seqlen=4096	Flash-Attention v2.0.9	401.7
B=16, seqlen=4096	Flash-Decoding	57
B=8, seqlen=8192	PyTorch Eager	3173.1
B=8, seqlen=8192	Flash-Attention v2.0.9	529.2
B=8, seqlen=8192	Flash-Decoding	56.4
B=4, seqlen=16384	PyTorch Eager	3223
B=4, seqlen=16384	Flash-Attention v2.0.9	582.7
B=4, seqlen=16384	Flash-Decoding	58.2
B=2, seqlen=32768	PyTorch Eager	3224.1
B=2, seqlen=32768	Flash-Attention v2.0.9	1156.1
B=2, seqlen=32768	Flash-Decoding	60.3
B=1, seqlen=65536	PyTorch Eager	1335.6
B=1, seqlen=65536	Flash-Attention v2.0.9	2300.6
B=1, seqlen=65536	Flash-Decoding	64.4
B=1, seqlen=131072	PyTorch Eager	2664
B=1, seqlen=131072	Flash-Attention v2.0.9	4592.2
B=1, seqlen=131072	Flash-Decoding	106.6

多头注意力机制的微基准测试，运行时间为 us。Flash-Decoding 在序列长度扩展到 64k 时，运行时间几乎保持不变。

先前测量的端到端高达 8 倍的加速是可能的，因为注意力机制本身比 FlashAttention 快 50 倍。在序列长度达到 32k 之前，注意力时间大致保持不变，因为 Flash-Decoding 设法充分利用了 GPU。

如何使用 Flash-Decoding

Flash-decoding 可在以下位置使用：

• FlashAttention 包，从 2.2 版本开始。
• xFormers，从 0.0.22 版本开始，通过 xformers.ops.memory_efficient_attention。调度器将根据问题大小自动使用 Flash-Decoding 或 FlashAttention 方法。当这些方法不受支持时，它可以调度到一个高效的 triton 内核，该内核实现了 Flash-Decoding 算法。

LLaMa v2 / CodeLLaMa 的完整解码示例可在 FlashAttention 存储库 这里 和 xFormers 存储库 这里 找到。我们还提供了一个 LLaMa v1/v2 模型的最小高效解码代码示例，旨在快速、易于阅读、具有教育意义和可修改性。

致谢

感谢 Erich Elsen、Ashish Vaswani 和 Michaël Benesty 建议了将 KV 缓存加载拆分的这个想法。我们要感谢 Jeremy Reizenstein、Patrick Labatut 和 Andrew Tulloch 的宝贵讨论，以及 Quentin Carbonneaux 为 xFormers 提供高效的解码示例。我们还要感谢 Geeta Chauhan 和 Gregory Chanan 在写作方面提供的帮助，以及更广泛地为将本文发表在 PyTorch 博客上做出的贡献。

参考文献：

• Flash-Decoding for long-context inference | PyTorch

你是否想要提升深度学习模型的性能？面对这个问题，很多人会习惯性地使用一些“秘诀”，比如“使用就地操作！将梯度设置为 None！安装 PyTorch 1.10.0，但不要安装 1.10.1！”。

这种“灵机一动”式的优化方法虽然看似有效，但实际上却缺乏理论基础。深度学习模型的性能优化并非炼金术，而是需要从基本原理出发进行分析。

理解性能瓶颈

深度学习模型的性能优化可以从三个方面进行考虑：

• 计算 (Compute): GPU 上执行浮点运算 (FLOPS) 所花费的时间。
• 内存 (Memory): 在 GPU 内传输张量所花费的时间。
• 开销 (Overhead): 其他所有时间。

就像训练机器学习模型一样，了解模型所处的性能瓶颈可以帮助我们找到最有效的优化方法。例如，如果你的模型主要受限于内存带宽，那么提升 GPU 的 FLOPS 性能将毫无用处。相反，如果你的模型主要受限于计算能力，那么将模型逻辑重写为 C++ 代码来降低开销也无济于事。

计算：GPU 的核心能力

从某种程度上来说，优化深度学习系统就是最大化模型在计算受限状态下运行的时间。你为那些强大的 GPU 付出了高昂的代价，理所当然地希望它们能够发挥出全部的计算能力。然而，为了充分利用 GPU 的矩阵乘法能力，我们需要减少其他方面的开销，例如内存传输和系统开销。

为什么我们更关注最大化计算能力，而不是内存带宽？原因很简单：我们可以降低开销或内存成本，但我们无法在不改变实际操作的情况下减少计算量。

计算与内存带宽的矛盾

更糟糕的是，计算能力的增长速度远超内存带宽。下表展示了 CPU FLOPS 和内存带宽的翻倍时间：

组件	翻倍时间
CPU FLOPS	18 个月
内存带宽	24 个月

我们可以将计算能力想象成一个工厂。我们向工厂发送指令（开销），并提供原材料（内存带宽），以确保工厂高效运转（计算）。

如果工厂的效率提升速度快于我们提供原材料的速度，那么工厂将难以达到峰值效率。

即使工厂的规模（FLOPS）翻倍，如果我们的带宽跟不上，那么性能也不会翻倍。

这种计算能力增长速度快于内存带宽的趋势，一方面意味着机器学习系统工程师将长期拥有稳定的工作，另一方面也使得理解性能瓶颈变得更加重要。

现代 GPU 的特殊性

现代机器学习加速器都拥有专门用于矩阵乘法的硬件，例如英伟达的“Tensor Cores”。

如果你没有进行矩阵乘法，那么你只能达到 19.5 teraflops，而不是标称的 312 teraflops。需要注意的是，这并非 GPU 独有，事实上，TPU 的通用性甚至比 GPU 更低。

GPU 在非矩阵乘法运算方面速度较慢，这似乎是一个问题。那么，其他运算，例如层归一化或激活函数呢？事实是，这些运算在 FLOPS 方面几乎可以忽略不计。例如，让我们看一下这篇论文中关于 BERT 不同运算类型 FLOPS 数量的表格，其中“张量收缩”= 矩阵乘法。

运算类型	FLOPS 占比
张量收缩	99.8%
归一化	0.1%
点运算	0.1%

可以看到，非矩阵乘法运算只占总 FLOPS 的 0.2%，因此 GPU 在非矩阵乘法运算方面速度慢 15 倍并不重要。

内存带宽：数据传输的瓶颈

然而，在实际应用中，非矩阵乘法运算却往往比预期花费更多时间。罪魁祸首通常是数据在工厂和仓库之间传输的时间，也就是内存带宽成本。

内存带宽成本是指将数据从一个地方移动到另一个地方所花费的成本。这可能包括将数据从 CPU 移动到 GPU，从一个节点移动到另一个节点，甚至从 CUDA 全局内存移动到 CUDA 共享内存。其中，最后一种情况通常被称为“带宽成本”或“内存带宽成本”。

我们可以再次使用工厂的比喻来理解内存带宽成本。

虽然工厂是进行实际工作的场所，但它并不适合作为大规模存储单元。主要原因是，由于我们在工厂进行实际工作，因此所有存储都针对快速使用进行了优化（SRAM），而不是拥有大量的存储空间。

那么，我们应该在哪里存储实际结果和原材料呢？典型的做法是建立一个仓库，可能位于土地便宜且空间充足的地方（DRAM）。然后，我们可以将物资运送到工厂和从工厂运出（内存带宽）。

将物资运送到计算单元和从计算单元运出的成本就是所谓的“内存带宽”成本。顺便说一下，你的 GPU 的 DRAM 会显示在 nvidia-smi 中，它是你遇到“CUDA 内存不足”错误的主要原因。

需要注意的是，每次执行 GPU 内核时，都需要将数据从 GPU 的 DRAM（即仓库）中读取出来，并在执行完内核后将结果写入 DRAM。

现在，想象一下执行像 torch.cos 这样的单目运算会发生什么。我们需要将数据从存储单元运送到仓库，然后对每条数据进行少量计算，最后将数据运回存储单元。运输数据的成本非常高。因此，我们几乎所有时间都花在了运输数据上，而不是实际的计算本身。

由于我们所有时间都花在了内存带宽上，因此这种运算被称为内存受限运算，这意味着我们没有花太多时间进行计算。

操作融合：优化内存带宽

那么，我们该如何解决这个问题呢？让我们看一下一系列运算的执行过程。

一系列点运算的执行过程可能如下所示：

这种安排非常愚蠢。为什么我们要将相同的数据反复发送到全局内存，然后再发送回计算单元？我们应该将数据保留在工厂中，执行所有计算，然后再将其发送回存储单元。

我们可以通过操作融合来解决这个问题。操作融合是深度学习编译器中最重要的优化方法。简单来说，与其将数据写入全局内存，然后再读取回来，不如将多个计算合并在一起执行，从而避免额外的内存访问。

例如，如果我们执行 x.cos().cos()，通常需要进行 4 次全局读写操作。

x1 = x.cos()  # 从全局内存中读取 x，写入 x1
x2 = x1.cos()  # 从全局内存中读取 x1，写入 x2

但是，使用操作融合，我们只需要进行 2 次全局内存读写操作！因此，操作融合可以将速度提高 2 倍。

x2 = x.cos().cos()  # 从全局内存中读取 x，写入 x2

操作融合可以显著提高性能，但它也有一些限制。首先，GPU 需要在执行当前操作时知道下一步要执行的操作。因此，我们无法在 Eager 模式下执行操作融合，因为 PyTorch 在 Eager 模式下会逐个执行操作。其次，我们需要为此生成 CUDA 代码，这会带来新的挑战。

并非所有操作融合都像点运算那样简单。你可以将点运算融合到约简运算中，或者将点运算融合到矩阵乘法中。甚至矩阵乘法本身也可以看作是广播乘法和约简的融合。

如果你有兴趣编写自定义 CUDA 内核，那么操作融合将是你最受益的优化方法。任何两个 PyTorch 操作都存在融合的可能性，从而节省了在它们之间读写全局内存的内存带宽成本。此外，许多现有的编译器可以执行“简单的”融合，例如 NVFuser 和 XLA。然而，自动系统无法与人类的智慧相提并论，因此，如果你想尝试编写一些自定义 CUDA 内核，那么 Triton 是一个不错的起点。

最后，操作融合会带来一些意想不到的结果。例如，融合后的 x.cos().cos() 的执行时间几乎与单独调用 x.cos() 的时间相同。这就是为什么激活函数的成本几乎都相同，尽管 gelu 明显比 relu 包含更多的操作。

这一事实导致了重计算/激活检查点的一些有趣结果。本质上，进行额外的重计算可能会导致更少的内存带宽，从而减少运行时间。因此，我们可以通过重计算来降低内存和运行时间，我们在 AOTAutograd 中利用它构建了一个巧妙的最小割优化过程。你可以在此处阅读更多相关内容（也可能在未来的博客文章中进行介绍！）。

如何判断内存带宽受限

在判断你的操作是否内存带宽受限时，计算器可以起到很大的作用。

对于简单的操作，我们可以直接计算内存带宽。例如，A100 拥有 1.5 terabytes/second 的全局内存带宽，可以执行 19.5 teraflops/second 的计算。因此，如果你使用 32 位浮点数（即 4 字节），那么你可以在 GPU 执行 20 万亿次操作的同时加载 4000 亿个数字。此外，为了执行简单的单目运算（例如将张量乘以 2），实际上需要将张量写回全局内存。

因此，除非你的单目运算执行了大约一百次操作，否则你会花费更多时间执行内存访问，而不是实际的计算。

借助 NVFuser 这样的融合编译器，我们可以很容易地测量内存带宽。你可以在此处查看 Colab 代码。

如果我们使用 PyTorch 函数，例如：

def f(x: Tensor[N]):
    for _ in range(repeat):
        x = x * 2
    return x

并使用融合编译器对其进行基准测试，那么我们可以计算出不同 repeat 值下达成的 FLOPS 和内存带宽。增加 repeat 是增加计算量而不增加内存访问次数的一种简单方法，这也被称为增加计算强度。

具体来说，假设我们对这段代码进行基准测试，并找到每秒执行的迭代次数。那么，作为 N（张量大小）的函数，我们将执行 2*N 次内存访问，以及 N * repeat 次 FLOP。因此，达成的内存带宽将为 bytes_per_elem * 2 * N * itrs_per_second，达成的 FLOPS 将为 N * repeat * itrs_per_second。

现在，让我们绘制运行时间、FLOPS 和达成的内存带宽作为计算强度的函数。注意，所有内容都以对数对数刻度显示。

首先，请注意，直到我们执行 64 次乘法运算，运行时间才明显增加。这意味着，在此之前，我们主要受限于内存带宽，我们的计算能力大部分处于闲置状态。

因此，我们最初达成的 FLOPS 只有 0.2 teraflops。随着我们使计算强度翻倍，这个数字线性增长，直到我们接近 9.75 teraflops 的峰值 [1]。一旦我们接近峰值 teraflops，我们就认为是“计算受限”。

最后，你可以看到，我们达成的内存带宽最初接近峰值，随着我们增加计算强度，它开始下降。这正是我们所期望的，因为我们花费了越来越多的时间执行实际计算，而不是访问内存。

在这种情况下，很容易看出我们何时是计算受限的，何时是内存带宽受限的。对于 repeat < 32，我们饱和了内存带宽，而我们的计算能力未得到充分利用。相反，一旦 repeat > 64，我们看到我们的计算能力已经饱和（即接近峰值 FLOPS），而我们利用的内存带宽开始下降。

对于更大的系统，我们通常很难判断是计算受限还是内存带宽受限，因为它们通常包含计算受限和内存带宽受限的组件。

衡量计算受限程度的一种常见方法是测量达成的 FLOPS 占峰值 FLOPS 的百分比。例如，如果你达成了峰值 FLOPS 的 80%，那么你就知道你至少是 80% 的计算受限，这已经相当不错了！你剩下的时间可能都花在了执行内存带宽操作上。[2]

然而，除了内存带宽成本之外，还有一件事可能会导致你的 GPU 无法“Brrr”起来。

开销：系统瓶颈

开销是指你的代码花费在传输张量或计算之外的所有时间。例如，在 Python 解释器中花费的时间？开销。在 PyTorch 框架中花费的时间？开销。启动 CUDA 内核（但没有执行它们）花费的时间？也是开销。

开销之所以是一个棘手的问题，主要是因为现代 GPU 非常快。A100 每秒可以执行 312 万亿次浮点运算（312 TeraFLOPS）。相比之下，Python 非常慢。在本地进行基准测试，Python 每秒可以执行 3200 万次加法运算。

这意味着，在 Python 执行一次 FLOP 的时间里，A100 可以执行 975 万次 FLOPS。

更糟糕的是，Python 解释器并不是开销的唯一来源，像 PyTorch 这样的框架在到达实际内核之前也有很多层调度。如果你使用 PyTorch 执行相同的实验，我们每秒只能执行 28 万次操作。当然，PyTorch 的设计目的不是处理微小的张量，但是，如果你使用的是微小的张量（例如在科学计算中），你可能会发现 PyTorch 比 C++ 慢得多。

例如，看一下 PyTorch 执行一次加法的火焰图配置文件。那个方框？那是执行实际计算的部分。其他所有部分都是纯粹的开销。

鉴于此，你可能会惊讶于为什么有人会使用 PyTorch，但请记住，现代深度学习模型通常会执行大规模运算。此外，像 PyTorch 这样的框架是异步执行的。也就是说，当 PyTorch 运行 CUDA 内核时，它可以继续运行并在其后面排队更多 CUDA 内核。因此，只要 PyTorch 可以“领先于”CUDA 内核，大部分框架开销就会完全隐藏起来！

如果我们的 GPU 操作足够大，那么我们的 CPU 可以领先于 GPU（因此 CPU 开销无关紧要）。另一方面，如果我们的 GPU 操作太小，那么我们的 GPU 将大部分时间都浪费在昂贵的“纸镇”上。

那么，如何判断你是否处于这种状态呢？由于开销通常不会随着问题规模的增加而增加（而计算和内存会增加），因此判断开销受限的最简单方法是增加数据的规模。如果运行时间没有按比例增加，那么你就是开销受限的。例如，如果你将批次大小增加一倍，但运行时间只增加了 10%，那么你很可能是开销受限的。[3]

另一种方法是使用 PyTorch 分析器。在这里，粉红色的线实际上显示了 CPU 内核与 GPU 内核的匹配情况。

GPU 上有很多间隙，因为它正在等待 CPU 开销

我们的 CPU 运行速度远远超过 GPU

另一个补充说明 – nvidia-smi 中的“GPU-Util”（不是“Volatile GPU-Util”）条目基本上是测量底部的行中实际运行 GPU 内核的百分比。因此，这也是评估开销的一个好方法。

开销存在的主要原因是像 PyTorch 这样的框架具有很大的灵活性。本质上，需要花费大量时间来“弄清楚要做什么”。

这可能来自 Python（查找属性或调度到正确的函数）或 PyTorch 中的代码（PyTorch 的所有调度器）。例如，当你执行 a + b 时，需要执行以下步骤。

Python 需要查找 a 上的 add 调度到什么。
PyTorch 需要确定张量的许多属性（例如数据类型、设备以及是否需要自动微分）以确定要调用哪个内核。
PyTorch 需要实际启动内核。

从根本上说，这种开销来自于在每一步都能够执行不同操作的灵活性。如果你不需要这种灵活性，那么解决这种灵活性的方法之一是将其追踪出来，例如使用 jit.trace、FX 或 jax.jit。或者，你也可以在更低级别使用像 CUDA Graphs 这样的工具来实现。

不幸的是，这样做会失去灵活性。我期待的一种能够兼顾两者的方法是编写类似于“真正的”JIT 的工具，通过在 VM 级别进行内省来实现。有关详细信息，请参阅 TorchDynamo。

总结

如果你想加速深度学习系统，最重要的是要了解模型的瓶颈是什么。瓶颈决定了加速系统的最佳方法。

我经常看到研究人员和其他对加速 PyTorch 代码感兴趣的人盲目地尝试各种方法，而没有了解模型所处的性能瓶颈。

性能瓶颈	可行的解决方案
开销受限	追踪、操作融合、不要使用 Python、真正的 JIT :^)
带宽受限	操作融合
计算受限	使用 Tensor Cores、给英伟达更多钱

当然，可以说，用户需要考虑这些问题本身就反映了框架的失败。PyTorch 的编译器或分析 API 并不总是那么容易使用，尽管这是一个积极的关注领域。

无论如何，我发现理解系统的基本原理总是很有用，希望这对你也有帮助。

PS：如果你喜欢这篇文章，我将在 thonking.ai 上发布我未来的大部分文章。

致谢

感谢 Emily Shen、Qian Huang 和 EleutherAI 的成员阅读这篇博文的早期草稿并提供反馈。

BibTeX 引用

@article{he2022brrrrfromfirstprinciples,
  author={Horace He},
  title={Making Deep Learning Go Brrrr From First Principles},
  year={2022},
  url={https://horace.io/brrr_intro.html},
}

[1] 这可能与规格说明书上显示的 19.5 teraflops 不同。原因是 GPU 拥有更多专门用于融合乘加 (FMA) 指令的硬件。因此，对于完全通用的计算，A100 实际上只能达到 9.75 teraflops。 ↩︎

[2] 现在在 PyTorch 中以一种很好的方式计算 FLOPS 的方法有很多，请参阅 https://dev-discuss.pytorch.org/t/the-ideal-pytorch-flop-counter-with-torch-dispatch/505 ↩︎

在之前的文章中，我们介绍了 KV 缓存，这是一种常见的 LLM 推理过程优化方法，它可以使（自）注意力机制的计算需求随总序列长度（提示 + 生成的完成）线性增长，而不是二次增长。

更具体地说，KV 缓存通过在生成过程中计算时将过去标记的键和值张量存储（“缓存”）在 GPU 内存中，来避免在每个生成步骤中重新计算这些张量。

KV 缓存是一种折衷方案：我们用内存换取计算。在这篇文章中，我们将了解 KV 缓存可以增长到多大，它会带来什么挑战，以及解决这些挑战最常用的策略。

KV 缓存到底可以增长到多大？

这很简单：对于批次中每个序列的每个标记，我们需要为每个注意力层的每个注意力头存储两个向量张量（一个键张量和一个值张量），每个张量的尺寸为 。每个张量参数所需的存储空间取决于精度：全精度（FP32）为 4 字节/参数，半精度（BF16，FP16）为 2 字节/参数，8 位数据类型（INT8，FP8）为 1 字节/参数，等等。

假设 为批次大小， 为总序列长度（提示 + 完成）， 为解码器块/注意力层的数量， 为每个注意力层的注意力头的数量， 为注意力层的隐藏维度， 为精度。多头注意力（MHA）模型的 KV 缓存的每个标记的内存消耗（以字节为单位）为：

   

注意： 我们提醒，在 MHA 模型中，，但我们不会使用它来简化上面的公式。

因此，KV 缓存的总大小（以字节为单位）为：

   

KV 缓存的第一个挑战出现了：它随着批次大小线性增长，更重要的是随着总序列长度线性增长。由于它随着总序列长度增长，因此 KV 缓存的大小实际上是无界的，而我们的 GPU 内存显然是有限的。更糟糕的是，由于总序列长度无法提前知道，因此 KV 缓存的内存需求也是未知的，这使得内存管理特别具有挑战性。

让我们看看一些流行的 MHA 模型（表 1）的数字，包括 Meta 的 Llama-2 [1] 和 OPT [2]，MosaicML 的 MPT [3] 和 BigScience 的 BLOOM [4]：

模型	参数	层数	头数	隐藏维度
Llama-2–7B	7B	32	32	4096
OPT-1.3B	1.3B	24	16	768
MPT-7B	7B	24	32	4096
BLOOM-176B	176B	60	80	12288

表 1 – 流行的多头注意力 (MHA) 模型规格

假设参数存储在半精度（FP16，BF16）中，并选择一个较小的模型（Llama-2–7B）和一个较大的模型（BLOOM-176B）。对于 Llama-2–7B（分别为 BLOOM-176B），KV 缓存的内存消耗约为 0.5MB/标记（分别为 4MB/标记）。

让我们关注 Llama-2–7B。使用半精度，加载模型权重消耗约 14GB 内存，与缓存 28k 个标记的键和值相同。28k 个标记可能对应于 56 个长度为 512 的序列的批次，这并不算特别极端。

从上面的数字可以看出，KV 缓存的内存消耗会变得非常大，甚至可能超过加载大型序列模型权重所需的内存量。

现在让我们将这些数字与常用的 NVIDIA 数据中心 GPU 的内存容量进行比较（表 2）：

GPU	内存
A10	24GB
A100	40GB
H100	80GB
H200	96GB

表 2 – 常用于训练和/或服务 LLM 的 NVIDIA 数据中心 GPU 规格

让我们选择性价比比较高的 A10 GPU，坚持使用 Llama-2–7B 并计算最大 KV 缓存容量。加载模型权重后，还有 24-2×7=10 GB 可用于 KV 缓存，即总容量约为 20k 个标记，包括提示，这显然不允许在处理或生成长序列时，特别是当处理或生成长序列时，无法处理大量并发请求。

我们现在了解到，KV 缓存阻止我们处理或生成非常长的序列（即阻碍长上下文窗口）和/或处理大型批次，因此无法最大限度地提高硬件效率。

从这个角度来看，最大限度地提高处理能力意味着为 KV 缓存留出尽可能多的空间，这可以通过以下方法实现：

• 减少模型权重的内存占用（权重量化）
• 减少 KV 缓存的内存占用（见下文）
• 通过将模型在多个 GPU 上分片，以网络通信为代价，从多个设备中汇集内存（模型并行），或使用其他类型的存储，如 CPU 内存或磁盘（卸载）。

由于模型权重和不断增长的 KV 缓存必须在每次前向传递时加载，因此解码步骤涉及非常大的数据传输，正如我们将在后面的文章中看到，实际上是内存带宽受限的，即我们实际上花费更多的时间移动数据，而不是进行有用的工作，即计算。在这种情况下，延迟只能通过拥有更多内存带宽（即更好的硬件）或传输更少数据来改善。更小的模型权重和 KV 缓存释放了内存，可以容纳更多序列，因此能够提高吞吐量（和/或最大序列长度）。

在这方面，内存占用减少策略具有三重用途，因为它们允许我们提高硬件利用率，从而提高成本效益，同时降低延迟并提高吞吐量。

为什么我需要为输入标记付费？

在这一点上，你应该明白为什么你需要为输入和输出标记付费。一旦输入提示被处理，即在预填充阶段结束时，我们已经消耗了 GPU 内存（用于存储每个输入标记的键和值张量）和计算（用于将提示标记通过模型）。

让我们看看一些实际数字。假设 P 参数模型的前向传递的总 FLOPs 数量约为 2.P FLOPs/标记 [5]，使用 Llama-2-7B 处理提示消耗约 0.5 MB/标记的 GPU 内存（见上文）和约 14 GFLOPs/标记的 GPU 计算。对于一个 1000 个标记的提示（略小于两页），那就是约 500MB 的内存和 14 TFLOPs 的计算，而且我们还没有生成任何内容。

现在让我们看看所有可以减少 KV 缓存的内存占用的方法，方法是采用上面的公式，依次查看每个项：

减少批次大小？

在大多数情况下，我们不想减小批次大小，因为虽然它有助于减少 KV 缓存的内存占用，从而降低延迟，但它会降低硬件利用率，从而降低成本效益。在后面的文章中，我们将看到，相反，我们希望尽可能地增加批次大小。

减少对总序列长度的依赖？

不为序列中的所有标记存储键和值的一个原因是，我们明确选择在每次迭代时重新计算缺失的键和值，因为花费 FLOPs 比消耗 GPU 内存更划算（例如，因为我们是内存带宽受限的，在自回归阶段就是这样）。据我所知，这在实践中并不是我所知道的，所以我们不会深入研究这个方向。

另一个角度是，我们可以不必为模型不关注或很少关注的标记存储键和值。对于经过训练只关注序列一部分的模型（例如，使用 Mistral AI 的 Mistral-7B），或者作为内存消耗和模型精度之间的一种折衷方案，这可能是这种情况。让我解释一下。

像 Mistral-7B [6] 这样的模型经过训练不会关注整个序列。Mistral-7B 的注意力层实际上通过只关注最后（4096）个相邻标记来构建标记表示。这种注意力机制的变体称为滑动窗口注意力（SWA）或局部注意力。通过设计，局部注意力保证我们永远不会在 KV 缓存中为每个序列存储超过窗口大小（例如 4096）的张量对。

另一种方法是利用注意力层在序列中将注意力分散到标记上的方式中的模式。众所周知，注意力模块不成比例地且始终如一地将更多注意力分配给序列中的一小部分标记（图 1）。相反，许多标记始终对输出的贡献很小，那么为什么要费心存储它们的键和值呢？

通过丢弃这些标记，我们实际上将相应的注意力分数设置为零，并用更稀疏的矩阵来近似注意力矩阵。成功的近似将最大限度地减少近似误差，从而最大限度地减少对模型精度的影响（例如，使用困惑度来衡量）。

让我们看看过去几个月出现的一些方法，这些方法可以直接应用，无需任何重新训练或微调：StreamingLLM 框架、H2O（重击者预言机）、剪刀手和 FastGen。据我所知，目前还没有任何流行的 LLM 推理框架支持它们。

针对使用有限长度上下文窗口训练的模型，StreamingLLM 框架 [7] 基于这样一个观察结果：初始标记会收集大量的注意力。因此，该框架通过只在缓存中保留最开始的几个位置标记（“汇聚标记”）和最后一个相邻标记（局部注意力）来构建一个滑动窗口。因此，StreamingLLM KV 缓存的长度是固定的，既有固定部分（通常为 1 到 4 个标记），也有滑动部分。

类似的 H2O [8] 和剪刀手 [9] 方法明确地旨在通过设置一个最大缓存标记数（预算）并通过在每次达到缓存预算时丢弃标记来压缩 KV 缓存。H2O 算法一次只丢弃一个标记，而剪刀手则根据目标压缩率（例如，减少 30% 的 KV 缓存大小）丢弃尽可能多的标记。

这两种方法都基于这样一个观察结果：在给定步骤中具有影响力的标记（“关键标记”或“重击者”）在未来步骤中仍然具有影响力（剪刀手作者称之为“重要性持久性假设”）。换句话说，我们可以确保被丢弃的低影响力标记在未来步骤中将仍然相对被忽略，因此可以安全地丢弃。

这两种算法的一个关键方面显然是缓存驱逐策略。剪刀手只保留最新的标记和历史窗口内注意力分数最高的标记。H2O 丢弃累积注意力分数最低的标记，因此只保留在迭代过程中始终获得高注意力分数的标记。这两个作者团队都表明，他们的算法可以实现高达 80% 的 KV 缓存大小减少，而模型精度损失可以忽略不计。

FastGen 方法 [10]（不要与无关的 DeepSpeed-FastGen 混淆）仍然基于注意力模式，但采用另一种方法，即不设置缓存预算，而是设置注意力矩阵的最大近似误差，因此侧重于保持模型精度。

FastGen 是一种两步法：首先，在预填充阶段结束时对模型的注意力层进行分析，以确定允许满足误差目标的压缩策略集。与其他方法一样，它假设所识别的注意力模式将在未来的生成步骤中保持不变。压缩策略包括：保留特殊标记、保留标点符号标记、保留最后一个相邻标记（局部注意力）等等（图 2）。如果误差目标过于严格，无法满足，FastGen 将回退到常规的 KV 缓存。然后，在每个生成步骤中将所选的压缩策略应用于 KV 缓存。

图 2 – FastGen 论文中压缩策略集的示例：特殊标记（绿色）+ 标点符号标记（橙色）+ 局部注意力（蓝色）。被丢弃的标记用灰色表示。

请注意，与其他方法不同，FastGen 为每个提示构建一个量身定制的压缩策略。FastGen 作者表明，对于给定的 KV 缓存压缩率，他们比 H2O 和剪刀手更好地保持模型精度。

在任何情况下，打破对不可预测的总序列长度的依赖都是一种解脱，因为它允许为每个序列分配一个内存预算，从而大大简化内存管理。由于数据传输是延迟的主要贡献者，因此没有一个随着序列长度线性增长的 KV 缓存，尤其是在较长的序列长度下，可以带来惊人的加速。

减少层数？

这里没有太多可以获得的。较小的模型通常具有较少的层（表 4），因此如果较小的模型在你的用例中表现良好，只需选择它即可。

模型	参数	层数	头数	隐藏维度
Llama-2–7B	7B	32	32	4096
Llama-2–13B	13B	40	40	4096
Llama-2–70B	70B	60	64	4096
Llama-2–13B	13B	80	80	4096

表 4 – Llama-2 模型规格

减少注意力头的数量？

由于对于给定的模型架构，模型大小主要由层数和头数控制，因此减少头数可能意味着选择一个更小的模型（见表 4）。

但是，如果我们仔细观察，我们会注意到，我们只需要减少键和值头的数量，查询头的数量不会影响 KV 缓存的大小。这正是多查询注意力（MQA）[11] 和分组查询注意力（GQA）[12] 架构背后的理念。这些多头注意力（MHA）变体的唯一动机是减少 KV 缓存的大小。

MQA 最早于 2019 年被引入。在 MQA 中，所有查询头共享相同的单个键和值头。换句话说，所有查询头使用相同的键计算它们的注意力分数，所有头的输出使用相同的值（但不是相同的注意力分数）计算（图 3）。

图 3 – 多头注意力（上）与多查询注意力（下）（两个注意力头）

然而，剥离所有头对于较大的模型来说相对更激进。例如，从 64 个头减少到 1 个头，与从 32 个头减少到 1 个头相比，模型的表示能力的削减更大。GQA 通过提供一种折衷方案来解决这个问题：我们不是让所有查询头共享相同的唯一 KV 头，而是将它们分成 g 个查询头的组，同一个组的查询头共享相同的唯一 KV 头。换句话说，我们不是从 缩减到 1 个 KV 头，而是将 KV 头的数量从 缩减到 。

从这个角度来看，MHA 和 MQA 都是 GQA 的特例（分别为 和 ）。QGA 允许在两个极端情况（MHA 和 MQA）之间更平滑地导航模型精度/KV 缓存大小（与延迟和吞吐量都相关）的折衷方案。

考虑到这个新的参数 g，KV 缓存大小公式变为：

   

在实践中，MQA/GQA 架构已被 Google Research 的 PaLM [13]，TII 的 Falcon [14] 模型，Meta 的 Llama-2 [1]（仅 70B）和 Mistral AI 的 Mistral-7B [7]（表 5）显著地实现。

模型	参数	层数	头数	隐藏维度	注意力类型
PaLM	540B	52	64	16384	MQA
Falcon	40B	40	64	5120	MQA
Llama-2–70B	70B	60	64	4096	GQA
Mistral-7B	7B	24	32	4096	GQA

表 5 – 使用 MQA 或 GQA 的模型系列

注意力头的隐藏维度？

再次强调，如果你不准备选择另一个模型，这里没有太多可以获得的。根据模型系列的不同，头隐藏维度在不同模型大小之间可能是恒定的（例如 Llama-2，Falcon），因此选择同一个系列中较小的变体不会有帮助。

使用更少的字节数/参数？

量化 KV 缓存确实是一种可以大幅减少其大小的好方法。但是，像 AWQ [15] 或 GPTQ [16] 这样的仅权重量化算法从定义上来说不会有帮助。只有量化权重和“激活”（即任何不是权重的东西）的算法，例如 LLM.int8()[17] 或 SmoothQuant [18]，才会产生量化的 KV 缓存。

请注意，在权重和激活上都起作用的量化算法的意图之一是在较低的精度下执行计算密集型矩阵乘法。如果像训练期间那样是计算受限的，这会带来性能提升，但正如我们将在后面的文章中看到，推理的自回归阶段实际上是内存带宽受限的，因此能够更快地计算不会带来太多价值。

我们实际上只对内存占用的减少感兴趣，因为它意味着更少的数据传输。从这个角度来看，像 LLM.int8() 或 SmoothQuant 这样的量化算法有点过分：在将缓存的张量移动到 GPU 内存之前对其进行量化，并在从 GPU 内存中获取相同的张量后对其进行反量化（以额外的开销为代价）应该就足够了。

一些 LLM 推理系统已经包含了这样的 KV 缓存量化功能。例如，FlexGen [19] 将 KV 缓存和模型权重都量化并存储在 4 位数据格式中。NVIDIA TensorRT-LLM 能够将 KV 缓存量化到 8 位数据格式（INT8 或 FP8）中。流行的 vLLM 框架从 0.3.0 版本开始也支持 KV 缓存（FP8）量化。由于量化是在每次迭代时动态执行的，因此不需要校准步骤。

高效内存管理的重要性

到目前为止，我们隐含地假设内存没有浪费：所有保留的内存都用于存储标记，所有可用的内存都可以被保留。在实践中，简单的内存管理策略会导致大量内存被浪费（PagedAttention 论文 [20] 表明，实际的有效内存利用率可能低至 20%，即 80% 的浪费！）：

• 由于请求的总序列长度事先未知，我们可以保留能够容纳最大序列长度的连续内存块。这种分配中很大一部分肯定永远不会被使用，并且由于无法用于其他请求，因此被浪费了（内部内存碎片）。
• 即使序列长度事先已知，由于内存是逐渐消耗的，但内存块是为请求的生命周期保留的，因此较短的请求无法使用仍然未使用的内存块。
• 如果我们使用像束搜索这样的解码策略，为每个请求生成多个序列，那么多个候选序列实际上可以部分共享它们的 KV 缓存。如果我们不考虑这种情况，我们不可避免地会通过存储本可以共享的重复 KV 条目来浪费内存。

这些缺点正是现在流行的 PagedAttention 算法旨在解决的问题。PagedAttention 分配固定大小且相对较小的内存块，称为块。每个块可以包含固定数量的标记，如果需要，可以跨不同的请求共享。按需分配和较小的块大小减轻了内部内存碎片，而相同大小的块消除了外部内存碎片。

总的来说，PagedAttention 实现了 KV 缓存内存的近乎零浪费（低于 4% [21]）。以前被浪费的内存现在可以用来容纳更多请求，从而提高吞吐量。PagedAttention 推出时，吞吐量提高的数字与当时内存浪费的程度一样惊人。

PagedAttention 最初由 vLLM 推理系统实现，但现在所有主要的推理框架都支持它（例如 HuggingFace TGI、NVIDIA TensorRT-LLM、LMDeploy TurboMind 等）。

PagedAttention 未涵盖的另一个可能的优化是在请求之间重用键值缓存。这将适用于提示共享公共前缀的情况，这种情况在像聊天和代理这样的多轮用例中很常见，或者在使用提示模板时（图 4）。

图 4 – SGLang 论文中 KV 缓存共享示例（多轮聊天），总共四个生成请求。蓝色框表示可共享的提示部分。

能够在请求之间重用 KV 缓存将能够显著降低延迟（尤其是第一个标记的延迟）和吞吐量（通过大大减少具有共享前缀的并发请求的内存占用）。

LMSYS SGLang 论文 [22] 中介绍的 RadixAttention 技术实现了这种 KV 缓存重用。

RadixAttention 算法不是在完成生成请求后丢弃 KV 缓存，而是将其保留在 GPU 内存中，并向一个专门的数据结构（基数树）添加一个新条目，该结构将标记序列映射到它们的 KV 缓存张量。当一个新的请求进来时，调度程序使用基数树进行前缀匹配。如果有缓存命中，调度程序将重用缓存的 KV 张量来满足请求。

由于 GPU 内存有限，缓存的 KV 张量不能永远保留。因此，RadixAttention 算法包含一个驱逐策略（例如，最近最少使用 (LRU) 驱逐策略）。最佳缓存重用可能与先到先得的调度程序不兼容。因此，RadixAttention 带有一个修改后的调度程序，该调度程序优先考虑与缓存的前缀匹配的请求（缓存感知调度）。

注意： PagedAttention 和 RadixAttention 的命名有点误导，因为与人们可能认为的相反，它们不是模型注意力层的优化（如 FlashAttention），而是在模型服务器级别运行（它们帮助服务应用程序更好地管理主机上的 KV 缓存）。

如果 GPU 内存不足，为什么不“只”使用多个 GPU？或者卸载到 CPU 内存甚至磁盘？

这些是两种不同的但有效的方案。

首先关于卸载到更丰富但更慢的存储器（CPU 内存和磁盘）。并非所有推理框架都支持此功能，让我们列举 HuggingFace Accelerate、DeepSpeed-Inference 和更先进的 FlexGen。由于它涉及使用更慢的存储器，因此卸载会以较大的延迟损失为代价，因此对于对延迟敏感的用例来说，显然不应优先考虑此选项。卸载系统通常用于面向吞吐量的用例，例如离线批处理。

关于使用多个 GPU（对于较大的模型来说不可避免），将模型在多个设备上分片可以释放内存压力，因为可以从聚合的内存容量和内存带宽中获益。

如果选择管道并行 [23]，模型和 KV 缓存都将在层维度上分片。如果选择张量并行 [24]（更常见于推理），KV 缓存将在头维度上分片。请注意，MQA 在这种设置下效率很低：由于我们无法跨多个设备分片单个头，因此 KV 缓存必须在所有设备上复制，从而失去了 MQA 的优势。对于实现 MQA 的模型，另一种选择是跨批次大小维度分片 KV 缓存 [25]。

在任何情况下，以上所有情况都假设单个主机，我们仍然受限于我们可以得到的最大多 GPU 实例的存储容量。据我所知，目前还没有推理框架支持多主机模型并行。如果我们能够将模型和 KV 缓存都分片到多个主机上，那么我们可以处理的可用内存量和最大序列长度将实际上变得无限。这是 Infinite-LLM 论文 [26] 旨在通过引入一种新的分布式注意力算法 (DistAttention) 和调整 Ray 框架来构建多主机分布式 KV 缓存管理和调度系统 (DistKV-LLM) 来解决的问题。

总结

在这篇文章中，我们学习了选择 KV 缓存会带来额外的挑战。多头注意力 (MHA) 模型的 KV 缓存确实会消耗大量 GPU 内存，大约为 1MB/标记，并且很容易超过模型权重的内存占用。

鉴于 GPU 内存的有限性，KV 缓存内存压力导致了许多不同方向的举措：新的注意力架构（MQA、GQA、SWA）、缓存压缩策略（H2O、剪刀手、FastGen）、高效内存管理（PagedAttention、RadixAttention）、量化和存储容量扩展（卸载系统、单主机和多主机模型并行）。

正如我们将在后面的文章中看到，减少 KV 缓存的大小至关重要，不仅因为 GPU 内存有限，而且因为数据移动量实际上是每个自回归步骤的延迟的主要贡献者，因此也是整个生成过程的延迟的主要贡献者。

在下一篇文章中，我们将看看可能影响模型延迟和吞吐量的不同类型的瓶颈。我们下次再见！

参考文献：

[1]: Llama 2: Open Foundation and Fine-Tuned Chat Models (Touvron 等人，2023)
[2]: OPT: Open Pre-trained Transformer Language Models (Zhang 等人，2022)
[3]: MPT-7B（2023 年 5 月）和 MPT-30B（2023 年 6 月）的发布博客文章
[4]: BLOOM: A 176B-Parameter Open-Access Multilingual Language Model (BigScience，2023)
[5]: Scaling Laws for Neural Language Models (Kaplan 等人，2020)
[6]: Mistral 7B (Jiang 等人，2023)
[7]: Efficient Streaming Language Models with Attention Sinks (Xiao 等人，2023) + GitHub 仓库
[8]: H_2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models (Zhang 等人，2023) + GitHub 仓库
[9]: Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time (Liu 等人，2023)
[10]: Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs (Ge 等人，2023)
[11]: Fast Transformer Decoding: One Write-Head is All You Need (Shazeer，2019)
[12]: GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (Ainslie 等人，2023)
[13]: PaLM: Scaling Language Modeling with Pathways (Chowdhery 等人，2022)
[14]: The Falcon Series of Open Language Models (Almazrouei 等人，2023)
[15]: AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration (Lin 等人，2023) + GitHub 仓库
[16]: GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers (Frantar 等人，2022) + GitHub 仓库
[17]: LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale (Dettmers 等人，2022) + GitHub 仓库
[18]: SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models (Xiao 等人，2022) + GitHub 仓库
[19]: FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU (Sheng 等人，2023) + GitHub 仓库
[20]: Efficient Memory Management for Large Language Model Serving with PagedAttention (Kwon 等人，2023) + GitHub 仓库
[21]: vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttention (Kwon 等人，2023)
[22]: Efficiently Programming Large Language Models using SGLang (Zheng 等人，2023) + 博客文章
[23]: GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism (Huang 等人，2018)
[24]: Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (Narayanan 等人，2021)
[25]: Efficiently Scaling Transformer Inference (Pope 等人，2022)
[26]: Infinite-LLM: Efficient LLM Service for Long Context with DistAttention and Distributed KVCache (Lin 等人，2024)

大型语言模型（LLM）在自然语言处理领域取得了巨大成功，但其高昂的计算成本和庞大的内存需求也成为了限制其应用的瓶颈。矩阵乘法（MatMul）是LLM中最主要的计算操作，其占用了大部分的计算时间和内存资源。为了解决这一问题，来自加州大学圣克鲁兹分校的研究人员提出了一种全新的，可扩展的无矩阵乘法语言模型（MatMul-free LM），该模型在保持强大性能的同时，彻底消除了所有矩阵乘法操作。

为什么矩阵乘法如此重要？

矩阵乘法在神经网络中无处不在，从密集层到卷积层，再到自注意力机制，都离不开矩阵乘法。这主要是因为现代图形处理单元（GPU）对矩阵乘法操作进行了高度优化。通过利用CUDA和cuBLAS等线性代数库，矩阵乘法可以被高效地并行化和加速。这使得AlexNet在2012年ImageNet竞赛中取得了胜利，并推动了深度学习的快速发展。

然而，矩阵乘法操作也带来了巨大的计算成本和内存消耗。在训练和推理阶段，矩阵乘法通常占用了绝大部分的执行时间和内存访问。因此，研究人员一直在探索用更简单的操作来替代矩阵乘法。

现有方法的局限性

目前，替代矩阵乘法的方法主要有两种：

1. 用基本运算替代矩阵乘法： 例如AdderNet用带符号的加法来替代卷积神经网络中的乘法。但AdderNet主要针对计算机视觉任务，在语言建模方面效果不佳。
2. 二值化或三值化： 将矩阵中的元素量化为二进制或三进制，从而将矩阵乘法简化为简单的加减运算。这种方法可以应用于激活值或权重。例如，脉冲神经网络（SNN）使用二值化的激活值，而二值化神经网络（BNN）和三值化神经网络（TNN）使用量化的权重。

近年来，BitNet等语言模型证明了量化方法的可扩展性，将所有密集层权重替换为二进制或三进制值，支持高达30亿个参数。然而，BitNet仍然保留了自注意力机制，而自注意力机制仍然依赖于昂贵的矩阵乘法。

MatMul-free LM的创新之处

为了彻底消除LLM中的矩阵乘法，研究人员提出了MatMul-free LM，该模型利用了密集层中的加法运算和自注意力机制中的逐元素哈达玛积。

1. 三值化权重： 类似于BNN，MatMul-free LM将密集层中的权重限制为{-1, 0, +1}，从而将矩阵乘法转换为简单的加减运算。

2. 无矩阵乘法线性门控循环单元（MLGRU）： 为了消除自注意力机制中的矩阵乘法，研究人员对门控循环单元（GRU）进行了优化，使其仅依赖于逐元素乘法。

3. 无矩阵乘法门控线性单元（GLU）： MatMul-free LM使用GLU作为通道混合器，并将其中的密集层替换为三值化权重，从而消除了矩阵乘法。

MatMul-free LM的优势

MatMul-free LM具有以下优势：

• 计算效率更高： 消除了矩阵乘法操作，大幅减少了计算时间。
• 内存需求更低： 三值化权重减少了模型的内存占用。
• 硬件友好： 更适合在FPGA等专用硬件上实现。

实验结果

研究人员对MatMul-free LM进行了实验，并将其与Transformer++模型进行了比较，结果表明：

• MatMul-free LM在性能上与Transformer++相当，甚至在某些情况下表现更佳。
• MatMul-free LM在训练和推理阶段的效率都更高，内存占用和延迟都更低。
• MatMul-free LM在FPGA上的实现也取得了成功，其效率接近人脑。

未来展望

MatMul-free LM的出现，为构建更高效、更节能的LLM开辟了新的道路。随着LLM在各种平台上的应用不断扩展，MatMul-free LM将成为构建高效、可扩展的LLM的重要方向。

参考文献:

[1] A. Krizhevsky, I. Sutskever, and G. E. Hinton. Imagenet classification with deep convolutional neural networks. In Advances in neural information processing systems, pp. 1097–1105, 2012.

[2] Y. LeCun, Y. Bengio, and G. Hinton. Deep learning. Nature, 521(7553):436–444, 2015.

[3] S. Han, H. Mao, and W. J. Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. In International Conference on Learning Representations, 2016.

[4] S. B. Furber. Neuromorphic engineering. The MIT Press, 2016.

[5] G. Indiveri, B. Linares-Barranco, R. Legenstein, D. Chicca, G. Indiveri, B. Linares-Barranco, R. Legenstein, D. Chicca, and A. Hamilton. Neuromorphic silicon. Springer, 2011.

[6] T. Masquelier, S. Thornton, S. B. Furber, and J. V. Pulvermüller. A spiking neural network model of word recognition in the human brain. PLoS computational biology, 10(12):e1003974, 2014.

[7] I. Hubara, M. Courbariaux, D. Soudry, R. El-Yaniv, and Y. Bengio. Binarized neural networks. In Advances in neural information processing systems, pp. 4107–4115, 2016.

[8] M. Courbariaux, Y. Bengio, and J.-P. Salinas. Binaryconnect: Training deep neural networks with binary weights during backpropagation. In Advances in neural information processing systems, pp. 4107–4115, 2015.

[9] R. Zhu, Y. Zhang, E. Sifferman, T. Sheaves, Y. Wang, D. Richmond, P. Zhou, and J. K. Eshraghian. Scalable MatMul-free Language Modeling. arXiv preprint arXiv:2406.02528, 2024.

[10] L. Pei, S. Li, S. Zhang, J. Li, and S. Liu. BitNet: A Billion-Parameter Binary and Ternary Neural Network for Language Modeling. arXiv preprint arXiv:2302.03633, 2023.

[11] L. Pei, S. Li, S. Zhang, J. Li, and S. Liu. BitNet: A Billion-Parameter Binary and Ternary Neural Network for Language Modeling. arXiv preprint arXiv:2302.03633, 2023.

[12] T. Chen, Z. Du, N. Sun, J. Wang, C. Wu, Y. Chen, H. Li, Y. Tang, Y. Wang, and X. Lin. Training deep neural networks with 8-bit floating point numbers. In Advances in Neural Information Processing Systems, pp. 6083–6092, 2018.

[13] K. Cho, B. van Merriënboer, C. Gulcehre, D. Bahdanau, F. Bougares, H. Schwenk, and Y. Bengio. Learning phrase representations using RNN encoder-decoder for statistical machine translation. In Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 1724–1734, 2014.

[14] D. M. K. Pramanik, S. Bhattacharyya, and P. Das. Ternary BERT: Low-Precision BERT for Resource-Constrained Devices. arXiv preprint arXiv:2004.06633, 2020.

[15] Y. Sun, Y. Zhang, Z. Liu, Y. Liu, and J. Tang. Quantized BERT: Efficient BERT for Resource-Constrained Devices. arXiv preprint arXiv:1910.04432, 2019.

[16] Y. Wang, Y. Zhang, Y. Sun, J. Tang, and Z. Liu. Incremental Network Quantization: Towards Lossless CNNs with Low-Precision Weights. arXiv preprint arXiv:2002.08150, 2020.

[17] J. Zhou, Z. Sun, A. Zou, Q. Liu, and Y. Gong. Training Low-Precision Deep Neural Networks via Quantization-Aware Training. arXiv preprint arXiv:1905.04893, 2019.

[18] C. Lee, S. Lee, H. Kim, and J. Shin. Spikformer: Spiking Neural Networks for Efficient Transformer. arXiv preprint arXiv:2103.13518, 2021.

[19] C. Lee, S. Lee, H. Kim, and J. Shin. Spikformer: Spiking Neural Networks for Efficient Transformer. arXiv preprint arXiv:2103.13518, 2021.

[20] A. S. M. A. Saleh, A. A. M. Al-Jumaily, and A. Al-Ani. Spike-Driven Transformer for Image Classification. arXiv preprint arXiv:2203.08669, 2022.

[21] S. M. A. Saleh, A. A. M. Al-Jumaily, and A. Al-Ani. Spike-Driven Transformer for Image Classification. arXiv preprint arXiv:2203.08669, 2022.

[22] X. Li, L. Huang, J. Li, and Y. Chen. Spiking-BERT: A Spiking Neural Network for Sentiment Analysis. arXiv preprint arXiv:2106.07442, 2021.

[23] M. A. Saleh, A. A. M. Al-Jumaily, and A. Al-Ani. SpikeBERT: A Spiking Neural Network for Sentiment Analysis. arXiv preprint arXiv:2110.03458, 2021.

[24] A. Parmar, A. Vaswani, J. Uszkoreit, L. Kaiser, N. Shazeer, A. Ku, and D. Tran. Image transformer. In International Conference on Machine Learning, pp. 3887–3896, 2018.

[25] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998–6008, 2017.

[26] S. Liu, Z. Chen, Y. Li, Z. Liu, and W. Zhang. Mamba: A Low-Resource and Efficient Transformer. arXiv preprint arXiv:2106.02256, 2021.

[27] D. Bahdanau, K. Cho, and Y. Bengio. Neural machine translation by jointly learning to align and translate. In International Conference on Learning Representations, 2015.

[28] P. Ramachandran, B. Zoph, and Q. Le. Swish: A self-gated activation function. arXiv preprint arXiv:1710.05941, 2017.

[29] A. Courbariaux, R. Bengio, and J.-P. Salinas. Binaryconnect: Training deep neural networks with binary weights during backpropagation. In Proceedings of the 2015 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 4107–4115, 2015.

[30] I. Hubara, M. Courbariaux, D. Soudry, R. El-Yaniv, and Y. Bengio. Binarized neural networks. In Advances in Neural Information Processing Systems, pp. 4107–4115, 2016.

[31] S. Hochreiter and J. Schmidhuber. Long short-term memory. Neural computation, 9(8):1735–1780, 1997.

[32] S. Z. Li, Y. H. Zou, T. Y. Liu, and B. Zhang. Linear recurrent unit. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pp. 1780–1789, 2018.

[33] A. Radford, J. Wu, R. Sutskever, and I. Sutskever. Learning to generate text with recurrent neural networks. arXiv preprint arXiv:1701.07875, 2017.

[34] B. Wang. RWKV: A New Type of Recurrent Neural Network. arXiv preprint arXiv:2210.00078, 2022.

[35] L. Pei, S. Li, S. Zhang, J. Li, and S. Liu. BitNet: A Billion-Parameter Binary and Ternary Neural Network for Language Modeling. arXiv preprint arXiv:2302.03633, 2023.

[36] A. Touvron, T. Bernard, J. Caillard, G. Lample, and E. Lemoine. Llama: Open and Efficient Large Language Models. arXiv preprint arXiv:2302.09492, 2023.

[37] A. Touvron, T. Bernard, J. Caillard, G. Lample, and E. Lemoine. Llama: Open and Efficient Large Language Models. arXiv preprint arXiv:2302.09492, 2023.

[38] A. Touvron, T. Bernard, J. Caillard, G. Lample, and E. Lemoine. Llama: Open and Efficient Large Language Models. arXiv preprint arXiv:2302.09492, 2023.

[39] Mistral AI. Mistral AI: A New Open Source LLM. https://mistral.ai, 2023.

[40] Y. Bengio, N. L. C. da Silva, J. S. S. Souza, and P. Vincent. Neural networks for speech recognition. The Journal of the Acoustical Society of America, 131(3):1852–1865, 2012.

[41] M. Courbariaux, Y. Bengio, and J.-P. Salinas. Binaryconnect: Training deep neural networks with binary weights during backpropagation. In Proceedings of the 2015 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 4107–4115, 2015.

近年来，变形金刚（Transformers）凭借其强大的能力，在自然语言处理领域取得了巨大成功，成为深度学习的主流架构之一。然而，近年来，结构化状态空间模型（SSMs），例如Mamba，在小型到中型模型上展现出与变形金刚相媲美甚至超越的性能。

本文将深入探讨变形金刚和状态空间模型之间的深层联系，揭示它们之间的结构化状态空间对偶性（SSD），并以此为基础，提出了一种名为Mamba-2的新型架构，它不仅在语言建模方面与变形金刚竞争，而且训练速度更快，效率更高。

变形金刚和 SSM 的深层联系

变形金刚的核心是注意力机制，它通过对序列中每个元素之间的关系进行评分，让每个元素能够“关注”其他元素。然而，注意力机制的计算复杂度较高，特别是当序列长度较长时，训练和推理的效率会急剧下降。

另一方面，状态空间模型则通过递归的方式来处理序列，它们将输入序列映射到一个隐含的状态空间，并通过对状态空间进行操作来生成输出。与注意力机制相比，状态空间模型的计算复杂度更低，特别是在处理长序列时。

那么，变形金刚和 SSM 之间到底有什么联系呢？本文揭示了它们之间的结构化状态空间对偶性（SSD），它通过结构化矩阵的概念，将两者联系起来。

结构化矩阵：高效计算的桥梁

结构化矩阵是指那些具有亚二次参数和乘法复杂度的矩阵。它们可以被压缩表示，并通过高效的算法进行操作。

本文证明了状态空间模型等价于一类被称为半可分离矩阵（Semiseparable Matrices）的结构化矩阵。这种等价性揭示了状态空间模型的全新特性，并为其设计高效算法提供了新的思路。

结构化掩码注意力：线性注意力的扩展

线性注意力（Linear Attention）是注意力机制的一种变体，它通过将注意力矩阵分解为两个矩阵的乘积，并利用矩阵乘法的结合律，将注意力机制的计算复杂度从二次降到线性。

本文通过张量收缩的概念，对线性注意力进行了重新推导，并将其扩展到更一般的结构化掩码注意力（SMA）。SMA 允许使用任何具有亚二次乘法复杂度的结构化矩阵作为掩码，从而进一步提高了注意力的效率。

结构化状态空间对偶性：将 SSM 和 SMA 联系起来

本文的核心结论是，状态空间模型和结构化掩码注意力之间存在着一种结构化状态空间对偶性（SSD）。

具体而言，当状态空间模型中的状态转移矩阵具有标量乘以单位矩阵的结构，而结构化掩码注意力中的掩码矩阵为1-半可分离矩阵时，它们就具有相同的线性形式和二次形式。

这种对偶性意味着，我们可以将设计用于注意力的算法和优化技术应用到状态空间模型中，反之亦然。

Mamba-2 架构：将 SSD 应用于语言建模

基于 SSD 框架，本文提出了一种名为Mamba-2的新型架构，它将 SSD 层作为核心层，并结合了其他优化技术，例如并行参数投影、额外归一化、多头模式等。

Mamba-2 在语言建模方面展现出与变形金刚相媲美的性能，同时训练速度更快，效率更高。

总结

本文揭示了变形金刚和状态空间模型之间的深层联系，并提出了一个名为 SSD 的理论框架，它通过结构化矩阵的概念，将两者联系起来。SSD 框架为设计更高效的序列模型提供了新的思路，并为 Mamba-2 架构的提出提供了理论基础。

参考文献

[1] Dao, Tri, and Albert Gu. “Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality.” arXiv preprint arXiv:2305.11024 (2023).

---

Q: 这篇论文试图解决什么问题？

A: 这篇论文的主要目标是建立结构化状态空间模型（SSMs）和注意力机制（Transformers中的一个关键组件）之间的理论联系。具体来说，它试图解决以下问题：

1. 模型效率问题：Transformers在处理长序列时面临效率问题，因为它们的注意力机制在序列长度上呈二次方增长。而SSMs提供了一种线性复杂度的序列处理方式，但它们的发展似乎与Transformers的优化和硬件适应性改进相脱节。
2. 理论联系缺失：SSMs和Transformers的优化工作在社区中相对独立，缺乏理论上的联系和相互借鉴，这使得SSMs相比Transformers更难以理解和实验。
3. 算法和系统优化转移：论文希望将为Transformers开发的算法和系统优化转移到SSMs中，以构建性能更好、序列长度扩展性更高的基础模型。
4. 设计新架构：通过理论联系，论文提出了一种新的架构（Mamba-2），其核心层是对Mamba的SSM的选择性模型的改进，旨在提高计算速度，同时在语言建模方面与Transformers保持竞争力。
5. 系统优化：论文还探讨了如何将SSMs与Transformers相连接，从而利用为Transformers开发的系统优化工作，例如张量并行性（TP）、序列并行性等。

总的来说，论文试图通过建立SSMs和注意力机制之间的理论桥梁，促进两者的相互理解和优化，推动更高效、更可扩展的深度学习模型的发展。

Q: 有哪些相关研究？

A: 论文中提到了多个与结构化状态空间模型（SSMs）、注意力机制以及相关算法和系统优化相关的研究。以下是一些主要的相关研究：

1. Transformers：特别是解码器仅模型（如GPT和Llama），这些模型是深度学习在语言建模方面成功的主要驱动力。
2. Mamba：一种结构化SSM，它在小到中等规模上展示了与Transformers相匹配或更好的性能。
3. Linear Attention (LA) 框架：由Katharopoulos等人提出，它展示了自回归注意力和线性RNNs之间的联系。
4. FlashAttention-2：由Dao提出，这是一种优化的Softmax注意力实现。
5. Megatron：Shoeybi等人提出的一种模型并行性技术，用于训练大型Transformer模型。
6. S4：Gu等人提出的结构化SSM，用于长序列建模。
7. GateLoop：Katsch提出的一个模型，它使用输入依赖的衰减因子。
8. Gated Linear Attention (GLA)：Yang等人提出的一个模型，它在线性注意力中引入了数据依赖的门控机制。
9. HGRN：Qin等人提出的一个RNN模型，它结合了输入依赖的门控和状态扩展。
10. Griffin：De等人提出的一个模型，它结合了门控线性递归和局部注意力。
11. xLSTM：Beck等人提出的一个改进的长短期记忆网络（LSTM）模型。
12. RWKV：Peng等人提出的一个基于不同线性注意力近似的RNN模型。

这些研究为论文提出的结构化状态空间对偶（SSD）框架提供了理论基础和技术支持。论文通过与这些相关工作的比较和联系，展示了SSD框架的优势和潜在的应用前景。

Q: 论文如何解决这个问题？

A: 论文通过提出一个名为结构化状态空间对偶（Structured State Space Duality, SSD）的框架来解决上述问题。SSD框架旨在建立结构化状态空间模型（SSMs）和注意力机制之间的理论联系，并开发出新的算法和架构设计。具体来说，论文采取了以下几个步骤：

1. 理论联系：展示了SSMs和注意力机制实际上是密切相关的，通过研究一类被称为结构化半可分矩阵的特定矩阵类来建立联系。
2. 状态空间模型（SSM）与结构化矩阵的等价性：证明了SSMs可以表示为结构化矩阵，特别是半可分矩阵，这揭示了SSMs的新属性和算法。
3. 线性注意力理论的改进：扩展了线性注意力的理论，提供了通过张量收缩语言的递归形式的新证明，并将线性注意力推广到新的结构化掩蔽注意力（Structured Masked Attention, SMA）。
4. SSMs与SMA的连接：展示了SSMs和SMA之间存在大量的交集，并且这些模型是彼此的对偶，具有SSM类的线性和注意力类的二次形式。
5. 高效算法：基于SSD框架，提出了新的算法来计算SSMs，这些算法利用了线性SSM递归和二次对偶形式，实现了在训练和推理计算、内存使用以及利用现代硬件上的矩阵乘法单元方面的最优权衡。
6. 架构设计：利用SSD框架，将Transformers的优化和系统技术转移到SSMs中，设计了新的架构Mamba-2，该架构在语言建模方面与Transformers竞争，同时在序列长度上具有更好的扩展性。
7. 系统优化：SSD框架允许利用为Transformers开发的系统优化工作，例如张量并行性（TP）、序列并行性等，来提高SSMs的效率。
8. 实验验证：通过实验验证了Mamba-2在语言建模、训练效率和多查询关联记忆任务上的性能。

通过这些步骤，论文不仅在理论上建立了SSMs和注意力机制之间的联系，还开发了新的算法和架构，使得SSMs能够更高效地运行，并且能够与Transformers相竞争。

Q: 论文做了哪些实验？

A: 论文中进行了一系列的实验来验证所提出的Mamba-2架构和SSD算法的有效性。以下是实验的主要内容：

1. 合成关联记忆任务（Associative Recall）：
   • 使用了一个更难版本的MQAR（Multi-Query Associative Recall）任务，该任务要求模型记忆多个键值对关联，并在提示时正确完成。
   • 实验结果表明，Mamba-2在所有设置中表现良好，即使在控制状态大小（N=16）的情况下，也比Mamba-1有显著提升。
2. 语言建模（Language Modeling）：
   • 在标准的自回归语言建模任务上训练和评估Mamba-2架构，并与其他架构进行比较。
   • 使用了Pile数据集，并遵循了GPT-3的规格进行训练。
3. 扩展性规律（Scaling Laws）：
   • 对于不同大小的模型，从约125M到约1.3B参数，进行了扩展性规律的实验。
   • 实验结果表明，Mamba-2在性能（困惑度）、理论FLOPs和实际墙钟时间上与Mamba和Transformer++相当或更好。
4. 下游评估（Downstream Evaluations）：
   • 在多种流行的零样本评估任务上，比较了Mamba-2与其他知名开源模型的性能。
   • 实验结果显示，Mamba-2在各个任务上的表现与Pythia相当，甚至在某些情况下超过了Mamba。
5. 混合模型（Hybrid Models）：
   • 探索了将SSD层与MLP和注意力层结合的不同方式，以了解每种组件的好处。
   • 实验发现，大约10%的层是注意力层时表现最佳。
6. 速度基准测试（Speed Benchmarks）：
   • 对SSD算法的速度进行了基准测试，与Mamba的扫描实现和FlashAttention-2进行了比较。
   • SSD算法由于利用了GPU上的矩阵乘法单元，比Mamba的融合关联扫描快2-8倍，并且在序列长度为2K及以上时比FlashAttention-2更快。
7. 架构消融研究（Architecture Ablations）：
   • 对Mamba-2块的设计进行了消融研究，包括参数投影、额外的归一化层等。
   • 验证了并行投影节省参数并提高了性能，额外的归一化层也略微提高了性能并有助于训练稳定性。
8. 多头结构（Head Structure）：
   • 对Mamba-2架构中多头结构的选择进行了消融研究，包括多值（MVA）、多查询（MQA）、多键（MKA）和多头（MHA）模式。
   • 发现多值（MVA）模式表现最佳。
9. 注意力核近似（Attention Kernel Approximations）：
   • 对于SSD结合线性注意力文献中的各种核近似方法进行了消融研究。
   • 测试了不同的核激活函数，包括cosFormer、Random Feature Attention和Performer等。

这些实验全面评估了Mamba-2架构和SSD算法在不同方面的表现，包括记忆能力、语言建模性能、扩展性、速度和架构设计选择。通过这些实验，论文展示了Mamba-2在多个任务上与现有模型相比具有竞争力或优越性。

Q: 有什么可以进一步探索的点？

A: 论文在第10节中提出了一些可能的研究方向，这些方向可以用来进一步探索和改进结构化状态空间模型（SSMs）和注意力机制。以下是一些主要的探索点：

1. 提高SSMs的表达性：研究如何扩展SSMs以包含更复杂的动态，例如通过引入非线性或更复杂的状态转换。
2. 改进算法效率：开发新的算法来进一步提高SSMs的计算效率，特别是在处理长序列时。
3. 理论分析：深入研究SSMs和注意力机制的理论基础，包括它们的性质、稳定性和收敛性。
4. 系统优化：探索如何将SSMs与现代硬件架构更有效地结合，以实现更好的性能和扩展性。
5. 注意力机制的改进：研究如何改进注意力机制，以解决如“注意力下沉”（attention sink）等问题，并提高其效率。
6. 结构化矩阵的进一步研究：研究其他类型的结构化矩阵，以及如何将它们应用于深度学习模型中，以提高效率和表达性。
7. 模型压缩和加速：探索模型压缩技术，以减少SSMs和注意力模型的参数数量和计算复杂度。
8. 跨领域应用：将SSMs和改进的注意力机制应用于其他领域，如计算机视觉、自然语言处理之外的任务。
9. 模型解释性：提高SSMs和注意力模型的解释性，帮助理解模型的决策过程。
10. 混合模型：研究如何将SSMs与Transformers或其他类型的模型结合，以创建混合模型，这些模型结合了各自的优点。
11. 优化硬件实现：针对特定硬件平台（如GPU、TPU）优化SSMs和注意力机制的实现，以充分利用硬件资源。
12. 模型泛化能力：研究SSMs和注意力模型在不同任务和数据集上的泛化能力，以及如何提高它们对新情况的适应性。

这些探索点为未来的研究提供了丰富的方向，有助于推动深度学习模型在效率、表达性和应用范围上的进一步发展。

Q: 总结一下论文的主要内容

A: 这篇论文提出了一个名为结构化状态空间对偶（Structured State Space Duality, SSD）的理论框架，旨在建立结构化状态空间模型（SSMs）和注意力机制之间的联系。以下是论文的主要内容概述：

1. 背景介绍：论文首先介绍了Transformers和SSMs在深度学习，尤其是语言建模方面的成功应用。尽管SSMs在小到中等规模上展现出与Transformers相媲美或更优的性能，但它们的开发似乎与Transformers的优化工作相脱节。
2. SSD框架：作者提出了SSD框架，通过结构化矩阵的概念，展示了SSMs和注意力机制之间的理论联系。这一框架允许将Transformers的算法和系统优化转移到SSMs中。
3. SSMs与结构化矩阵：论文证明了SSMs与一类被称为半可分矩阵的结构化矩阵等价。这种联系揭示了SSMs的新属性和算法。
4. 线性注意力理论的改进：作者改进了线性注意力的理论，并将其推广到一个新的结构化掩蔽注意力（Structured Masked Attention, SMA）家族。
5. SSMs与SMA的联系：展示了SSMs和SMA之间存在大量交集，并且这些模型是彼此的对偶，具有SSM类的线性和注意力类的二次形式。
6. 高效算法：论文提出了新的SSD算法，该算法利用半可分矩阵的块分解，实现了在所有主要效率轴上的最佳权衡。
7. 架构设计：利用SSD框架，作者设计了新的架构Mamba-2，该架构在语言建模方面与Transformers竞争，同时在序列长度上具有更好的扩展性。
8. 系统优化：SSD框架允许利用为Transformers开发的系统优化工作，如张量并行性和序列并行性，来提高SSMs的效率。
9. 实验验证：论文通过实验验证了Mamba-2在语言建模、训练效率和多查询关联记忆任务上的性能。
10. 相关工作和未来方向：论文最后讨论了与SSD框架相关的工作，并提出了未来研究的可能方向。

总的来说，这篇论文通过提出SSD框架，为理解和改进序列模型提供了新的视角，并展示了如何将SSMs和注意力机制的优势结合起来，以构建性能更优、扩展性更好的深度学习模型。

循环神经网络（RNN）长期以来一直是建模时序数据的经典架构，但它们在训练长序列时却臭名昭著地难以训练。这是因为随着时间的推移，向后传播的误差信号往往会消失或爆炸。近年来，状态空间模型（SSM）作为 RNN 的一个子类，成功地克服了这些困难，这挑战了我们对 RNN 的理论理解。

本文深入探讨了 RNN 的优化挑战，发现随着网络记忆的增加，其参数的变化会导致输出变化越来越大，使得基于梯度的学习高度敏感，即使没有爆炸梯度。我们的分析进一步揭示了逐元素递归设计模式与精心参数化的重要性，它们可以减轻这种影响。这种特性存在于 SSM 中，也存在于其他架构中，例如 LSTM。总体而言，我们的见解为 RNN 基于梯度学习的一些困难提供了一个新的解释，以及为什么某些架构比其他架构表现更好的原因。

消失和爆炸梯度：一个老问题

循环神经网络的训练难点在于，随着时间的推移，误差信号会不断累积，导致梯度要么消失，要么爆炸。这就像一个传声筒，声音在传递过程中不断衰减或放大。

公式说明：

假设一个 RNN 的隐藏状态为 ht，更新函数为 fθ，参数为 θ，输入序列为 (xt)t。网络的平均性能由损失函数 L 来衡量。

ht+1 = fθ(ht, xt+1)  
L = E[ (1/T) * Σ(t=1 to T) Lt(ht) ]

其中，Lt 表示在时间步 t 的损失。

梯度消失和爆炸的根源：

梯度消失和爆炸的根本原因在于，在反向传播过程中，误差信号会乘以一个称为雅可比矩阵（Jacobian matrix）的矩阵。如果该矩阵的特征值小于 1，误差信号就会不断衰减，导致梯度消失；反之，如果特征值大于 1，误差信号就会不断放大，导致梯度爆炸。

解决方法：

为了解决这个问题，人们提出了各种方法，例如：

• LSTM 和 GRU： 这些架构引入了记忆单元，可以存储和检索信息，并通过特殊的门控机制控制误差信号的流动，从而避免梯度消失或爆炸。
• 梯度裁剪： 将梯度的大小限制在一个范围内，防止其过大或过小。
• 活动归一化： 通过对神经元的活动进行归一化，使梯度保持在一个合理的范围内。
• 权重初始化： 使用合适的权重初始化方法，可以避免梯度消失或爆炸。
• 架构约束： 通过对网络结构进行约束，例如使用分层处理、正交权重矩阵或振荡机制，可以改善梯度传播。

记忆的诅咒：一个新挑战

传统的观点认为，解决消失和爆炸梯度问题就足以让 RNN 学习长期依赖关系。然而，我们发现，即使解决了这些问题，RNN 仍然面临着一个新的挑战：记忆的诅咒。

直观解释：

RNN 的更新函数 fθ 会反复应用，这意味着修改参数 θ 不仅会影响一次更新，还会影响所有未来的更新。随着网络记忆的增加，隐藏状态会保留更多更新的影响，从而变得对参数变化更加敏感。

数学分析：

我们可以通过分析隐藏状态 ht 对参数 θ 的敏感度来理解记忆的诅咒：

dht/dθ = Σ(t'≤t) dht/dht' * ∂fθ/∂θ(ht'-1, xt')

当网络记忆更长时，雅可比矩阵 dht’/dht 的非零项数量会增加，导致敏感度增加。

线性对角循环神经网络的分析：

为了更深入地理解记忆的诅咒，我们分析了线性对角循环神经网络，其更新函数为：

ht+1 = λ ⊙ ht + xt+1

其中，λ 是一个向量，⊙ 表示逐元素乘积。

结论：

我们发现，随着 |λ| 接近 1，隐藏状态和反向传播误差的方差都会爆炸。这意味着，当网络试图编码更长的依赖关系时，其对参数变化的敏感度会急剧增加，从而阻碍学习过程。

减轻记忆的诅咒：解决方案

为了减轻记忆的诅咒，我们可以采用以下方法：

• 输入归一化： 通过对输入进行归一化，可以控制隐藏状态的方差，避免其爆炸。
• 特征值重新参数化： 通过对 λ 进行重新参数化，可以控制梯度的大小，避免其爆炸。

状态空间模型和门控 RNN：

状态空间模型和门控 RNN 隐含地包含了这些机制。状态空间模型通过对连续时间微分方程进行离散化来实现，而门控 RNN 通过门控机制来控制信息的流动。

线性教师-学生分析：一个简单模型

为了验证我们的理论，我们设计了一个线性教师-学生任务，其中教师和学生都遵循线性循环动态。

结果：

我们发现，当教师编码更长的依赖关系时，线性 RNN 难以学习，而具有对角连接的 LRU 则表现良好。

对自适应学习率的重要性：

我们还发现，对角连接可以简化损失函数的结构，使自适应学习率优化器能够更好地处理记忆的诅咒。

深度循环神经网络的初始化信号传播：实际验证

为了验证我们的理论在实际应用中的有效性，我们分析了深度循环神经网络在初始化时的信号传播。

结果：

我们发现，在初始化时，复杂对角 RNN 会受到记忆的诅咒的影响，而 LRU 则可以有效地减轻这种影响。

结论

消失和爆炸梯度并非循环神经网络训练的全部故事。记忆的诅咒是一个新的挑战，它源于 RNN 的迭代性质，并在网络接近动态稳定性边缘时出现。重新参数化和自适应学习率可以有效地减轻这种行为，而对角化递归可以简化这两个过程。

我们的分析还揭示了学习复数特征值的复杂性，这可能解释了为什么最近大多数状态空间模型架构中没有发现复数有用。

未来研究方向：

理解如何提高小型线性模块的表达能力，同时保持其优良的优化特性，是一个很有前景的研究方向。

---

Q: 这篇论文试图解决什么问题？

A: 这篇论文探讨了循环神经网络（Recurrent Neural Networks, RNNs）在学习长期记忆时面临的挑战，尤其是梯度消失和梯度爆炸问题。尽管这些问题已经得到了广泛的研究，但论文指出，即使解决了这些问题，RNNs 在学习过程中仍然会遇到新的困难，即所谓的“记忆的诅咒”（curse of memory）。具体来说，论文的主要贡献和解决的问题包括：

1. 记忆的诅咒：论文提出了一个新的问题，即随着网络记忆能力的增加，网络活动对其参数变化的敏感度也在增加，这使得基于梯度的学习变得非常敏感，即使在没有梯度爆炸的情况下也是如此。
2. 状态空间模型（SSMs）的有效性：论文探讨了状态空间模型（一种特殊的 RNN），这类模型在捕捉长期依赖性方面表现出色。作者试图分析这种有效性是否可以完全归因于它们避免了梯度消失和爆炸。
3. 元素级递归设计模式的重要性：论文分析了元素级递归设计模式与精心设计的参数化在减轻记忆诅咒效应中的重要性。
4. RNNs 的优化挑战：论文深入探讨了 RNNs 的优化问题，并提出了一些新的见解，解释了为什么某些架构（如 LSTMs 和 SSMs）在基于梯度的学习中表现更好。
5. 理论分析与实验验证：论文不仅在理论上分析了上述问题，还通过简单的教师-学生任务和深度循环网络的实验来验证理论分析的正确性。

总的来说，论文试图提供一个更全面的视角来理解 RNNs 在学习长期依赖性时面临的挑战，并探讨了除了梯度消失和爆炸之外的其他复杂性。

Q: 有哪些相关研究？

A: 论文中提到了多个与循环神经网络（RNNs）及其训练挑战相关的研究领域和具体工作。以下是一些主要的相关研究和它们的贡献：

1. 梯度消失和梯度爆炸问题：
   • Hochreiter (1991) [5], Bengio et al. (1994) [6], Hochreiter et al. (2001) [7] 等研究者早期工作强调了在时间反向传播过程中梯度消失或爆炸的困难。
   • Pascanu et al. (2013) [8] 对此问题进行了进一步的分析。
2. 解决梯度问题的尝试：
   • LSTM (Long Short-Term Memory) 单元 [3] 和 GRU (Gated Recurrent Unit) [23] 通过使用记忆神经元来促进直接信息存储和检索，从而解决了这个问题。
3. 其他方法：
   • 梯度裁剪 [24, 8], 活动归一化 [25–27], 谨慎的权重初始化 [28, 29] 或者架构约束（如分层处理 [30, 31], 正交权重矩阵 [32–34] 和振荡 [35–37]）。
4. 注意力机制和 Transformer：
   • Bahdanau et al. (2015) [9] 提出了用于神经机器翻译的注意力机制。
   • Vaswani et al. (2017) [10] 提出了 Transformer 模型，它通过直接的 token-to-token 通信简化了长序列信号传播。
5. Transformer 的效率改进：
   • 一系列研究 [11–15] 旨在使 Transformer 更加高效，包括更快的自回归 Transformer [12], FlashAttention [13], Switch Transformers [14] 等。
6. 状态空间模型 (SSMs)：
   • 一系列研究 [16–22] 探讨了状态空间模型，这是一种新型的线性循环网络，以其捕捉长期依赖性的能力而著称。
7. 优化和学习率理论：
   • 有关学习率和优化的理论研究，例如 Noci et al. (2024) [41] 和 Yang et al. (2020) [42]。
8. 深度学习中的信号传播：
   • 有关深度神经网络中信号传播的研究，例如 Chen et al. (2018) [50] 和 Cohen et al. (2020) [54]。
9. 在线学习和组合泛化：
   • 有关在线学习 [58] 和组合泛化 [59] 的研究。

这些研究为理解 RNNs 的训练挑战、改进它们的性能以及开发新的架构提供了理论基础和实证支持。论文通过引用这些相关工作，建立了其研究的理论背景和实践意义。

Q: 论文如何解决这个问题？

A: 论文通过以下几个步骤来解决循环神经网络（RNNs）在学习长期记忆时遇到的问题：

1. 理论分析：首先，论文对RNNs的优化挑战进行了深入的理论分析，特别是针对“记忆的诅咒”（curse of memory）这一新问题。作者探讨了随着网络记忆的增加，网络活动对参数变化的敏感度如何增加，以及这如何影响梯度的学习。
2. 状态空间模型（SSMs）：论文研究了SSMs这一特殊类型的RNN，这些模型通过避免梯度消失和爆炸问题，以及采用元素级递归设计模式和精心设计的参数化，有效地捕捉长期依赖性。
3. 元素级递归设计模式：论文强调了元素级递归设计模式的重要性，这种设计模式结合了精心的参数化，可以减轻记忆诅咒效应。
4. 教师-学生任务分析：通过一个简单的教师-学生任务，论文展示了线性循环网络的学习复杂性，并分析了所提出的解决方案在这种简单设置中的有效性。
5. 实验验证：论文通过实验验证了理论分析的正确性。作者在深度循环网络中研究了初始化时的信号传播，并比较了不同循环层（如复杂对角线RNN、LRU和LSTM）的性能。
6. 参数化和归一化策略：论文提出了特定的参数化和归一化策略，如输入归一化和特征值重参数化，以控制隐藏状态和梯度的幅度，从而减轻记忆诅咒的影响。
7. 适应性学习率优化器：论文指出，适应性学习率优化器（如Adam）可以有效地解决由于记忆诅咒导致的额外敏感性问题，特别是当这些敏感性集中在少数几个参数上时。

通过这些方法，论文不仅提供了对RNNs训练难题的新解释，还揭示了为什么某些架构（如LSTMs和SSMs）在基于梯度的学习中表现更好，并为设计更易于优化的RNN架构提供了指导。

Q: 论文做了哪些实验？

A: 论文中进行了一系列实验来验证理论分析，并探讨循环神经网络（RNNs）在学习和优化过程中的行为。以下是论文中提到的主要实验：

1. 线性教师-学生任务（Linear Teacher-Student Task）：
   • 这是一个简化的设置，用于直观展示记忆诅咒和梯度消失问题是如何相互作用的。
   • 作者考虑了一个一维动力学的学生模型，其目标是复现具有复数参数λ*的教师模型的隐藏状态。
   • 通过改变λ*的值，研究了在不同记忆长度下，学生模型学习教师行为的难度。
2. 对角连接简化优化（Diagonal Connectivity Simplifies Optimization）：
   • 实验中，作者比较了线性RNN和具有对角连接的LRU（Linear Recurrent Unit）在学习长期依赖性时的性能差异。
   • 通过调整教师模型的记忆长度（通过控制隐藏维度n和变换参数ν），观察了学生模型在复现教师行为时的困难程度。
3. 自适应学习率的重要性（Importance of Adaptive Learning Rates）：
   • 研究了直接参数化循环连接矩阵的复杂特征值如何影响损失景观，并使具有自适应学习率的优化器能够补偿这些行为。
   • 通过分析损失函数的Hessian矩阵，展示了在优化过程中，自适应学习率如何帮助处理额外的敏感性。
4. 深度循环网络中的信号传播（Signal Propagation in Deep Recurrent Networks at Initialization）：
   • 在深度循环网络中，研究了在初始化时信号如何在网络中传播，以及记忆参数ν如何影响隐藏状态和梯度的方差。
   • 比较了不同类型的循环层（如复杂对角线RNN、LRU和LSTM）在信号传播和梯度爆炸方面的表现。
5. 损失景观的结构分析（Structure of the Loss Landscape）：
   • 通过可视化和理论分析，研究了在优化过程中损失景观的结构，特别是在对角化和非对角化情况下的差异。
6. 集中特征值分布的影响（Impact of Concentrating Eigenvalue Distributions）：
   • 探讨了当特征值分布集中时，对学习动态的影响，特别是在对角RNN和LRU架构中。

这些实验涉及了理论分析、数值模拟和可视化，旨在深入理解RNNs在学习和优化过程中的复杂性，并验证作者提出的理论观点。通过这些实验，论文展示了记忆诅咒问题的存在，并提出了相应的解决方案。

Q: 有什么可以进一步探索的点？

A: 论文中提出了多个可以进一步探索的点，以下是一些主要的研究方向：

1. 复杂参数化：论文讨论了对于复数参数的复杂参数化问题，特别是角度参数的困难。研究如何改进这种参数化，以便更好地优化和学习可能是一个有价值的方向。
2. 自适应学习率优化器：论文指出自适应学习率优化器（如Adam）在处理敏感性方面表现出色。进一步研究这些优化器与RNN架构的相互作用，以及如何改进它们以更好地处理长期依赖性，可能是一个有前景的领域。
3. 状态空间模型（SSMs）：SSMs在捕捉长期依赖性方面表现出色，但它们的表达能力有限。研究如何增加SSMs的表达能力，同时保持它们的快速训练速度，可能是一个重要的研究方向。
4. 深度循环网络的初始化：论文在初始化时研究了深度循环网络中的信号传播。探索不同的初始化策略，以及它们如何影响网络的训练动态和最终性能，可能是一个有价值的研究领域。
5. 损失景观的结构：论文分析了损失景观的结构，特别是在对角化和非对角化情况下的差异。进一步研究这些结构如何影响优化过程，以及如何设计损失函数以促进更好的学习，可能是一个有趣的方向。
6. 集中特征值分布的影响：论文探讨了集中特征值分布对学习动态的影响。研究如何管理这些分布，以及它们如何影响不同架构（如对角RNN和LRU）的性能，可能是一个值得探索的领域。
7. 模块化和独立机制：论文提出了神经元和自适应学习率优化器之间的共生关系。研究如何设计具有模块化和独立机制的网络，以及这些设计如何促进在线学习、组合泛化和并行化，可能是一个有潜力的研究方向。
8. 理论与实践的结合：论文中的一些理论分析和实验结果表明，理论与实践之间存在差距。进一步探索这种差距的原因，并开发能够弥合这一差距的方法和理论，可能是一个重要的研究领域。
9. 长期依赖性学习：尽管论文讨论了长期依赖性学习的一些挑战，但这一领域仍然存在许多未解决的问题。研究新的架构和技术，以更有效地学习长期依赖性，可能是一个持续的研究主题。

这些探索点为未来的研究提供了丰富的思路，有助于推动循环神经网络和深度学习领域的进一步发展。

Q: 总结一下论文的主要内容

A: 这篇论文深入探讨了循环神经网络（Recurrent Neural Networks, RNNs）在学习长期依赖性时面临的挑战，并提出了新的见解和解决方案。以下是论文的主要内容总结：

1. 问题背景：RNNs 在学习长期记忆时，由于梯度消失和梯度爆炸问题，训练变得非常困难。尽管现有研究已经提出了多种解决这些问题的方法，但论文指出，即使这些问题得到解决，RNNs 仍然面临其他挑战。
2. 记忆的诅咒：论文提出了一个新的问题——“记忆的诅咒”。随着网络记忆的增加，网络活动对参数变化的敏感度也在增加，这使得基于梯度的学习变得非常敏感。
3. 状态空间模型（SSMs）：论文研究了状态空间模型，这是一类特殊的 RNN，它们能够有效地捕捉长期依赖性。作者探讨了 SSMs 的成功是否可以完全归因于它们避免了梯度消失和爆炸。
4. 元素级递归设计模式：论文强调了元素级递归设计模式的重要性，并指出结合精心设计的参数化可以减轻记忆诅咒效应。
5. 理论分析：论文通过理论分析，展示了 RNNs 在信号传播方面的行为，特别是在对角化和非对角化情况下的差异。
6. 实验验证：通过一系列实验，包括线性教师-学生任务、深度循环网络的信号传播分析等，论文验证了理论分析的正确性，并展示了所提出解决方案的有效性。
7. 自适应学习率优化器：论文指出，自适应学习率优化器（如 Adam）可以有效地解决由于记忆诅咒导致的额外敏感性问题。
8. 损失景观结构：论文分析了损失景观的结构，特别是在对角化和非对角化情况下的差异，以及如何通过参数化和归一化策略来控制隐藏状态和梯度的幅度。
9. 未来研究方向：论文提出了多个未来研究方向，包括改进参数化、优化器与架构的相互作用、损失函数设计、初始化策略等。
10. 结论：论文得出结论，虽然梯度消失和爆炸问题是 RNNs 训练中的挑战，但解决这些问题并不足以确保良好的学习性能。通过揭示新的挑战和解决方案，论文为理解和改进 RNNs 的训练提供了新的视角。

总的来说，这篇论文通过理论分析和实验验证，为理解和解决 RNNs 在学习长期依赖性时面临的挑战提供了深入的见解，并提出了一系列有价值的研究方向。

---