Implement Softmax Cross-Entropy Forward and Backward Passes and a Training Loop

This is an "ML coding" round for a machine learning engineer role. Implement, from scratch, the forward pass and the backward pass of a softmax layer followed by cross-entropy loss. You may use an array library such as NumPy, but not automatic differentiation. As the follow-up, use your implementation to write a complete training loop. ### Clarifying Questions - What are the input shapes: a batch of logits of shape `(N, C)` for `N` examples and `C` classes? - Are labels integer class indices or one-hot vectors? - Is the loss averaged over the batch or summed? - In the training loop, what model produces the logits (for example a single linear layer), what data is given, and which optimizer and stopping rule are expected? ### Part 1 — Forward and backward Write `forward(logits, labels)` that returns the mean cross-entropy loss of the softmax probabilities against the labels, and `backward(...)` that returns the gradient of that loss with respect to the logits. Explain the gradient formula and show how you would convince yourself it is correct. ```hint Large logits Try your forward pass on logits such as 1000 and 0 before trusting it. ``` ```hint Differentiate the pair, not the pieces Work out the gradient of the combined softmax-plus-loss with respect to one logit; the result is much simpler than chaining the two Jacobians separately. ``` #### What This Part Should Cover - A numerically stable forward pass for softmax and for the log of the probabilities - The correct gradient with respect to the logits, including the batch averaging factor - Vectorized code without Python loops over examples or classes - A verification method, such as a finite-difference gradient check ### Part 2 — Training loop Using Part 1, write the training loop for a model that produces logits, for example a linear classifier `logits = X W + b`: batching, forward pass, loss, backward pass through the model's parameters, and parameter updates, over several epochs. ```hint Beyond the loss layer The gradient you computed stops at the logits; the parameters still need their own gradients. ``` #### What This Part Should Cover - Correct backpropagation from the logits to the model parameters - Minibatch iteration with reshuffling every epoch - A parameter update rule with a sensible learning rate and optional regularization - Monitoring the loss (and ideally validation accuracy) to confirm training works ### What a Strong Answer Covers - Clean separation between the loss layer, the model, and the loop - Shape bookkeeping stated out loud at each step - Numerical stability and a gradient check as evidence of correctness - Discussion of why deep learning libraries fuse softmax and cross-entropy into one operation ### Follow-up Questions - How do the forward and backward passes change with label smoothing? - How would you add per-class weights for an imbalanced dataset? - The loss becomes NaN after a few hundred steps. How do you debug it? - How would you extend the loop with a validation set and early stopping?