Masked Self-Attention, Explained Through Avengers: Endgame

If you've studied transformers, you've run into this sentence a dozen times:

"Masked self-attention prevents a token from attending to future positions."

Fine. But why does that matter, and what actually happens inside the model when you mask something? I found the cleanest way to internalize it wasn't through more equations. It was through a movie I already knew scene-by-scene: Avengers: Endgame.

Here's the full mapping, with the actual math and code underneath every metaphor.

---

1. Self-Attention = The Infinity Stones

In a standard self-attention layer, every token in a sequence can look at every other token, including ones ahead of it in the sequence. For each token, we compute:

Attention(Q, K, V) = softmax(QKᵀ / √d_k) V

Q (query), K (key), and V (value) are just learned projections of the input. The QKᵀ term produces a score between every pair of tokens - how much token i should "attend to" token j, for all i and j, regardless of order.

That's enormous power. Think of it like holding all six Infinity Stones at once: Space, Time, Mind, Power, Reality, Soul. Any token can reach into any part of the sequence, past, present, or future, and pull in whatever information helps it the most.

Powerful. Also, in one very specific context, dangerous.

2. Data Leakage = Thanos's Snap

Here's where it gets dangerous. Transformers used for autoregressive language modeling (GPT-style models) are trained to predict the next token given everything before it. The training objective looks like:

P(token_t | token_1, token_2, ..., token_t-1)

Now imagine during training, the attention mechanism has full, unmasked access to the entire sequence, including token_t, the exact word it's supposed to predict. The model doesn't need to learn anything. It just looks at the answer key sitting right there in the input and copies it.

This is data leakage, and it's a well-known failure mode any time a model has access to information it wouldn't have at inference time. During training you have the whole sequence sitting in memory; during real-world inference (actually generating text word by word) you obviously don't have the future yet. If your model trained as if it did, it will collapse the moment it has to generate something real.

Thanos does the same thing, structurally. He doesn't discover the future; he uses the Time Stone to rearrange the present so that his already-known outcome comes true. He's not solving anything. He's cheating the timeline. And just like a leaky model, that "solution" only holds up as long as the shortcut exists. The moment reality has to run forward on its own (his snap gets undone), everything falls apart, because nothing was actually learned. It was just forced.

3. Masked Self-Attention = Iron Man's Counter-Snap

The fix is almost insultingly simple, and that's what makes it elegant.

Before the softmax step, we take the raw attention score matrix and set every score that corresponds to a "future" position to negative infinity:

scores[i][j] = -∞   for all j > i

Then we apply softmax. Since softmax(-∞) = 0, every future position's contribution vanishes completely. Token i is left attending only to positions 0 through i, itself and everything before it. Nothing after.

This is done with a causal mask, and visually, it looks like this for a 5-token sequence, where v means "visible / allowed to attend" and x means "blocked, this is a future position":

        tok0  tok1  tok2  tok3  tok4
tok0  [  v     x     x     x     x  ]
tok1  [  v     v     x     x     x  ]
tok2  [  v     v     v     x     x  ]
tok3  [  v     v     v     v     x  ]
tok4  [  v     v     v     v     v  ]

Each row can only "see" columns up to and including its own position. This is the lower-triangular structure you'll see referred to constantly in transformer papers and code (tf.linalg.band_part, torch.triu, "causal mask," "look-ahead mask", all the same idea, just implemented differently depending on the framework).

The important part, and the reason the Endgame analogy actually holds up structurally and not just narratively: masking doesn't introduce a new mechanism. It's the exact same Q, K, V computation, the exact same softmax. You're not adding new machinery, you're constraining the existing one. Iron Man doesn't get a seventh stone. He uses the same six, wielded correctly, to undo what an unconstrained use of that same power caused.

4. The Code

Here's a minimal, from-scratch implementation of masked self-attention in TensorFlow:

import tensorflow as tf

