DEV Community: Ishita Garg

Masked Self-Attention, Explained Through Avengers: Endgame

Ishita Garg — Sat, 15 Aug 2026 07:25:23 +0000

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.

I Compared 3 Ways to Do Transfer Learning - Here's What Actually Reduced Overfitting

Ishita Garg — Tue, 21 Jul 2026 09:22:35 +0000

If you've trained a deep learning model from scratch, you already know the two problems that hit you immediately: you need a lot of labeled data, and even if you have it, training takes forever. I ran into both while working on a cats vs. dogs image classifier, and it's what pushed me to actually understand transfer learning instead of just using it as a buzzword.

This post walks through four experiments I ran on the same dataset - a CNN I built and trained entirely from scratch, followed by three transfer learning approaches using VGG16 - to see how much transfer learning actually helps, and why the order you apply these techniques matters.

Why not just train your own model?

Deep learning models are data-hungry. To train something reliable from scratch, you typically need thousands of labeled images - and labeling isn't free. Someone (or something) has to go through every image and mark whether it's a cat or a dog, which costs time and, at scale, money.

Even with enough data, training a CNN from scratch on a reasonably large dataset takes significant compute time. Both of these - data scarcity and training cost - are exactly what transfer learning is designed to solve.

What is transfer learning, actually?

The simplest way I can put it: transfer learning means taking a model that already learned something useful from one problem, and reusing that knowledge on a different but related problem.

It's the same idea as learning to ride a bicycle before a motorcycle - the balance and coordination your brain already built don't get thrown away, they get reused. Or if you play violin, picking up guitar is easier because you already understand musical notes, rhythm, and practice discipline.

In deep learning, this looks like: take a CNN that's already been trained on a huge, general image dataset (like ImageNet, ~1.4 million images across 1000 categories), and reuse its learned visual understanding on your own, much smaller dataset.

Why does this actually work? Because convolutional layers learn hierarchically. Early layers pick up primitive features - edges, colors, simple textures - which are common across almost any real-world image. Later layers combine those into more complex, task-specific patterns. So the early general knowledge doesn't need to be relearned every time; only the final, task-specific classification layers need to adapt to your actual problem.

My experiment: Cats vs. Dogs, from scratch vs. transfer learning

I used the Cats vs Dogs dataset from Kaggle and ran four experiments: first a CNN trained entirely from scratch as my baseline, then three transfer learning approaches, each building on lessons from the last.

Baseline: A CNN trained from scratch (not transfer learning)

Before touching transfer learning at all, I trained my own CNN architecture from scratch, with no pretrained weights involved - this is the point of comparison for everything that follows.

My first version had no regularization. Result: 99% training accuracy, but only 78% validation accuracy, with validation loss climbing past 1.3 as training progressed. That gap is a textbook overfitting signature - the model was memorizing training images rather than learning generalizable features.

Adding dropout and batch normalization changed the picture: 87% train / 76% val, with validation loss staying controlled instead of diverging. Notice train accuracy actually dropped - that's expected and healthy. The model gave up some memorization in exchange for learning patterns that generalize better.

Even with regularization, a meaningful train-val gap remained. That gap is what pushed me toward transfer learning.

Transfer Learning Approach 1: VGG16 Feature Extraction

Here's where transfer learning actually starts. I used VGG16 - a CNN pretrained on ImageNet - as a frozen feature extractor. I removed VGG16's original classification head, froze its convolutional base entirely, and attached my own dense layers on top:

model.add(Dense(256, activation='relu'))
model.add(Dense(1, activation='sigmoid'))

Result: 98% train / 89% val accuracy - a big jump over the from-scratch baseline. But the train-val gap widened again, since the dense head was learning fast on top of already-excellent frozen features, with nothing to keep it in check.

Transfer Learning Approach 2: Feature Extraction + Data Augmentation

This was the most interesting result of the whole project. I applied the same frozen-VGG16 setup, but added image augmentation (rotation, flips, zoom) during training.

Result: 93% train / 92% val accuracy.

Look closely - train accuracy actually went down compared to the previous approach (98% → 93%), but validation accuracy went up (89% → 92%), and the train-val gap nearly disappeared entirely.

This is the clearest illustration I've seen of the difference between memorization and generalization. Augmentation forced the model to work harder on each training example (since it never sees the exact same image twice), which cost it some raw training accuracy - but that trade bought real generalization.

Transfer Learning Approach 3: Fine-tuning

For the final experiment, I unfroze VGG16's last convolutional block (block5) and trained it jointly with the dense head, using a much lower learning rate than before.

Two details mattered a lot here:

I only unfroze block5 after the dense head was already well-trained from the earlier approaches. Unfreezing pretrained layers before the head is trained would send large, noisy gradients backward and disrupt the pretrained weights - essentially undoing the benefit of transfer learning.
The low learning rate is deliberate. Fine-tuning pretrained weights with a normal learning rate risks catastrophic forgetting - wiping out the general knowledge the model already had.

Result: 99% train / 95% val accuracy - the best of all four experiments.

Results at a glance

Approach	Type	Train Acc	Val Acc	Key takeaway
Custom CNN (dropout + batch norm)	From scratch	87%	76%	Overfitting still visible
VGG16 Feature Extraction	Transfer learning	98%	89%	Big jump, but wider train-val gap
Feature Extraction + Data Augmentation	Transfer learning	93%	92%	Overfitting nearly eliminated
Fine-tuning (unfreezing block5)	Transfer learning	99%	95%	Best result overall

What I'd try next

Swap VGG16 for a more modern backbone (ResNet50, EfficientNetB0) and compare
Test the same pipeline with far less data, to see transfer learning's advantage more starkly in a genuinely low-data setting
Try discriminative learning rates - smaller LR for early unfrozen layers, larger for later ones - instead of one flat rate during fine-tuning

Code

All four notebooks, plus the results and graphs above, are on GitHub:
👉 transfer-learning-cats-vs-dogs

If you're learning deep learning too, I'd genuinely love to hear what you'd try differently - drop a comment below.