The purpose of the triangular causal mask in the attention mechanism is to enforce the autoregressive property of the model during training and inference. This property ensures that the model can only attend to previous positions in the sequence, not future positions, in order to generate predictions sequentially.
While it's true that during inference you truncate the upper-left submatrix of the causal mask and omit future tokens, the causal mask still serves an important role. By setting the values in the causal mask to $0$ for positions that are "ahead" of the current position, you prevent the model from indirectly attending to those positions through subsequent attention steps. This constraint encourages the model to rely solely on the information available at each step, mimicking its behavior during training.
The purpose of the causal mask is not solely to prevent cheating or looking into the ground truth, but rather to ensure that the model learns to make predictions based only on the information it has seen so far. Without the causal mask, the model could potentially attend to future tokens or create dependencies between tokens that should be independent, leading to incorrect and less reliable predictions.
Regarding the second part of your question, the second row $[1, 1, 0, 0]$ in the causal mask is necessary to prevent the model from attending to tokens that come after the current position. Even though you're predicting the 5th token, the model should not have access to information from the 3rd and 4th tokens because it should make predictions based on the preceding tokens only. Allowing available tokens to attend to each other would violate the autoregressive property and potentially introduce information leakage from future tokens, leading to incorrect predictions.
Example: predicting optimal time to call reinforcments based on the history of enemy attacks
Consider a scenario where you have a military convoy moving through a battlefield, and you want to predict the optimal time to call for reinforcements based on the history of enemy attacks encountered so far. Each attack can be represented by a sequence of features, such as time of occurrence, location, and intensity.
To train a model to make accurate predictions, you feed it with a sequence of past attacks and corresponding features, up to a certain context length 𝑇. The model's task is to predict whether or not reinforcements should be called for the next attack.
Now, during inference, let's say you have the previous 4 attacks and their features, and you want to predict whether to call for reinforcements for the 5th attack. You truncate the upper-left submatrix of the causal mask, as you mentioned, to create a triangular mask that prevents the model from attending to future attacks.
The reason you still need the causal mask in this case is to ensure that the model learns to make predictions solely based on the historical information it has seen so far. Even though you didn't provide the 5th attack as input, the model should not indirectly attend to that attack through subsequent attention steps.
By using the causal mask, you enforce the constraint that the model can only attend to the past attacks and their features, simulating the real-world scenario where decisions about reinforcement need to be made based on historical data. This prevents the model from "cheating" by peeking at future attacks that would not be available in real-time decision-making.