def masked_self_attention(Q, K, V):
    """
    Q, K, V: tensors of shape (batch, seq_len, d_k)
    Returns: attention output of shape (batch, seq_len, d_k), attention weights
    """
    d_k = tf.cast(tf.shape(Q)[-1], tf.float32)

    # Raw attention scores: how much each token attends to every other token
    scores = tf.matmul(Q, K, transpose_b=True) / tf.sqrt(d_k)   # shape: (batch, seq_len, seq_len)

    # Build the causal mask: 1 where j > i (future positions)
    seq_len = tf.shape(scores)[-1]
    causal_mask = 1 - tf.linalg.band_part(tf.ones((seq_len, seq_len)), -1, 0)

    # Set future positions to -inf before softmax
    scores = scores + (causal_mask * -1e9)

    # After softmax, masked positions collapse to ~0: future is erased
    attn_weights = tf.nn.softmax(scores, axis=-1)

    return tf.matmul(attn_weights, V), attn_weights


# Quick sanity check
tf.random.set_seed(0)
seq_len, d_k = 5, 8
Q = tf.random.uniform((1, seq_len, d_k))
K = tf.random.uniform((1, seq_len, d_k))
V = tf.random.uniform((1, seq_len, d_k))

output, weights = masked_self_attention(Q, K, V)
print(tf.round(weights[0] * 100) / 100)

Running that print statement gives:

tf.Tensor(
[[1.   0.   0.   0.   0.  ]
 [0.52 0.48 0.   0.   0.  ]
 [0.36 0.36 0.28 0.   0.  ]
 [0.25 0.27 0.21 0.27 0.  ]
 [0.2  0.2  0.19 0.22 0.18]], shape=(5, 5), dtype=float32)

This is the causal mask made real. Every entry above the diagonal is exactly 0, exactly matching the x positions in the diagram above. Token 0 can only attend to itself, so its full weight (1.0) goes there. Token 4, the last one, can attend everywhere, so it's the only row with all five positions filled in. Theory and output line up exactly.

A quick note if you run this yourself: the exact decimal values (like 0.52, 0.48 in row 2) will likely be slightly different each time, since Q and K are randomly initialized and tf.random.set_seed() alone doesn't always guarantee identical values across every rerun or environment. That's expected and fine. What's fixed and guaranteed, every single time, is the zero pattern in the upper triangle. That's the mask doing its job, not the randomness. If you want fully reproducible numbers, pass an explicit seed argument to each tf.random.uniform() call individually.

5. Why This Matters Beyond the Metaphor

This isn't just a training-time implementation detail. It's the reason autoregressive generation works at all:

• At inference time, the model generates one token at a time, and by construction it only ever has access to past tokens. Masking during training makes sure the model never learns to rely on information it won't have later. Train/inference consistency is the whole point.
• Encoder vs. decoder difference: this is exactly why BERT (encoder-only, bidirectional) doesn't use causal masking. It's trained with masked language modeling (a different kind of masking, hiding random tokens, not future tokens), and it's allowed to see the full sequence in both directions. GPT-style decoders use causal masking because their whole objective depends on next-token prediction being honest.
• It's cheap: masking adds essentially zero computational cost. It's one matrix operation and a softmax that was happening anyway. The "fix" for a fundamental correctness problem is a few lines of code, not a redesign.

6. The Takeaway

Sometimes the fastest way to actually retain a technical concept is to hang it on a structure you already have memorized. The Endgame analogy isn't just decoration here; the structural parallel is genuinely tight:

• Same underlying power (attention / the Stones)
• Misused, it breaks the thing it's supposed to build (data leakage / the Snap)
• The fix isn't new machinery, it's the same mechanism, constrained correctly (masking / the counter-snap)

Causality isn't a footnote in transformer architecture. It's the constraint that makes autoregressive generation honest, the difference between a model that has learned to predict the future and one that just memorized it.

---

If you found this useful, I write about ML/DL concepts and the projects I'm building. Check out more at dev.to/ishita_garg. Feedback and corrections welcome, always happy to be told where the analogy breaks down.