The final step aims to select token embeddings according to the sparse attention index matrices to form the input tensors of a multi-head sparse attention Transformer [1]}.
Motivated by [1]}, we use reshaped inputs for efficient vectorized sparse attention computation.
More specifically, we select the embeddings of tokens that correspond to non-zero elements in the sparse attention index matrix to form the input embedding tensors.
For each attention head, the input tensor size is \(N\times K\times \frac{D}{h}\) , where \(D\)  is the total hidden dimension and \(h\)  is the number of attention heads.
If there are multiple Transformer layers, we apply the three steps described above to each layer.
We train the multi-head sparse attention Transformer in the target tasks, and we use it to generate the final predictions on the test samples.
