Custom Gradient Scaling Function

mediumPROCustom Autograd FunctionsPyTorch Advanced

Problem Statement

Implement a gradient scaling function that acts as identity in the forward pass but scales the gradient by a factor in the backward pass.

Background

Gradient scaling is used in techniques like gradient reversal layers (for domain adaptation) or gradient scaling for multi-task learning. The forward pass is identity, but backward multiplies gradients by a scalar.

Your Task

The starter code defines a GradientScale class and test harness. Implement a function that passes input through unchanged in the forward pass but multiplies the gradient by a given scale factor in the backward pass. Note that the scale is a plain number, not a tensor.

The test applies scale=0.5 and scale=-1.0 (gradient reversal) to verify both cases.

Output Format

The function returns a dictionary with "grad_half" and "grad_reverse".

Example

Input

None

Output

{'grad_half': [0.5, 0.5, 0.5, 0.5], 'grad_reverse': [-1.0, -1.0, -1.0, -1.0]}

Reasoning

Input: the harness builds $x = [1.0, 2.0, 3.0, 4.0]$ . Each run passes $x$ and a plain number $s$ (first $s = 0.5$ , then $s = -1.0$ ) through GradientScale, then sums the result into the loss $L = \sum_i y_i$ .
Forward: the map is the identity, $y = x$ , so $y = [1, 2, 3, 4]$ and $L = 10$ . The value of $s$ changes nothing here.
Plain chain rule: $\frac{\partial L}{\partial y_i} = 1$ and $\frac{\partial y_i}{\partial x_i} = 1$ , so a plain identity gives $\frac{\partial L}{\partial x} = [1, 1, 1, 1]$ .
Scaled backward: this layer deliberately reports $s$ times the incoming gradient instead. With $s = 0.5$ : $0.5 \cdot 1 = 0.5$ per element, giving [0.5, 0.5, 0.5, 0.5]. With $s = -1.0$ : every entry flips sign, giving [-1.0, -1.0, -1.0, -1.0]. That is gradient reversal.
Notice: the gradient does not depend on the values $1, 2, 3, 4$ at all, only on $s$ . The harness resets the stored gradient between runs; otherwise the second result would accumulate to $0.5 + (-1.0) = -0.5$ .

Constraints

Forward must be identity (return input unchanged)
Backward must scale gradient by given factor
Store scale using ctx.scale, not save_for_backward (scale is not a tensor)

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 primerPREVIEW

Background Knowledge

Introduction to Autograd

Autograd is a key component of PyTorch, responsible for computing gradients of outputs with respect to inputs in a computational graph. This is crucial for training neural networks using backpropagation and gradient descent. In PyTorch, autograd is used to automatically compute gradients, allowing users to focus on defining the forward pass of their models.

Custom Autograd Functions

PyTorch provides the ability to define custom autograd functions by subclassing torch.autograd.Function. This allows users to implement custom forward and backward passes for specific operations, which can be useful for a variety of tasks, such as implementing new layer types or modifying the behavior of existing ones. Custom autograd functions are particularly useful when working with complex models or specialized techniques that require fine-grained control over the computation of gradients.

Gradient Scaling and Reversal

Gradient scaling and reversal are techniques used in various deep learning applications, including domain adaptation and multi-task learning. Gradient scaling involves multiplying the gradients of a model by a scalar value, while gradient reversal involves multiplying the gradients by a negative scalar value. These techniques can be used to modify the behavior of a model during training, allowing it to adapt to different tasks or datasets. In the context of this problem, we will implement a custom autograd function that applies gradient scaling to the input gradients during the backward pass.

Algorithm/Approach

The general approach to solving this problem involves defining a custom autograd function that applies gradient scaling to the input gradients during the backward pass. This will involve subclassing torch.autograd.Function and implementing the forward and backward methods. The forward method should apply an identity transformation to the input, while the backward method should multiply the input gradients by a scalar value.

Step-by-Step Strategy

To implement the solution, follow these steps:

Define a custom autograd function GradientScale by subclassing torch.autograd.Function.
Implement the forward method to apply an identity transformation to the input and save the scale factor.
Implement the backward method to multiply the input gradients by the saved scale factor.
Create an input tensor with requires_grad=True and apply the GradientScale function with a scale factor of 0.5.
Compute the gradients of the output with respect to the input using the sum and backward methods.
Store the computed gradients as grad_half.
Repeat steps 4-6 with a scale factor of -1.0 to compute the gradients for gradient reversal.
Store the computed gradients as grad_reverse.
Return a dictionary containing the computed gradients as lists.

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