Katkat
Back

Implement a Transformer Decoder from Scratch Without Starter Code

ByteDanceMachine Learning Engineer · Technical Screen
HardMachine Learning EngineerTechnical ScreenMachine Learning

Last updated

Implement a transformer decoder from scratch with no starter code in about 30 minutes, including embeddings, masked multi-head self-attention, feed-forward layers, residual connections and layer normalization. Tests tensor-shape fluency, causal masking and hands-on deep learning coding for a machine learning engineer interview.

With no starter code, no skeleton and no other scaffolding, implement a transformer decoder from scratch in about 30 minutes. Write it in Python with a tensor library such as PyTorch or NumPy, and explain the tensor shapes as you go.

The decoder takes a batch of token-ID sequences and returns next-token logits for every position. It should contain:

  • token embeddings plus positional information;
  • masked (causal) multi-head self-attention that you write yourself, so that each position attends only to itself and to earlier positions;
  • a position-wise feed-forward network;
  • residual connections and layer normalization around the sub-layers;
  • a stack of identical decoder blocks, followed by a projection to vocabulary logits.

Constraints and Clarifications

  • Input: integer token IDs of shape (batch, seq_len), with seq_len at most a maximum context length passed to the constructor.
  • Output: logits of shape (batch, seq_len, vocab_size). The logits at position t must depend only on tokens 0 through t.
  • Assume basic building blocks (linear layers, embedding tables, layer normalization, softmax) are allowed, but a ready-made attention or transformer module is not. Confirm this with the interviewer.
  • Model width, number of heads, number of layers and feed-forward width are constructor arguments.

Clarifying Questions

  • Is this a decoder-only language model, or the decoder half of an encoder-decoder transformer that also attends to encoder outputs?
  • Is a forward pass enough, or should I also write the training loss and a generation loop?
  • Can the inputs contain padding that must be masked as well?
  • Is there a preference between pre-norm and post-norm blocks, or between learned and fixed positional encodings?

Hints

Hint 1

Write the shapes down first: Before any code, write the shape of the activations entering and leaving each sub-layer and the shape of the attention-score tensor. Most bugs in this exercise are a wrong transpose or reshape.

Hint 2

Where the mask goes: Decide at which step of the attention computation the causal mask is applied, and what value a masked position takes, so that a later token can never influence an earlier position.

Hint 3

Heads without a loop: Think about how a single projection can be reshaped so that every head is computed by one batched matrix multiplication, and how the heads are merged back afterwards.

Follow-up questions

  • Why are the attention scores divided by the square root of the head dimension, and what happens without it?
  • How would you add a key-value cache for autoregressive generation, and what does it save?
  • What changes to turn this into the decoder of an encoder-decoder model?
  • How would you write a unit test that proves the causal mask works?