Causal vs. Bidirectional Architectu
The symptom was bizarre: during training, the validation loss looked incredible—almost too good to be true. But the second I hit model.generate(), the output was a repetitive loop of the same three tokens.
The error didn't throw a traditional Python traceback; it was a logical failure. I was seeing something like this in my logs:
# Expected output: "The weather is sunny"
# Actual output: "The the the the the"
# Logits for next token were skewed heavily toward the previous token
# because the model had "cheated" during training.I diagnosed it by stripping the model down and printing the attention maps. I noticed that the tokens at position $t$ were attending to tokens at position $t+1$ and beyond. In a causal architecture (like GPT), the mask should be a lower triangular matrix to prevent the model from "looking into the future." Because I had bidirectional attention enabled, the model wasn't learning to predict the next token based on history; it was simply learning to copy the token that already existed in the target sequence.
To fix this, I had to explicitly force a causal mask in the forward pass. Here is the snippet that actually solved the leakage:
import torch
def apply_causal_mask(attn_weights):
# Create a mask that blocks the upper triangle
mask = torch.triu(torch.ones(attn_weights.shape[-2:], device=attn_weights.device), diagonal=1).bool()
attn_weights.masked_fill_(mask, float('-inf'))
return attn_weightsThe trade-off here is a classic architectural headache. Bidirectional models (like BERT) are objectively better for understanding context because they see the whole sentence at once, which is why my training loss was so low. But the moment you move to generation, that "global" view becomes a liability. You can't use a bidirectional encoder to generate text token-by-token because the model expects to see the future tokens that don't exist yet.
Key takeaways from this mess:
Loss doesn't equal performance. A plummeting loss curve in a generative model is often a red flag for data leakage or masking errors rather than a sign of a "god-tier" model.
Attention visualization is non-negotiable. If you aren't plotting your attention heads when debugging, you're basically flying blind. I only found the bug once I saw the attention weights bleeding into the future indices.
Architecture choice is binary for generation. You either use a pure causal decoder or an encoder-decoder setup (like T5). Trying to "blend" them by toggling masks manually is a recipe for the kind of headache I just had.
All Replies (0)
No replies yet — be the first!
