Neural Mastery
← Practice
Two-Layer MLP Training Stephard

Two-Layer MLP Training Step

hard Neural Networks by Hand 25 min

NumPy allowed · plain Python earns 10 bonus points

Execute an end-to-end training iteration: forward pass with ReLU, Softmax cross-entropy, backpropagation, and gradient updates.

Related Skills

Task

Implement two_layer_mlp_train_step(...): for batch x and targets y, compute hidden activations a1 = ReLU(x @ w1 + b1). Compute logits z2 = a1 @ w2 + b2 and softmax probabilities p. Compute cross-entropy loss -(1/N) sum(log(p_i[y_i])). Backpropagate gradients: dlogits = (p - one_hot(y)) / N, dw2 = a1.T @ dlogits, db2 = sum(dlogits). Backpropagate through ReLU: da1 = dlogits @ w2.T, dz1 = da1 * (z1 > 0), dw1 = x.T @ dz1, db1 = sum(dz1). Update parameters using SGD: param_new = param - lr * dparam. Return {'loss': ..., 'w1': ..., 'b1': ..., 'w2': ..., 'b2': ...} with all floats rounded to 4 decimals.

Function Signature

two_layer_mlp_train_step(x: list[list[float]], y: list[int], w1: list[list[float]], b1: list[float], w2: list[list[float]], b2: list[float], lr: float = 0.1) -> dict

Examples

Example 1: Single sample 2-input 2-hidden 2-class
Input: x = [[1, -1]] y = [0] w1 = [[0.5, -0.5], [0.2, 0.3]] b1 = [0.1, 0] w2 = [[0.4, -0.4], [0.1, 0.2]] b2 = [0, 0] lr = 0.1
Output: {'loss': 0.5459, 'w1': [[0.5337, -0.5], [0.1663, 0.3]], 'b1': [0.1337, 0], 'w2': [[0.4168, -0.4168], [0.1, 0.2]], 'b2': [0.0421, -0.0421]}
Example 2: Batch of 2 samples XOR-like
Input: x = [[0, 1], [1, 0]] y = [1, 0] w1 = [[0.1, 0.5], [-0.2, 0.4]] b1 = [0, 0.1] w2 = [[0.5, -0.5], [-0.2, 0.3]] b2 = [0.1, -0.1] lr = 0.05
Output: {'loss': 0.6808, 'w1': [[0.1125, 0.4938], [-0.2, 0.4061]], 'b1': [0.0125, 0.0998], 'w2': [[0.5012, -0.5012], [-0.1986, 0.2986]], 'b2': [0.1003, -0.1003]}
Example 3: 3-class classification step
Input: x = [[1, 2]] y = [2] w1 = [[0.3, -0.2], [0.1, 0.4]] b1 = [0, 0] w2 = [[0.1, 0.2, -0.3], [0.2, -0.1, 0.4]] b2 = [0, 0, 0] lr = 0.2
Output: {'loss': 1.1101, 'w1': [[0.2401, -0.1544], [-0.0198, 0.4913]], 'b1': [-0.0599, 0.0456], 'w2': [[0.0643, 0.1687, -0.233], [0.1572, -0.1376, 0.4805]], 'b2': [-0.0714, -0.0627, 0.1341]}

Constraints

  • 1 <= N <= 50
  • 1 <= D_in, H, C <= 32
  • 0 <= y[i] < C
  • lr > 0.0
Python3Saved ✓
def two_layer_mlp_train_step(x: list[list[float]], y: list[int], w1: list[list[float]], b1: list[float], w2: list[list[float]], b2: list[float], lr: float = 0.1) -> dict:
"""Execute complete forward pass, loss computation, backpropagation, and SGD update for a 2-layer MLP.

Args:
x: Input matrix of shape (N, D_in).
y: Ground truth integer class indices of length N.
w1: First layer weights of shape (D_in, H).
b1: First layer bias of length H.
w2: Second layer weights of shape (H, C).
b2: Second layer bias of length C.
lr: Learning rate.

Returns:
Dictionary with 'loss' (float) and updated 'w1', 'b1', 'w2', 'b2', all rounded to 4 decimals.
"""
pass

Case 1: Single sample 2-input 2-hidden 2-class
Input: x = [[1, -1]] y = [0] w1 = [[0.5, -0.5], [0.2, 0.3]] b1 = [0.1, 0] w2 = [[0.4, -0.4], [0.1, 0.2]] b2 = [0, 0] lr = 0.1
Expected: {'loss': 0.5459, 'w1': [[0.5337, -0.5], [0.1663, 0.3]], 'b1': [0.1337, 0], 'w2': [[0.4168, -0.4168], [0.1, 0.2]], 'b2': [0.0421, -0.0421]}
Case 2: Batch of 2 samples XOR-like
Input: x = [[0, 1], [1, 0]] y = [1, 0] w1 = [[0.1, 0.5], [-0.2, 0.4]] b1 = [0, 0.1] w2 = [[0.5, -0.5], [-0.2, 0.3]] b2 = [0.1, -0.1] lr = 0.05
Expected: {'loss': 0.6808, 'w1': [[0.1125, 0.4938], [-0.2, 0.4061]], 'b1': [0.0125, 0.0998], 'w2': [[0.5012, -0.5012], [-0.1986, 0.2986]], 'b2': [0.1003, -0.1003]}
Case 3: 3-class classification step
Input: x = [[1, 2]] y = [2] w1 = [[0.3, -0.2], [0.1, 0.4]] b1 = [0, 0] w2 = [[0.1, 0.2, -0.3], [0.2, -0.1, 0.4]] b2 = [0, 0, 0] lr = 0.2
Expected: {'loss': 1.1101, 'w1': [[0.2401, -0.1544], [-0.0198, 0.4913]], 'b1': [-0.0599, 0.0456], 'w2': [[0.0643, 0.1687, -0.233], [0.1572, -0.1376, 0.4805]], 'b2': [-0.0714, -0.0627, 0.1341]}