anschen用户9113
大模型里的Attention机制
输入“/”快速插入内容
大模型里的Attention机制
大模型里的Attention机制
用户9113用户91132025年3月2日修改
引言
近年来，随着深度学习技术的快速发展，大模型（Large Language Models, LLMs）在自然语言处理（NLP）领域取得了显著的进展。这些模型通过大规模数据训练，能够生成高质量的文本、完成复杂的语言理解任务，并在多个领域展现出强大的能力。然而，随着模型规模的不断扩大，如何高效地处理长序列数据、降低计算复杂度以及减少内存占用成为了研究的热点问题。
在大模型中，Attention 机制（注意力机制）是核心组件之一。自 Transformer 模型提出以来，Attention 机制凭借其强大的序列建模能力，逐渐取代了传统的循环神经网络（RNN）和卷积神经网络（CNN）。然而，标准的 Attention 机制（如 Multi-Head Attention, MHA）在处理长序列时面临着计算复杂度和内存占用的双重挑战，这使得其在处理长文本时效率低下。
为了解决这些问题，研究者们提出了多种改进的 Attention 机制，例如 FlashAttention、Multi-Query Attention (MQA)、Group Query Attention (GQA) 以及 Sparse Attention 等。这些方法通过优化计算流程、减少内存访问次数或引入稀疏性，显著提升了 Attention 机制的效率。此外，DeepSeek-V2 提出的 Multi-Head Latent Attention (MLA) 进一步通过低维映射和算子融合技术，降低了计算复杂度和内存占用，为大模型的规模化应用提供了新的思路。
本文将深入探讨这些 Attention 技术的原理、实现细节及其在大模型中的应用。我们将从标准的 MHA 出发，逐步介绍其变体（如 MQA 和 GQA），并详细分析 FlashAttention 如何通过分块计算和重计算技术优化性能。此外，我们还将探讨 Sparse Attention 和 DeepSeek-MLA 的创新设计，以及它们如何在大规模语言模型中实现高效的长序列处理。
通过对这些技术的全面分析，本文旨在为读者提供一个系统的视角，理解大模型中 Attention 机制的演进及其在实际应用中的优化策略。我们相信，随着这些技术的不断发展，大模型将在更多领域展现出其强大的潜力，并为人工智能的未来开辟新的可能性。
Attention技术
MHA(Multi Head Attention)
💡
Attention is all you need.
Single Head
在基础的MHA，给定一个Sequence的表示
，
是Sequence的长度，
是表征的维度。
构造三个Tensor,
, 那么对应的，
Attention的结果
则变成了：
，这边
代表按行softmax操作。
Multi Head
Single Head往Multi Head的推广则只需要做如下的改变就可以了
是Head的个数。下面这段提供了 Multi Head的Demo实现的两种方式：
1.
使用Single Head的方式，然后拼接起来
2.
直接完整的张量计算方式
代码块
import torch
batch_size = 8
seq_len = 16
d = 128
h = 4
sub_d = d // h
x = torch.randn(size=(batch_size, seq_len, d))
w_qs = [torch.randn(size=(d, sub_d)) for _ in range(h)]
w_ks = [torch.randn(size=(d, sub_d)) for _ in range(h)]
w_vs = [torch.randn(size=(d, sub_d)) for _ in range(h)]
print("========From Single Head START=========")
os = []
for head in range(h):
q = torch.matmul(x, w_qs[head])
k = torch.matmul(x, w_ks[head])
v = torch.matmul(x, w_vs[head])
att = torch.einsum('bnd,bmd->bnm', q, k) / (sub_d ** 0.5)
att_sep.append(att)
att = torch.softmax(att, dim=-1)
att_sep_soft.append(att)
o = torch.einsum('bnm,bmd->bnd', att, v)
os.append(o)
output = torch.cat(os, dim=-1)
print("output shape: ", output.shape)
print("========From Single Head END=========")
print("========VEC START=========")
w_qs_vec = torch.stack(w_qs, dim=-1) # (d, sub_d, h)
w_ks_vec = torch.stack(w_ks, dim=-1)
w_vs_vec = torch.stack(w_vs, dim=-1)
q_vec = torch.einsum('bnd,dsh->bnsh', x, w_qs_vec) # (b, n, d, h)
k_vec = torch.einsum('bnd,dsh->bnsh', x, w_ks_vec)
v_vec = torch.einsum('bnd,dsh->bnsh', x, w_vs_vec)
att = torch.einsum('bnsh,bmsh->bnmh', q_vec, k_vec) / (sub_d ** 0.5)
att = torch.softmax(att, dim=-2)
o = torch.einsum('bnmh,bmsh->bnhs', att, v_vec)
output2 = torch.reshape(o, shape=(batch_size, seq_len, d))
print("output2 shape: ", output2.shape)
print("========VEC END=========")
err = torch.sum(torch.abs(output - output2))
print("err: ", err)