Apply Kaiming Normal Initialization

easyFREEWeight InitializationPyTorch Advanced

Problem Statement

Apply Kaiming (He) normal initialization, designed for ReLU networks.

Background

Kaiming initialization samples from N(0, std) where std = gain / sqrt(fan_in) with gain=sqrt(2) for ReLU. This prevents the variance from shrinking in deep ReLU networks.

Your Task

The starter code creates an nn.Linear(8, 4) layer. Apply Kaiming normal initialization designed for ReLU networks to the layer's weights.

The rest (computing statistics and comparing to theoretical std) is pre-filled.

Output Format

Returns a dictionary with "weight_shape", "weight_mean", "weight_std", "expected_std", and "std_close".

Example

Input

None

Output

{'weight_shape': [4, 8], 'weight_mean': 0.0267, 'weight_std': 0.5526, 'expected_std': 0.5, 'std_close': True}

Reasoning

There are no arguments: the function fixes the random seed, builds a linear layer mapping 8 input features to 4 outputs, and reports statistics of its weight matrix. That matrix has one row per output and one column per input, so weight_shape is [4, 8] and each output neuron sums over $fan\_in = 8$ inputs.
For ReLU the gain is $\sqrt{2}$ , so the target spread is $\sigma = \frac{\sqrt{2}}{\sqrt{8}} = \sqrt{2/8} = \sqrt{0.25} = 0.5$ , giving expected_std $= 0.5$ . Every weight should come from $\mathcal{N}(0, 0.5^2)$ .
The $4 \times 8 = 32$ drawn weights have sample mean $0.0267$ (near the target $0$ ) and sample standard deviation $0.5526$ . With only 32 draws the estimate drifts a little from $0.5$ .
The gap is $|0.5526 - 0.5| = 0.0526 < 0.2$ , so std_close is True.
Notice that $\sigma$ depends only on the input width: a layer with 32 inputs would target $\sqrt{2/32} = 0.25$ . Wider inputs get smaller weights, so the variance of each pre-activation stays about the same.

Constraints

Use nn.init.kaiming_normal_
mode='fan_in', nonlinearity='relu'
Compare actual vs theoretical std

Learn this in

PyTorch FundamentalsFoundations

More from PyTorch Advanced

Implement SGD with Momentumhard Implement Custom SGD Optimizerhard Weight Decay for Specific Layersmedium Freeze Backbone, Train Headmedium Per-Layer Learning Rateseasy

Concept primer

Background Knowledge

Introduction to Weight Initialization

Weight initialization is a crucial step in training neural networks. It involves setting the initial values of the model's weights, which can significantly impact the convergence and performance of the network. Different initialization techniques can be used, and the choice of technique depends on the type of network and the activation functions used.

Kaiming Initialization

Kaiming initialization, also known as He initialization, is a technique specifically designed for ReLU (Rectified Linear Unit) networks. It was introduced by Kaiming He et al. in their 2015 paper "Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification". The key idea behind Kaiming initialization is to initialize the weights in a way that prevents the variance of the activations from shrinking as the input passes through multiple layers. This is achieved by sampling the weights from a normal distribution with a standard deviation of $\frac{\text{gain}}{\sqrt{\text{fan_in}}}$ , where $\text{gain}$ is a constant that depends on the activation function, and $\text{fan_in}$ is the number of input features.

Importance of Initialization

Proper weight initialization is essential for training deep neural networks. If the weights are initialized too small, the activations may become too small, leading to vanishing gradients. On the other hand, if the weights are initialized too large, the activations may become too large, leading to exploding gradients. Kaiming initialization helps to avoid these issues by initializing the weights in a way that maintains a stable variance throughout the network.

Algorithm/Approach

The general approach to solving this problem involves:

Understanding the concept of Kaiming initialization and its importance in deep neural networks
Familiarity with PyTorch's nn.init.kaiming_normal_ function and its parameters
Ability to calculate the theoretical standard deviation of the weights using the Kaiming initialization formula
Comparing the actual standard deviation of the weights with the theoretical standard deviation to verify the correctness of the implementation

Step-by-Step Strategy

To solve this problem, follow these steps:

Import the necessary PyTorch modules and set the random seed using torch.manual_seed(42).
Create a PyTorch nn.Linear layer with the specified input and output dimensions.
Apply Kaiming normal initialization to the layer's weights using nn.init.kaiming_normal_.
Calculate the theoretical standard deviation of the weights using the Kaiming initialization formula.
Calculate the actual mean and standard deviation of the weights.
Compare the actual standard deviation with the theoretical standard deviation to determine if they are within a certain tolerance.

Common Pitfalls

Some common pitfalls to watch out for when implementing this solution include:

Forgetting to set the random seed, which can lead to different results each time the code is run
Using the wrong mode or nonlinearity when applying Kaiming initialization
Incorrectly calculating the theoretical standard deviation of the weights
Failing to compare the actual standard deviation with the theoretical standard deviation

Time & Space Complexity

The time complexity of this solution is O(1), since it involves a constant number of operations regardless of the input size. The space complexity is also O(1), since it only requires a constant amount of memory to store the weights and other variables. However, the space complexity of the PyTorch nn.Linear layer itself is O(n), where n is the number of input features, since it needs to store the weights and biases.

solution.py

Python 3.13

Test results