Katkat
Back

Implement a Transformer Block with SwiGLU

GoogleMachine Learning Engineer · Technical Screen
MediumMachine Learning EngineerTechnical ScreenCoding & Algorithms

Last updated

This question evaluates implementation and reasoning skills for Transformer internals, testing competency in multi-head self-attention, SwiGLU gated feed-forward layers, residual connections, and layer normalization within the Machine Learning / deep learning domain.

Implement a Transformer-style neural network block in Python using either NumPy or PyTorch.

Your implementation should include:

  1. Multi-head self-attention
  • Input tensor shape: (batch_size, sequence_length, hidden_dim).
  • Split the hidden dimension into num_heads attention heads.
  • Compute scaled dot-product attention.
  • Support an optional attention mask.
  • Concatenate heads and apply an output projection.
  1. SwiGLU feed-forward network
  • Implement a gated feed-forward layer using the SwiGLU formulation: SwiGLU(x) = Swish(xW_gate) * (xW_value), followed by an output projection.
  • The gate and value projections should be computed in parallel.
  • Use Swish(z) = z * sigmoid(z).
  1. Transformer block structure
  • Include residual connections.
  • Include layer normalization.
  • Return an output tensor with the same shape as the input.

Be prepared to explain tensor shapes, masking behavior, and the time and memory complexity of multi-head attention.