Weight Decay for Specific Layers

mediumPROCustom OptimizersPyTorch Advanced

Problem Statement

Apply weight decay only to weight parameters, not biases or normalization layers.

Background

Weight decay (L2 regularization) should typically only be applied to weight matrices, not bias terms or BatchNorm parameters. This is a common best practice.

Your Task

The starter code creates a model with Linear and BatchNorm layers. Separate the parameters into two groups: Linear weight matrices (which should get weight decay of 0.01) and everything else (no decay). Create an AdamW optimizer with these two parameter groups.

Output Format

Returns a dictionary with "decay_param_count", "no_decay_param_count", "decay_param_names", and "total_params".

Example

Input

None

Output

{'decay_param_count': 2, 'no_decay_param_count': 4, 'decay_param_names': ['0.weight', '2.weight'], 'total_params': 6}

Reasoning

The function selective_decay_test() creates a Sequential model with two Linear layers and one BatchNorm1d layer, resulting in 6 parameters: 2 weights and 2 biases from the Linear layers, and 2 parameters from the BatchNorm1d layer.
The parameters are then separated into two groups: decay (weight parameters of Linear layers) and no_decay (bias parameters and BatchNorm parameters), yielding 2 parameters in the decay group and 4 parameters in the no_decay group.
The names of the parameters with weight decay are 0.weight and 2.weight, corresponding to the weights of the two Linear layers.
The function returns a dictionary with the counts and names of parameters with and without weight decay, as well as the total parameter count, resulting in the output {'decay_param_count': 2, 'no_decay_param_count': 4, 'decay_param_names': ['0.weight', '2.weight'], 'total_params': 6}.

Constraints

Weight decay only on Linear weight matrices
No decay on biases and BatchNorm params
Use AdamW optimizer

More from PyTorch Advanced

Implement SGD with Momentumhard Implement Custom SGD Optimizerhard Freeze Backbone, Train Headmedium Per-Layer Learning Rateseasy Implement Bilinear Form with Einsumhard

Concept primerPREVIEW

Background Knowledge

Introduction to Weight Decay

Weight decay, also known as L2 regularization, is a technique used to prevent overfitting in neural networks. It works by adding a penalty term to the loss function that is proportional to the magnitude of the model's weights. This encourages the model to use smaller weights, which can help to reduce overfitting. The penalty term is typically proportional to the square of the weight values, hence the name L2 regularization.

Parameter Groups in Optimizers

In PyTorch, optimizers can be used with multiple parameter groups. This allows for different hyperparameters, such as learning rate and weight decay, to be applied to different sets of parameters. For example, it's common to apply weight decay only to the weights of linear layers, while leaving the biases and batch normalization parameters unchanged. This is because weight decay can help to prevent overfitting in the linear layers, but may not be beneficial for the biases and batch normalization parameters.

AdamW Optimizer

The AdamW optimizer is a variant of the Adam optimizer that includes weight decay. It is similar to the Adam optimizer, but with an additional term that applies weight decay to the parameters. The AdamW optimizer is often used in deep learning models because it can help to prevent overfitting and improve the stability of the training process.

Algorithm/Approach

The general approach to solving this problem is to:

Separate the model's parameters into two groups: one for parameters that should have weight decay applied, and one for parameters that should not.
Create an AdamW optimizer with the desired hyperparameters (e.g. learning rate, weight decay) for each parameter group.
Use the optimizer to update the model's parameters during training.

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.

PRO

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

Test results