Debug Transformer and Add KV Cache
Company: OpenAI
Role: Machine Learning Engineer
Category: Machine Learning
Difficulty: medium
Interview Round: Onsite
You are given a small decoder-only transformer implementation for autoregressive language modeling.
### Part 1: Debugging
The training code contains four bugs. Fix them so that training loss decreases and generated text becomes reasonable. The known bug categories are:
1. The language-model loss uses the wrong label shift.
2. Positional embeddings are initialized incorrectly.
3. The attention mask is wrong.
4. There is one additional typo-related bug somewhere in the code.
After your fixes, the model should train normally, the loss should go down, and sample outputs should look plausible.
### Part 2: KV cache
A cache class for attention keys and values is provided. Extend the model to support cached autoregressive decoding.
Update the implementation so that:
1. The attention module reads from and appends to the cache correctly.
2. Causal masking is only applied when there is no cache or the cache is empty.
3. Positional encodings use the correct offset based on the current cache length.
4. The `generate` function performs incremental decoding with the cache.
Finally, verify that deterministic generation with and without KV cache produces the same token sequence.
Quick Answer: This question evaluates debugging and implementation skills for transformer-based autoregressive language models, focusing on attention mechanics, positional embeddings, causal masking, and integrating a key-value (KV) cache for incremental decoding.