### 实现 Top-K 选择 为了在 Informer 模型中实现 top-k 选择，主要是在自注意力机制部分进行优化。Informer 使用了一种称为 ProbSparse 自注意力机制来减少计算量并提高效率[^1]。 ProbSparse 自注意力机制的核心思想是从输入序列中随机采样一部分 token 进行计算，而不是对整个序列进行全面的自我注意计算。具体来说，在计算 QK^T (查询矩阵和键矩阵相乘) 的过程中引入了稀疏化操作，从而只关注最有可能影响预测结果的部分元素。 对于 top-k 选择的具体实现方法如下： ```python import torch from torch import nn class ProbAttention(nn.Module): def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False): super(ProbAttention, self).__init__() self.factor = factor self.scale = scale self.mask_flag = mask_flag self.output_attention = output_attention self.dropout = nn.Dropout(attention_dropout) def _prob_QK(self, Q, K, sample_k, n_top): # n_top: sparsity level B, H, L_K, E = K.shape _, _, L_Q, _ = Q.shape # Calculate the sampled column index of keys U_col = torch.topk(torch.mean(torch.abs(K), dim=-1), sample_k, dim=-1)[1].unsqueeze(1).expand(-1, H, -1, -1) # Gather columns from key matrix based on indices obtained above K_sampled = torch.gather(K, dim=2, index=U_col.repeat([1,1,E,1]).permute((0,1,3,2))) # Compute similarity scores between queries and sampled keys Q_K_sampled = torch.matmul(Q.float(), K_sampled.transpose(-2,-1).float()) if self.mask_flag != False: pass # Apply masking here M = torch.max(Q_K_sampled,dim=-1,keepdim=True)[0] Q_K_sampled -= M exp_values = torch.exp(Q_K_sampled / self.scale) exp_values_sum = torch.sum(exp_values, dim=-1, keepdim=True) A = exp_values.div_(exp_values_sum) top_k_indices = torch.topk(A.view(B*H*L_Q,n_top), k=n_top, dim=-1)[1] return top_k_indices def forward(self, queries, keys, values, attn_mask): ... ``` 此代码片段展示了如何在一个简化版本的概率性自注意力层 `ProbAttention` 中执行 top-k 选取逻辑。这里的关键在于 `_prob_QK()` 函数内部的操作——通过对键向量取绝对值后的均值得到重要程度评分，并据此选出最重要的样本作为后续处理的对象[^2]。