一起读《大模型驱动的具身智能：架构、设计与实现》- 注意力机制Attension Mechanism

视频链接https://www.bilibili.com/video/BV1g7dZBpEJG/?vd_source5ba34935b7845cd15c65ef62c64ba82f仓库链接https://github.com/LitchiCheng/LLM-learning最近在阅读《大模型驱动的具身智能架构、设计与实现》其中 LLM 不得不提注意力机制已经有太多好的视频或文章针对于 [1706.03762] Attention Is All You Needhttps://arxiv.org/abs/1706.03762 进行分析介绍下面分享记录下自己学习过程和相关点的理解节奏上会比较跳跃不适合系统性学习最后用一段代码来测试注意力的计算过程。所有注意力都由 3 个向量计算最后让模型 “看重点”自动决定该关注权重高哪些信息、忽略权重低哪些信息Query查询 Q现在要找什么Key键 K这里有什么信息Value值 V信息的实际内容[B, L, D] 形状的张量是序列模型标准输入格式[batch_size, sequence_length, embedding_dimension]seq_len, embed_dim 4, 8 world_states torch.randn(1, seq_len, embed_dim)[ [ [0.2, -0.5, 0.1, 0.8, -0.3, 0.6, 0.0, 0.4], # 位置0: 球位置的8维特征 [-0.1, 0.3, 0.7, -0.2, 0.5, -0.4, 0.2, 0.1], # 位置1: 球速度的8维特征 [0.4, -0.1, 0.2, 0.3, -0.6, 0.8, -0.2, 0.5], # 位置2: 墙壁的8维特征 [-0.3, 0.6, -0.4, 0.1, 0.7, -0.2, 0.3, 0.0] # 位置3: 目标的8维特征 ] ]dim 设置为 8 是注意力机制的计算的一种编码形式没有什么物理意义相当于给每个概念球位置/速度/墙壁/目标分配一个 8 维的空间让模型在这个高维空间里用线性代数Q/K/V 投影、点积、softmax来计算谁该关注谁。训练后相似的物理状态会在 8 维空间中靠得更近默认的 linear 神经网络随机分配了不同的权重dim*dim可以用 Q_k.weight 查看每个 token 都会作为 Query 去查询所有 token 的 Key决定把注意力放在哪里Q_位置 · K_速度^T → 高权重 # 速度和位置强相关运动学 Q_位置 · K_墙壁^T → 中权重 # 墙壁限制位置范围 Q_位置 · K_目标^T → 高权重 # 位置相对目标的距离很重要位置的输出会混合速度、目标、墙壁的信息形成位置和前后速度距离目标都有关系的上下关联Q_速度 · K_位置^T → 高权重 # 速度改变位置 Q_速度 · K_墙壁^T → 高权重 # 即将撞墙时速度会反转反弹 Q_速度 · K_目标^T → 中权重 # 速度方向是否朝向目标模型理解到速度和墙壁有因果关联这是反事实推理的基础Q_墙壁 · K_位置^T → 高权重 # 只关注靠近墙壁的物体位置 Q_墙壁 · K_速度^T → 高权重 # 关注朝向墙壁运动的速度 Q_墙壁 · K_目标^T → 低权重 # 与目标关系不大墙壁知道什么时候该拦截什么时候无关完整的测试代码import os import torch import torch.nn as nn import torch.nn.functional as F import numpy as np import matplotlib matplotlib.use(Agg) import matplotlib.pyplot as plt os.makedirs(/mnt/agents/output, exist_okTrue) plt.rcParams[font.sans-serif] [SimHei, DejaVu Sans] plt.rcParams[axes.unicode_minus] False class SelfAttention(nn.Module): def __init__(self, embed_dim): super().__init__() self.scale embed_dim ** -0.5 # 缩放因子 1/√d_k self.W_q nn.Linear(embed_dim, embed_dim) # Query 投影 self.W_k nn.Linear(embed_dim, embed_dim) # Key 投影 self.W_v nn.Linear(embed_dim, embed_dim) # Value 投影 def forward(self, x, return_attentionFalse): Q self.W_q(x) # [B, seq, dim] print(self.W_q.weight) K self.W_k(x) # [B, seq, dim] V self.W_v(x) # [B, seq, dim] # QK^T / √d_k scores torch.matmul(Q, K.transpose(-2, -1)) * self.scale # scores: [B, seq, seq] - 每个位置对其他位置的关注成都 # Softmax 归一化: 权重高的关注, 权重低的忽略 attn_weights F.softmax(scores, dim-1) out torch.matmul(attn_weights, V) # [B, seq, dim] # 输出是 V 的加权组合,权重由 QK 相似度决定 if return_attention: return out, attn_weights return out torch.manual_seed(42) seq_len, embed_dim 4, 8 world_states torch.randn(1, seq_len, embed_dim) attn_layer SelfAttention(embed_dim) output, attention_map attn_layer(world_states, return_attentionTrue) print(fInput shape: {world_states.shape}) print(fAttention weights:\n{attention_map[0].detach().numpy().round(3)}) fig, ax plt.subplots(figsize(6, 5)) labels [Ball Pos, Ball Vel, Wall, Target] im ax.imshow(attention_map[0].detach().numpy(), cmapviridis, aspectauto) ax.set_xticks(range(seq_len)) ax.set_yticks(range(seq_len)) ax.set_xticklabels(labels) ax.set_yticklabels(labels) ax.set_title(Attention Weight Matrix: What the model focuses on) plt.colorbar(im, axax, labelAttention Weight) plt.tight_layout() plt.savefig(/mnt/agents/output/attention_map.png, dpi150, bbox_inchestight) plt.close(fig) print(Saved: /mnt/agents/output/attention_map.png)