Tensor & Batch（张量 ，批量）

### 多头张量和批量张量的概念 在深度学习框架中，多头张量（Multi-head Tensor）通常用于注意力机制中的多头自注意力模型。这种设计允许模型在同一层次上关注不同位置的不同表示空间[^1]。 #### 批量张量定义 批量张量（Batch Tensor）是指包含多个样本的数据结构，在神经网络处理过程中一次性传递给模型的一组数据。这有助于提高计算效率并通过并行化加速训练过程。对于形状为 `[batch_size, ...]` 的张量来说，第一个维度代表批次数目[^4]。 #### 多头张量的具体应用 当涉及到像Transformer这样的架构时，输入会被分割成若干个头部（heads），每个头部独立工作于不同的特征子集之上。具体而言： - 输入序列被映射到查询（Query）、键（Key）以及值（Value）三个矩阵； - 这些矩阵进一步拆分为预定数量的小型向量集合——即所谓的“头”，形成三维张量形式 `(num_heads, seq_len, head_dim)`； - 各个头分别执行点积相似度运算、softmax归一化操作，并最终聚合得到加权求和的结果作为输出。 ```python import torch from torch import nn class MultiHeadAttention(nn.Module): def __init__(self, embed_size, num_heads): super(MultiHeadAttention, self).__init__() self.embed_size = embed_size self.num_heads = num_heads assert embed_size % num_heads == 0, "Embedding size must be divisible by number of heads" self.head_dim = embed_size // num_heads # Linear transformations for Q,K,V self.q_linear = nn.Linear(embed_size, embed_size) self.k_linear = nn.Linear(embed_size, embed_size) self.v_linear = nn.Linear(embed_size, embed_size) self.out = nn.Linear(embed_size, embed_size) def forward(self, query, key, value, mask=None): N = query.shape[0] q = self.q_linear(query).view(N, -1, self.num_heads, self.head_dim).transpose(1,2) k = self.k_linear(key).view(N, -1, self.num_heads, self.head_dim).transpose(1,2) v = self.v_linear(value).view(N, -1, self.num_heads, self.head_dim).transpose(1,2) energy = torch.einsum("nqhd,nkhd->nhqk", [q,k]) / (self.head_dim ** (1/2)) if mask is not None: energy = energy.masked_fill(mask==0, float("-1e20")) attention = torch.softmax(energy, dim=-1) out = torch.einsum("nhql,nlhd->nqhd",[attention,v]).reshape(N,-1,self.embed_size) return self.out(out) ``` 上述代码展示了如何构建一个多头注意模块，其中包含了对批量大小的支持。这里的关键在于理解 `einsum` 函数的作用：它能够灵活地指定两个或多于两个张量之间的乘法规则，非常适合用来描述复杂的线性代数变换关系。