The AI Interview - Master AI/ML Interviews

Implement a dropout layer that applies random neuron deactivation during training to prevent overfitting in neural networks. The layer should randomly zero out a proportion of input elements based on a dropout rate p, scale the remaining values by 1/(1-p) to maintain expected values, and pass inputs unchanged during inference. During backpropagation, gradients must be masked with the same dropout pattern and scaled by the same factor to ensure proper gradient flow.

The dropout mask should be regenerated on each call to forward during training, and stored as self.mask so it can be reused in the backward pass.

Input:

p: Dropout rate (probability of zeroing out a neuron), must be between 0 and 1 (exclusive of 1)
x: Input numpy array of any shape
training: Boolean flag indicating if in training mode
grad: Gradient numpy array during backpropagation (same shape as x)

Output:

forward: Array of same shape as input with dropout applied (training) or unchanged (inference)
backward: Gradient array with the same dropout mask applied

Examples

Example 1:

Input: x = np.array([1.0, 2.0, 3.0, 4.0]), grad = np.array([0.1, 0.2, 0.3, 0.4]), p = 0.5

Output: forward output = array([2., 0., 6., 0.]), backward output = array([0.2, 0., 0.6, 0.])

Explanation: The Dropout layer randomly zeroes out elements of the input tensor with probability p during training. To maintain the expected value of the activations, the remaining elements are scaled by a factor of 1 / (1 - p). During inference, Dropout is disabled and the input is passed through unchanged. During backpropagation, the same dropout mask and scaling are applied to the gradients, ensuring the expected gradient magnitude is preserved.

Starter Code

import numpy as np

class DropoutLayer:
    def __init__(self, p: float):
        """Initialize the dropout layer.
        
        Attributes to set:
            self.p: the dropout rate
            self.mask: stores the dropout mask (initially None)
        """
        # Your code here

    def forward(self, x: np.ndarray, training: bool = True) -> np.ndarray:
        """Forward pass of the dropout layer.
        
        Generate a new mask on each training forward pass and store it in self.mask.
        """
        # Your code here

    def backward(self, grad: np.ndarray) -> np.ndarray:
        """Backward pass of the dropout layer.
        
        Use the stored self.mask from the most recent forward pass.
        """
        # Your code here

Dropout Layer

Examples

Starter Code