Backpropagation

import numpy as np
from typing import Dict, Tuple

class TwoLayerNeuralNetwork:
    """
    Two-Layer Fully Connected Neural Network with exact Chain Rule Backpropagation.
    Architecture: Input -> Dense(Hidden, ReLU) -> Dense(Output, Linear) -> MSE Loss.
    """
    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, lr: float = 0.05):
        self.lr = lr
        # He initialization for hidden layer
        self.W1 = np.random.randn(input_dim, hidden_dim) * np.sqrt(2.0 / input_dim)
        self.b1 = np.zeros((1, hidden_dim))
        # Xavier initialization for output layer
        self.W2 = np.random.randn(hidden_dim, output_dim) * np.sqrt(2.0 / (hidden_dim + output_dim))
        self.b2 = np.zeros((1, output_dim))

    def forward(self, X: np.ndarray) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:
        """Forward pass caching intermediate tensors."""
        z1 = np.dot(X, self.W1) + self.b1
        a1 = np.maximum(0.0, z1)  # ReLU
        z2 = np.dot(a1, self.W2) + self.b2
        y_hat = z2  # Linear output
        cache = {"X": X, "z1": z1, "a1": a1, "z2": z2, "y_hat": y_hat}
        return y_hat, cache

    def backward(self, cache: Dict[str, np.ndarray], y_true: np.ndarray) -> Dict[str, np.ndarray]:
        """
        Backward pass using Multivariable Chain Rule.
        Computes analytical gradients: dW1, db1, dW2, db2.
        """
        N = cache["X"].shape[0]
        y_hat = cache["y_hat"]
        a1 = cache["a1"]
        z1 = cache["z1"]
        X = cache["X"]

        # Step 1: Loss gradient w.r.t final pre-activation (MSE Loss: L = (1/2N) * sum((y_hat - y)^2))
        delta2 = (y_hat - y_true) / N  # [N, output_dim]

        # Step 2: Output layer parameter gradients
        dW2 = np.dot(a1.T, delta2)     # [hidden_dim, output_dim]
        db2 = np.sum(delta2, axis=0, keepdims=True)

        # Step 3: Backpropagate error delta to hidden layer through ReLU
        # delta1 = (delta2 @ W2.T) * relu'(z1)
        grad_a1 = np.dot(delta2, self.W2.T)  # [N, hidden_dim]
        delta1 = grad_a1 * (z1 > 0.0).astype(np.float64)  # [N, hidden_dim]

        # Step 4: Hidden layer parameter gradients
        dW1 = np.dot(X.T, delta1)      # [input_dim, hidden_dim]
        db1 = np.sum(delta1, axis=0, keepdims=True)

        return {"dW1": dW1, "db1": db1, "dW2": dW2, "db2": db2}

    def step(self, grads: Dict[str, np.ndarray]):
        """Apply gradient descent parameter updates."""
        self.W1 -= self.lr * grads["dW1"]
        self.b1 -= self.lr * grads["db1"]
        self.W2 -= self.lr * grads["dW2"]
        self.b2 -= self.lr * grads["db2"]

    def train_step(self, X: np.ndarray, y: np.ndarray) -> float:
        """Executes full Forward -> Backward -> Step training loop."""
        y_hat, cache = self.forward(X)
        loss = 0.5 * np.mean((y_hat - y) ** 2)
        grads = self.backward(cache, y)
        self.step(grads)
        return float(loss)