Gradient Clipping by Norm

easyFREEGradient ManipulationPyTorch Advanced

Problem Statement

Implement gradient clipping by norm to prevent exploding gradients.

Background

Norm-based gradient clipping scales all gradients so their total L2 norm doesn't exceed a threshold. This is essential for training RNNs and Transformers.

Your Task

The starter code creates a model, computes a loss, and calls backward. Compute the total gradient norm before clipping, apply norm-based gradient clipping with max_norm=1.0, then compute the norm again after clipping.

Output Format

Returns a dictionary with "norm_before", "norm_after", and "was_clipped".

Example

Input

None

Output

{'norm_before': 1463.896, 'norm_after': 1.0, 'was_clipped': True}

Reasoning

No arguments: a seeded linear layer maps 4 inputs to 2 outputs, with input $x = [10, 20, 30, 40]$ and target $[1, 1]$ . The prediction is $\hat y \approx [27.42, 5.00]$ , so the mean squared error is $\frac{26.42^2 + 4.00^2}{2} \approx 357.05$ .
Averaging over 2 outputs cancels the factor 2, so the bias gradient is $\hat y - y \approx [26.42, 4.00]$ . The weight gradient is its outer product with $x$ ; its first row is $26.42 \cdot [10, 20, 30, 40] \approx [264.2, 528.4, 792.7, 1056.9]$ .
The total norm treats all gradient entries as one long vector: $\sqrt{\lVert g_W \rVert^2 + \lVert g_b \rVert^2} = \sqrt{1463.652^2 + 26.723^2} \approx 1463.896$ , which is norm_before.
That exceeds the threshold $1.0$ , so every gradient is scaled by $1.0 / 1463.896 \approx 0.000683$ . Direction is kept and only the length shrinks, so the norm becomes $1.0$ (norm_after).
was_clipped is True because $1463.896 > 1.0$ . Notice that the large inputs inflate the weight gradient. A norm already below $1.0$ would be left alone, and both norms would match.

Constraints

Use torch.nn.utils.clip_grad_norm_
max_norm=1.0
Use MSE loss

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 Gradient Clipping

Gradient clipping is a technique used in deep learning to prevent exploding gradients, which can occur when training recurrent neural networks (RNNs) or transformers. Exploding gradients happen when the gradients of the loss function with respect to the model's parameters become very large, causing the model's weights to update excessively during backpropagation. This can lead to nan (not a number) or inf (infinity) values in the model's weights, causing the training process to fail.

Gradient Norm

The gradient norm is a measure of the magnitude of the gradients. It is calculated as the square root of the sum of the squares of the gradients. Gradient clipping by norm involves scaling the gradients so that their total norm does not exceed a certain threshold, known as the max norm. This helps to prevent exploding gradients and ensures that the model's weights are updated in a stable and controlled manner.

PyTorch Implementation

In PyTorch, gradient clipping by norm can be implemented using the torch.nn.utils.clip_grad_norm_ function. This function takes the model's parameters and the max norm as input, and scales the gradients so that their total norm does not exceed the max norm. The clip_grad_norm_ function modifies the gradients in-place, meaning that it changes the original gradients without returning a new tensor.

Algorithm/Approach

The general approach to solving this problem involves the following steps:

Create a simple model and seed the random number generator for reproducibility
Compute the loss and gradients using backpropagation
Record the total gradient norm before clipping
Clip the gradients using the clip_grad_norm_ function
Record the total gradient norm after clipping
Return the results, including the gradient norms and a boolean indicating whether the gradients were clipped

Step-by-Step Strategy

To implement the solution, follow these steps:

Create a simple model using nn.Linear(4, 2) and seed the random number generator using torch.manual_seed(42)
Create input and target tensors using the given values
Compute the MSE loss using the nn.MSELoss function and call backward to compute the gradients
Record the total gradient norm before clipping using the torch.norm function
Clip the gradients using the clip_grad_norm_ function with max_norm=1.0
Record the total gradient norm after clipping
Return a dictionary with the gradient norms and a boolean indicating whether the gradients were clipped

Common Pitfalls

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

Forgetting to seed the random number generator, which can cause the results to be non-reproducible
Using the wrong function to compute the gradient norm, such as torch.sum instead of torch.norm
Forgetting to call backward to compute the gradients before clipping
Using the wrong max_norm value, which can cause the gradients to be clipped too aggressively or not at all

Time & Space Complexity

The time complexity of this solution is O(n), where n is the number of parameters in the model, since we need to iterate over all parameters to compute the gradient norm and clip the gradients. The space complexity is O(1), since we only need to store a few scalar values (the gradient norms and the boolean indicating whether the gradients were clipped).

solution.py

Python 3.13

Test results