Implement SGD with Momentum

hardPROCustom OptimizersPyTorch Advanced

Problem Statement

Extend SGD with momentum, tracking velocity buffers in the optimizer state.

Background

Momentum SGD maintains a velocity: v = momentum * v + grad, then param -= lr * v. This smooths updates and helps escape local minima. The velocity buffer must persist across optimizer steps.

Your Task

The starter code defines a MomentumSGD class with init already implemented. Implement the step method that maintains a velocity buffer for each parameter and applies the momentum update rule.

The training loop and verification are pre-filled.

Output Format

Returns a dictionary with "initial_weight", "final_weight", "has_velocity", and "weight_changed".

Example

Input

None

Output

{'initial_weight': [0.7645, -0.8049, 0.2343], 'final_weight': [1.0577, 0.3795, 2.5983], 'has_velocity': True, 'weight_changed': True}

Reasoning

We initialize the MomentumSGD optimizer with a learning rate of 0.1 and momentum of 0.9, and create a nn.Linear(3, 1) model.
The initial weights of the model are recorded as [0.7645, -0.8049, 0.2343], which will be used to calculate the initial_weight output.
We run 5 training steps with MomentumSGD, updating the model parameters using the momentum update rule: $v = 0.9 \cdot v + \text{grad}$ , then $param -= 0.1 \cdot v$ , which changes the weights to [1.0577, 0.3795, 2.5983].
The optimizer state is checked for velocity buffers, and since MomentumSGD stores velocity in self.state[p], has_velocity is True, and since the weights have changed, weight_changed is also True.

Constraints

Store velocity in self.state[p]
Update: v = momentum * v + grad, param -= lr * v
Use @torch.no_grad() on step

Learn this in

PyTorch FundamentalsFoundations PretrainingLLM Study Plan

More from PyTorch Advanced

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

Concept primerPREVIEW

Background Knowledge

The problem involves implementing a custom optimizer in PyTorch, specifically Stochastic Gradient Descent (SGD) with Momentum. SGD is a fundamental optimization algorithm used in machine learning to minimize the loss function of a model. It iteratively updates the model's parameters in the direction of the negative gradient of the loss function. However, SGD can suffer from oscillations and slow convergence, especially in the presence of local minima. To address this, Momentum SGD introduces a momentum term that helps smooth the updates and escape local minima. The momentum term is calculated as a fraction of the previous update, which is added to the current update.

The Momentum SGD update rule can be written as: v = momentum * v + grad, where v is the velocity, momentum is the momentum coefficient, and grad is the gradient of the loss function. The parameter update is then param -= lr * v, where lr is the learning rate. This update rule helps the optimizer to "remember" the previous updates and make more informed decisions about the direction of the next update. The momentum term can be thought of as a "memory" that helps the optimizer to avoid oscillations and converge faster.

In the context of PyTorch, the torch.optim.Optimizer class provides a base class for implementing custom optimizers. To implement Momentum SGD, we need to define a custom optimizer class that inherits from torch.optim.Optimizer and overrides the init and step methods. The init method is used to initialize the optimizer's parameters, such as the learning rate and momentum coefficient. The step method is used to perform a single update step, which involves calculating the velocity and updating the parameters.

Algorithm/Approach

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

Define a custom optimizer class that inherits from torch.optim.Optimizer
Initialize the optimizer's parameters, such as the learning rate and momentum coefficient
Implement the step method to perform a single update step, which involves calculating the velocity and updating the parameters
Use the custom optimizer to train a model and verify that it works as expected

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