# Review a Compact GPT-Style Transformer Implementation You are reviewing a compact decoder-only Transformer implementation intended for next-token language modeling. Explain the expected tensor shapes and data flow, then identify the correctness checks you would perform in the attention, training-loss, and autoregressive-generation paths. You do not need to reproduce a specific framework's source code. ### Constraints & Assumptions - Inputs are integer token IDs shaped `[batch, time]` with `time` no greater than the configured context length. - The model uses token and positional embeddings, repeated pre-norm decoder blocks, multi-head causal self-attention, a feed-forward network, and a vocabulary projection. - The exact positional scheme is not fixed; learned absolute embeddings, rotary positions, and relative-position methods require different cropping and cache behavior. - Training predicts token `t+1` from positions through `t` using cross-entropy loss. - Generation may first be implemented without a key/value cache; discuss the cache as an optimization. ### Clarifying Questions to Ask - Are attention heads implemented separately or through a fused projection? - Are input embeddings and output projection weights tied? - Which positional scheme is used, and must generation support sequences beyond its trained context window? - Which dropout and normalization behavior differs between training and evaluation? - Must mixed precision and distributed training be reviewed? ### Part 1: Shapes and Causal Attention Trace shapes from embeddings through query/key/value projections, head splitting, attention scores, masking, head merging, and residual connections. #### What This Part Should Cover - Divisibility of model width by number of heads and correct transpose/reshape operations. - Scaled dot-product attention and a mask that prevents access to future keys. - Mask slicing for the actual sequence length, device/dtype compatibility, and residual shape preservation. ### Part 2: Training Objective Explain the shifted targets, flattening for cross-entropy, padding treatment if present, and checks for accidental information leakage. #### What This Part Should Cover - Logits and targets aligned so each position predicts the following token. - Correct vocabulary dimension and ignored padding where applicable. - Train/eval mode, deterministic small tests, and finite gradients. ### Part 3: Autoregressive Generation and Review Strategy Describe the generation loop, context cropping, sampling controls, and tests. Then explain how a key/value cache changes computation without changing outputs. #### What This Part Should Cover - Repeated next-token sampling conditioned only on the available prefix. - Temperature, top-k or greedy behavior, stop conditions, and seeded testing. - Cached keys/values per layer, with position and cropping rules tied explicitly to the chosen positional scheme. - The distinction between reindexing a recomputed learned-absolute window and maintaining cache-relative or rotary position offsets. - Parity tests against an uncached baseline using the same context policy. ### What a Strong Answer Covers - Exact shape reasoning and causal-mask semantics rather than Transformer buzzwords. - Correct next-token alignment, loss computation, and mode-specific behavior. - A disciplined code-review plan using tiny deterministic cases and invariants. - Honest complexity: full attention is quadratic in sequence length, while caching avoids recomputing prior projections during generation but does not remove all attention cost. ### Follow-up Questions - What symptoms would indicate the causal mask is transposed or applied with the wrong polarity? - Why can mixed-precision attention require numerically careful masking and softmax? - How would you test that cached and uncached generation agree token for token?