Katkat
Back

Implement NumPy neural-network layers

OpenAIData Scientist · Onsite
MediumData ScientistOnsiteMachine Learning

Last updated

This question tests your ability to implement fundamental neural network components using NumPy, focusing on matrix operations, broadcasting, and numerical stability.

Problem Statement

Implement a basic fully connected (Dense) layer and a Softmax activation function from scratch using only NumPy. You should structure these as Python classes with forward methods.

Requirements

  1. Dense Layer: Implement a Dense class that takes input_dim and output_dim as initialization parameters. It should maintain weights and biases as attributes. The forward method should perform the operation: y = x @ W + b.
  2. Softmax Layer: Implement a Softmax class. The forward method should take a 2D array of logits x (shape: batch_size, num_classes) and return the probabilities.

Crucial: Your Softmax implementation must be numerically stable. You should subtract the maximum value across the feature dimension for each sample to prevent exponentiation of large numbers, which would lead to overflow.

Example

# Dense Layer Example
layer = Dense(input_dim=3, output_dim=2)
layer.weights = np.array([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]])
layer.biases = np.array([0.1, 0.1])

x = np.array([[1.0, 2.0, 3.0]])
output = layer.forward(x)
# Output: [[2.3, 2.9]]

Constraints

  • x is a 2D NumPy array of shape (batch_size, input_dim).
  • Use np.exp and np.sum for the Softmax logic.
  • Ensure your broadcasting is explicit and correct.

Hints

Hint 1

For the Dense layer, ensure your weight matrix initialization has dimensions (input_dim, output_dim) to allow for the dot product with the input.

Hint 2

For Softmax stability, remember that exp(x - max(x)) / sum(exp(x - max(x))) is mathematically equivalent to exp(x) / sum(exp(x)) but prevents overflow.

Hint 3

When computing the Softmax sum, use `keepdims=True` in `np.sum` to ensure the resulting array shape is compatible with the input for division broadcasting.

Follow-up questions

  • How would you implement the `backward` pass for these layers, specifically calculating the gradient with respect to the weights?
  • If you were to implement a ReLU activation layer, how would you handle the gradient calculation during the backward pass?
  • Why is it common to initialize weights with a small variance (e.g., Xavier initialization) rather than initializing all weights to zero?