I have some confusion about the implementation of the triangle mask in the transformer decoder. I understand the reasoning for the mask, it prevents the network from 'cheating' by looking ahead at the next token, but the way it's implemented seems strange. First of all, why do we apply the mask to the product $QK^T$? Why not just set the rows representing unseen tokens to zero or some special token before sending the matrix to the transformer? Also, I don't understand the reasoning of a triangle mask in general. Suppose we are training with a sequence of length $\ell_s$, and we are trying to predict token $k$ with $k < \ell_s$. Why don't we have to mask all entries $QK^T_{i,j}$ if $i > k$ or $j > k$? With a standard triangle mask it seems like we still incorporate information computed from 'unseen' tokens.


1 Answer 1


During training you use teacher forcing. This means that you feed the entire $tgt$ sequence to the decoder. You want the decoder to look only at $tgt[:i]$ when predicting $y[i]$. You then use $y[i]$ only to compare it with $tgt[i+1]$ to compute the loss. Thus for computing the output at position $i$ you want all the future tokens to be masked, but not erased. So when you calculate the attention scores between $tgt[i]$ and all the rest you want to zero the attention scores with the future tokens.

The attention scores is a square matrix of shape $l_s \times l_s$, where $l_s$ is the length of the target sequence. The entry $A_{i,j}$ gives the attention score between $tgt[i]$ and $tgt[j]$. Basically, row $i$ gives you the attention scores for token $tgt[i]$. So you want to zero out the scores when $i>j$. That is the reason why we use a triangular mask - simply zero out the upper triangular part of the attention scores matrix.

If you want to read a very detailed blog post about the Transformer model in general that I wrote, feel free to check this out:


You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .