Implement decoder-only GPT-style transformer

Goal

Implement a simplified decoder-only Transformer language model (similar in spirit to GPT) for next-token prediction. The implementation should be modular, using four main classes:

1. MultiHeadAttention
2. FeedForward
3. DecoderLayer
4. GPT (the full model)

You may assume a deep learning framework (e.g., PyTorch or TensorFlow), but you should clearly specify tensor shapes and operations.

---

Model details and requirements

Assume:

• Batch size: B
• Sequence length: T
• Embedding dimension: d_model
• Number of attention heads: num_heads (assume d_model is divisible by num_heads)
• Vocabulary size: V

1. Multi-head self-attention (MultiHeadAttention)

Implement a class MultiHeadAttention that performs masked self-attention:

• Inputs:
  • x: Tensor of shape (B, T, d_model) (token embeddings or layer outputs).
• Internally learnable parameters:
  • Linear projections to compute queries Q, keys K, and values V for all heads.
• Operations:
  1. Project x into Q, K, and V with shapes (B, T, d_model).
  2. Split them into num_heads heads, each of dimension d_head = d_model / num_heads.
  3. Compute scaled dot-product attention for each head:

$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_{head}}} + M\right)V,$

 where `M` is a **causal mask** that prevents a position from attending to future positions (positions `> t`).

4. Concatenate the outputs of all heads back to shape (B, T, d_model). 5. Apply a final linear projection to return to (B, T, d_model).

• Output:
  • Tensor of shape (B, T, d_model).

Ensure you correctly implement the causal (autoregressive) mask so that position t can only attend to positions 0..t.

2. Position-wise feed-forward network (FeedForward)

Implement a class FeedForward for the position-wise MLP:

• Inputs:
  • x: Tensor of shape (B, T, d_model).
• Internal structure (typical choice):
  • Linear layer from d_model to d_ff (e.g., 4 * d_model).
  • Non-linear activation (e.g., GELU or ReLU).
  • Linear layer from d_ff back to d_model.
• All operations are applied independently to each position in the sequence.
• Output:
  • Tensor of shape (B, T, d_model).

3. Decoder layer (DecoderLayer)

Implement a single Transformer decoder block that combines attention, feed-forward, residual connections, and layer normalization:

• Inputs:
  • x: Tensor of shape (B, T, d_model).
• Components:
  1. LayerNorm ln1 before self-attention.
  2. MultiHeadAttention block (masked self-attention).
  3. Residual connection: x = x + attention_output.
  4. LayerNorm ln2 before feed-forward.
  5. FeedForward block.
  6. Residual connection: x = x + ff_output.
• Output:
  • Tensor of shape (B, T, d_model).

Your DecoderLayer should be reusable so that multiple layers can be stacked.

4. GPT model (GPT)

Implement a GPT class representing the full decoder-only Transformer:

• Inputs:

  • input_ids: Integer tensor of shape (B, T) representing token indices.

• Components:

  1. Token embedding layer mapping token IDs to vectors of size d_model.
  2. Positional encodings (learned or fixed) of shape (T, d_model) added to the token embeddings.
  3. A stack of N identical DecoderLayer blocks.
  4. A final linear layer mapping from d_model to vocabulary size V.

• Forward pass:

  1. Embed tokens, add positional encodings to obtain (B, T, d_model).
  2. Pass through the N decoder layers sequentially.
  3. Apply the final linear layer to obtain logits of shape (B, T, V).

• Output:

  • Logits for next-token prediction: (B, T, V).

Additional notes

• You should define the forward method for each class with correct tensor transformations.
• Be careful about tensor shapes when splitting/combining heads.
• Ensure the causal mask is correctly broadcast and applied in attention.
• You do not need to implement training loops or optimization; focus on correct and clean model implementation.

Part 1: Causal Multi-Head Self-Attention Forward Pass

Implement the forward pass of simplified GPT-style multi-head self-attention using nested Python lists instead of a deep learning framework. Given an input tensor x of shape B x T x d_model, projection matrices Wq, Wk, Wv, and Wo, compute causal scaled dot-product self-attention. Heads use contiguous chunks of the projected dimension. Position t may only attend to positions 0 through t. Use a numerically stable softmax and return every numeric value rounded to 6 decimal places.

Constraints

• 0 <= B <= 10
• 0 <= T <= 50
• 1 <= d_model <= 64 when T > 0
• 1 <= num_heads <= d_model
• d_model is divisible by num_heads
• All projection matrices are d_model x d_model
• Input values are finite numbers

Examples

Example 1

Input: ([], 1, [], [], [], [])
Output: []

Explanation: Empty batch produces an empty output.

Example 2

Input: ([[[0], [1]]], 1, [[1]], [[1]], [[1]], [[1]])
Output: [[[0], [0.731059]]]

Explanation: The second token attends to values [0, 1] with softmax scores [0, 1].

Hints

