What if we drop the causal mask in auto-regressive Transformer?

Question

I understand the triangular causal mask in the attention is used to prevent tokens from "looking into the future", but why do we want to prevent that?

Let's suppose we have a model with context length $T = 8$. At inference time, we want to predict the 5th token with the previous 4 ones, so we truncate the upper-left submatrix and the causal mask looks like the following

[1, 0, 0, 0]
[1, 1, 0, 0]
[1, 1, 1, 0]
[1, 1, 1, 1]

In this case, the model has no access to the 5th token because we did not feed it into the model in the first place, so it certainly cannot cheat by peeking ahead at the ground truth. Why do we still need the causal mask?

Moreover, the second row [1, 1, 0, 0] prevents the 2nd token from "attending to" the 3rd and 4th tokens. This makes sense if we are trying to predict the 3rd token, but we are actually predicting the 5th one. Why don't we allow available tokens to attend to each other?

Answer 1

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.

Answer 2

At inference time you only compute one embedding at a time, and you cache computed embeddings.

So let's say you are trying to predict the 5th token. While predicting the 4th token you used the embeddings of your first three tokens. You have those embeddings cached so now you only need to compute the embedding of the 4th token. Obviously the embedding of the 4-th token will have access only to the left context, that is all tokens with index i <= 4. Once you compute the embedding of that token you cache it for the next iteration.

You can see that using this procedure you actually emulate the causal mask, because all embeddings will be computed by only attending to tokens that came before them.

You could in fact recompute all of the embeddings at every step of the loop. So when predicting the 5th token instead of using the cached embeddings for tokens 1 to 3, you could recompute them so now they will have access to the full context, instead of only the left-context. The embedding for token 2 will be computed by attending to tokens 3 and 4, which is perfectly fine and in fact does not violate the auto-regressive property.

However, this causes another problem. Note that your model was not trained that way. During training each embedding only had access to the left context during prediction, and during inference you give access to the full context. This creates a training-inference mismatch which could really degrade the performance of your model.