CMU 10423 Generative AI：lec17（优化Attention机制的方法）

1 概述

这个文档主要介绍了如何通过优化Attention机制来加速Transformer的计算。以下是主要内容的概述：

1. 回顾：MoE与知识的稀疏性

• 在前面的课程中，讲解了如何使用Mixture of Experts (MoE) 来加速MLP层的计算。MoE的核心理念是知识的稀疏性，即在处理每个问题时只需要调动部分专家模型，而不是全部使用。
• 类似的思路可以用来优化Attention机制。

2. 加速Attention计算的技巧

• Flash Attention：一种通过分块计算和重新计算softmax来减少内存占用的注意力机制。
• Multi-Query Attention：通过减少多头注意力机制中的查询（query）的冗余计算，提升推理速度。
• Paged Attention：一种特殊的注意力机制，用于优化大规模上下文的注意力计算。

3. Multi-Head Attention（多头注意力机制）

• Transformer中的核心组件之一是多头注意力机制。给定一组输入向量，通过投影矩阵$Q$和$K$来计算注意力分数，并通过softmax函数将其归一化。
• 计算公式为：

$${\alpha }_{i,j}=\text{softmax}(v_i^{T}Q{K}^T{v}_j)+p_{i,j}$$

• 每个输入通过多个投影矩阵计算得到注意力分数，并通过线性变换和加权求和计算出最终的输出。

4. Transformer的架构

• 讨论了两种Transformer的变体：

  • Post-layernorm Transformer：在每个层的最后应用LayerNorm，并使用残差连接来增强梯度传播。
  • Pre-layernorm Transformer：在每个层的开头应用LayerNorm，从而避免梯度消失。

• 每个Transformer块的计算包括：多头注意力机制（MHA）和一个多层感知机（MLP），配合残差连接和LayerNorm。

5. 减少Attention计算的内存使用

• Attention的内存瓶颈在于存储$n\times n$大小的注意力矩阵，这会随着输入序列长度的增加而快速增长。
• Flash Attention通过在计算过程中避免存储整个softmax矩阵来减少内存使用，而是只按需计算部分结果。

6. 从"Stupid Attention"到Flash Attention

• Stupid Attention是一种简单但内存效率低的计算方式，因为它需要多次重新计算中间结果。
• Flash Attention通过引入分块计算，并巧妙地进行softmax归一化来提高计算效率，同时减少了对内存的占用。

7. 自回归训练与推理的优化

• 自回归模型在生成序列时，每个新的预测会基于之前的输出。为了提升推理速度，可以通过Multi-Query Attention缓存过去的计算结果，而不是每次都重新计算所有的查询和键值对。

8. Attention Mask与推理

• 在自回归模型中，使用注意力掩码（Attention Mask）来确保每个位置只关注它之前的序列。这种机制确保生成的文本符合因果关系。

9. Multi-Query Attention（多查询注意力）

• 传统的多头注意力机制需要为每个注意力头独立计算查询、键和值。Multi-Query Attention通过让所有注意力头共享同一个键和值，从而减少计算复杂度，提升推理速度。

总结：

本次课主要探讨了如何通过优化Attention机制来减少计算开销和内存使用，从而加速Transformer模型的推理过程。Flash Attention通过分块计算减少了内存使用，而Multi-Query Attention通过共享查询和键值对提升了推理效率。这些技术在大模型的推理中尤其重要，能够显著提升模型的实时性和计算效率。