1. After projection, head h uses dimensions h * d_head through (h + 1) * d_head - 1.
2. For each query position t, compute scores only against key positions j <= t; future positions should not appear in the softmax denominator.

Part 2: Position-Wise Feed-Forward Network

Implement the forward pass of the Transformer position-wise feed-forward network using nested lists. For every token vector independently, compute hidden = ReLU(x * W1 + b1), then output = hidden * W2 + b2. The same feed-forward network is applied to each batch item and sequence position.

Constraints

• 0 <= B <= 100
• 0 <= T <= 100
• 1 <= d_model <= 128 when T > 0
• 1 <= d_ff <= 512
• All numeric values are finite
• Use ReLU: ReLU(z) = max(0, z)

Examples

Example 1

Input: ([[[1, 2]]], [[1, -1, 0], [0, 2, 1]], [0, 1, -3], [[1, 0], [0, 1], [1, 1]], [0, 0])
Output: [[[1, 4]]]

Explanation: The hidden vector is ReLU([1, 4, -1]) = [1, 4, 0].

Example 2

Input: ([[[2], [-1]]], [[2, -1]], [0, 3], [[1], [-2]], [5])
Output: [[[7], [-3]]]

Explanation: Both sequence positions use the same MLP independently.

Hints

1. There is no mixing across sequence positions; process each vector x[b][t] independently.
2. Treat each vector as a row vector: output[j] = sum_i hidden[i] * W2[i][j] + b2[j].

Part 3: Pre-Norm Transformer Decoder Layer

Implement one simplified GPT decoder layer. The layer uses pre-layer-normalization, causal multi-head self-attention, residual addition, another layer normalization, a ReLU feed-forward network, and a second residual addition. This problem is framework-free: tensors are nested lists, linear layers have no bias except the feed-forward biases, and LayerNorm is computed over the last dimension. For LayerNorm, use variance without epsilon; if variance is zero, use normalized value 0 for every feature before applying gamma and beta.

Constraints

• 0 <= B <= 10
• 0 <= T <= 50
• 1 <= d_model <= 64 when T > 0
• 1 <= d_ff <= 256
• d_model is divisible by num_heads
• LayerNorm gamma and beta each have length d_model
• Use causal masking inside self-attention

Examples

Example 1

Input: ([[[1, 3]]], 1, [1, 1], [0, 0], [[0, 0], [0, 0]], [[0, 0], [0, 0]], [[1, 0], [0, 1]], [[1, 0], [0, 1]], [1, 1], [0, 0], [[1, 0], [0, 1]], [0, 0], [[1, 0], [0, 1]], [0, 0])
Output: [[[0, 5]]]

Explanation: LN([1,3]) = [-1,1], attention adds [-1,1], then the feed-forward residual adds [0,1].

Example 2

Input: ([[[1, 3], [3, 1]]], 1, [1, 1], [0, 0], [[0, 0], [0, 0]], [[0, 0], [0, 0]], [[1, 0], [0, 1]], [[1, 0], [0, 1]], [1, 1], [0, 0], [[0, 0], [0, 0]], [0, 0], [[0, 0], [0, 0]], [0, 0])
Output: [[[0, 4], [3, 1]]]

Explanation: The second position causally averages normalized values from positions 0 and 1; the feed-forward block is zero.

Hints

1. This is a pre-norm block: normalize before attention and before the feed-forward network, not after the residual.
2. Keep the original x for the first residual: x1 = x + attention(LN1(x)), then output = x1 + feedforward(LN2(x1)).

Part 4: Full GPT-Style Decoder-Only Forward Pass

Implement a simplified full GPT forward pass using nested lists. First add token embeddings and positional embeddings. Then apply N pre-norm decoder layers, each using causal multi-head self-attention and a ReLU feed-forward network. Finally project each hidden vector to vocabulary logits. Layer parameters are provided as parallel arrays indexed by layer number. Return logits rounded to 6 decimals.

Constraints

• 0 <= B <= 10
• 0 <= T <= 50
• 1 <= vocabulary size V <= 500
• 1 <= d_model <= 64 when T > 0
• 0 <= N <= 12
• d_model is divisible by num_heads
• Every input id is in the range [0, V - 1]
• All numeric parameters are finite

Examples

Example 1

Input: ([[]], [[1, 2], [3, 4]], [], 1, [], [], [], [], [], [], [], [], [], [], [], [], [[1, 0], [0, 1]], [0, 0])
Output: [[]]

Explanation: An empty sequence has no logits.

Example 2

Input: ([[1]], [[1, 0], [0, 2]], [[10, 10]], 1, [], [], [], [], [], [], [], [], [], [], [], [], [[1, 0, 1], [0, 1, -1]], [0, 1, 0])
Output: [[[10, 13, -2]]]

Explanation: With no decoder layers, logits are computed directly from token plus positional embedding.

Hints

1. The initial hidden state is token_embedding[input_ids[b][t]] + position_embedding[t].
2. Apply decoder layers sequentially; the output of layer i becomes the input to layer i + 1.