高级索引

定义

PyTorch 高级索引：

从笛卡尔积到点对点映射在基础索引（Basic Indexing，如 tensor[0:5, :]）中，我们定义的是子空间切片。

而高级索引（Advanced Indexing）是指使用整型张量或布尔张量作为索引。

数学描述：设目标张量为 $\mathcal{A}\in {\mathbb{R}}^{d_0\times d_1\times \cdots \times d_{n-1}}$。

当我们提供 $n$ 个索引张量 ${\mathcal{I}}^{(0)},{\mathcal{I}}^{(1)},\dots ,{\mathcal{I}}^{(n-1)}$ 时，这组索引张量必须是可广播（Broadcastable的。

设它们广播后的共同形状为 $(S_1,S_2,\dots ,S_k)$。
那么高级索引的操作本质上是定义了一个映射 $f$：

$${\mathcal{B}}_{i_1,i_2,\dots ,i_k}={\mathcal{A}}_{{\mathcal{I}}_{i_1,\dots ,i_k}^{(0)},{\mathcal{I}}_{i_1,\dots ,i_k}^{(1)},\dots ,{\mathcal{I}}_{i_1,\dots ,i_k}^{(n-1)}},i_j\in S_j$$

核心逻辑：输出张量 $\mathcal{B}$ 在位置 $(i_1,\dots ,i_k)$ 的值，是由 $n$ 个索引张量在相同位置提供的坐标共同指向的。

案例分析

        edge_index = self.connect_edges(res_X, batch)  # 2, E
        n_nodes = res_H.shape[0]
        n_edges = edge_index.shape[1]
        n_channels = res_X.shape[1]
        n_heads = self.n_heads
        row, col = edge_index
 
        R_ij = torch.cdist(res_X[row], res_X[col], p=2)  # (E, 14, 3) , (E, 14, 3) -> (E, 14, 14)
        dist_rep = self.distance_expansion(R_ij).view(
            row.shape[0], res_X.shape[1], res_X.shape[1], -1
        )  # E, 14, 14, d_edge
        Q = self.W_Q(res_H).view([n_nodes, n_channels, self.n_heads, self.d])  # (N, 14, n_heads, d)
        K = self.W_K(res_H).view([n_nodes, n_channels, self.n_heads, self.d])
        V = self.W_V(res_H).view([n_nodes, n_channels, self.n_heads, self.d])
 
        attend_logits = torch.matmul(Q[row].transpose(1, 2), K[col].permute(0, 2, 3, 1)).view(
            [n_edges, self.n_heads, n_channels, n_channels]
        )  # (E, n_heads, 14, d) @ (E, n_heads, d, 14) -> (E, n_heads, 14, 14)
        attend_logits /= np.sqrt(self.d)  # n_edges, n_heads, 14, 14
        attend_logits = attend_logits + self.sigma_D(dist_rep).permute(0, 3, 1, 2)  # distance bias
        attend_mask = torch.ones(n_edges, self.n_heads, n_channels, n_channels, device=res_H.device, dtype=torch.bool)
 
        # sparse attention, only keep top k=3
        attend_logits[torch.logical_not(attend_mask)] = -1e5  # do not sellect from entries not attend
        _, top_indices = torch.topk(attend_logits, self.sparse_k, dim=-1, largest=True)  # (E, n_heads, 14, 3)
        sparse_mask = torch.zeros_like(attend_logits, dtype=torch.bool)  # (E, n_heads, 14, 14)
        rows = torch.arange(n_edges).view(-1, 1, 1, 1).expand(-1, n_heads, n_channels, self.sparse_k) # (E, n_heads, 14, 3)
        depth = torch.arange(n_heads).view(1, -1, 1, 1).expand(n_edges, -1, n_channels, self.sparse_k) # (E, n_heads, 14, 3)
        height = torch.arange(n_channels).view(1, 1, -1, 1).expand(n_edges, n_heads, -1, self.sparse_k) # (E, n_heads, 14, 3)
        sparse_mask[rows, depth, height, top_indices] = True
        attend_logits = attend_logits * sparse_mask
        attend_logits = attend_logits * attend_mask  # (E, n_heads, 14, 14)

上述代码本质为坐标映射的构造. 当执行 sparse_mask[rows, depth, height, top_indices] = True 时，实际上是在手动构建这个 $n$ 维坐标映射。

• 维度对齐为了对形状为 $(E,H,C,W)$ 的四维张量进行精确赋值，需要提供 4 个坐标分量：

  • Row 轴 ($d_0$): $I^{(0)}=\text{rows}$
  • Depth 轴 ($d_1$): $I^{(1)}=\text{depth}$
  • Height 轴 ($d_2$): $I^{(2)}=\text{height}$
  • Width 轴 ($d_3$): $I^{(3)}=\text{top\_indices}$

• 广播机制的数学实现使用的 torch.arange(…).view(…).expand(…) 本质上是在构造基矢量张量：

  • rows: 形状 $(E,1,1,1)\stackrel{expand}{\to }(E,H,C,K)$，满足 ${\text{rows}}_{e,h,c,k}=e$
  • depth: 形状 $(1,H,1,1)\stackrel{expand}{\to }(E,H,C,K)$，满足 ${\text{depth}}_{e,h,c,k}=h$
  • height: 形状 $(1,1,C,1)\stackrel{expand}{\to }(E,H,C,K)$，满足 ${\text{height}}_{e,h,c,k}=c$
  • 这样，当它们与形状为 $(E,H,C,K)$ 的 top*indices 结合时，对于任意确定的索引位置 $(e,h,c,k)$，其指向的目标坐标就是：(e, \ h, \ c, \ \text{top_indices}*{e,h,c,k})