### 图解 Transformer 模型中的 Gather 操作 Gather 操作在 Transformer 中用于索引特定位置的数据，通常应用于嵌入层或注意力机制中。为了更好地理解这一过程，下面将通过文字描述和伪代码来展示其工作原理。 #### 1. 嵌入层中的 Gather 操作 当输入一批 token ID 到模型时，这些 IDs 被用来查找预定义好的词向量矩阵。假设有一个大小为 \(V \times d\) 的可学习参数矩阵 \(W_{emb}\)，其中 \(V\) 表示词汇表大小而 \(d\) 是每个单词对应的维度，则对于给定的一批 token 序列： \[X = [x_1, x_2,... ,x_n]\] 这里每一个 \(x_i\) 都是从范围 \([0,V)\) 内取值的一个整数形式的 token id 。此时 gather 就是用来获取这批 tokens 对应于权重矩阵里的那一行作为它们各自的 embedding 向量[^2]。 ```python import torch def get_embeddings(token_ids, embeddings_matrix): """ 获取token ids对应的embedding向量 参数: token_ids (Tensor): 形状为(batch_size, seq_len), 存储着batch内各条样本每一步骤上的token编号. embeddings_matrix (Tensor): 形状为(vocab_size, embed_dim), 即整个字典空间下所有可能token的embeddings. 返回: Tensor: shape=(batch_size, seq_len, embed_dim). """ return torch.gather(embeddings_matrix, 0, token_ids.unsqueeze(-1).expand(-1,-1,embeddings_matrix.size(1))) ``` #### 2. 注意力机制下的 Gather 操作 另一个典型的应用场景是在多头自注意模块里边，在完成 QKV 计算之后得到的是形状类似于 `(batch_size, num_heads, sequence_length, head_dims)` 的张量；而在某些情况下我们需要根据某个条件重新排列这个四维数组中的元素顺序——比如按照 softmax 权重分布去加权求和之前先挑选出最有可能关联的部分来做进一步处理。这时就可以利用 `torch.index_select` 或者更高效的 `gather` 函数实现这样的功能[^1]。 ```python def apply_attention_scores(scores, values): """ 使用attention scores对value进行加权平均 参数: scores (Tensor): attention score matrix with shape (batch_size, n_heads, query_seq_len, key_seq_len) values (Tensor): value vectors with shape (batch_size, n_heads, key_seq_len, dim_per_head) 返回: Tensor: 加权后的context vector tensor of shape (batch_size, n_heads, query_seq_len, dim_per_head) """ batch_size, n_heads, q_seq_len, k_seq_len = scores.shape _, _, v_seq_len, dim_per_head = values.shape # Apply softmax to normalize the weights over last dimension attn_weights = torch.softmax(scores / np.sqrt(dim_per_head), dim=-1) # Perform weighted sum using gather operation on 'values' based on indices from 'attn_weights' context_vectors = torch.matmul(attn_weights, values) return context_vectors ``` 上述两部分分别阐述了如何在不同阶段运用 gather 方法提取所需信息，并给出了相应的 Python 实现片段以便读者能够更加直观地感受其运作流程[^3]。