Transformer Attention Lab — PyTorch

The question behind the work.

Can a minimal, from-scratch implementation demonstrate and verify the different information-access properties of causal and bidirectional self-attention?

Dataset & environment

Deterministic synthetic token-embedding tensors (2 batches × 7 tokens × 32 features). No language-model dataset or pretrained model is used.

From the idea to the method.

1. Project token embeddings into queries, keys and values and split into multiple attention heads.
2. Implement scaled dot-product attention and the causal lower-triangular mask.
3. Implement key padding masks and a residual TransformerBlock with normalization and feedforward sublayers.
4. Test tensor shapes, attention-row normalization, forbidden future attention, perturbation invariance and finite gradients.
5. Run seeded reproducible validation checks and serialize the test results.

What was implemented.

• Implemented a self-contained PyTorch multi-head attention module and transformer block for the portfolio.
• Created automated checks for causal prefix invariance, masked attention weights and backpropagation.

How the work is evaluated.

Six deterministic validation checks pass, including future-token masking and gradient finiteness. These are correctness checks, not accuracy metrics on an NLP benchmark.

What is documented.

6 implementation checks passed (seed 2026). Causal prefix change after modifying future tokens: 0.0; bidirectional prefix change: 14.519463 (maximum absolute value on this seeded synthetic test). These are illustrative numerical checks, not a benchmark.

Synthetic, seed-controlled experiments. Not a standard benchmark and not evidence of performance on real datasets.

A foundation for the next experiment.

• The causal mask controls which key positions each query may access.
• Transformer implementation correctness can be tested before a full language-model training run.