Custom Polynomial Activation Function

mediumFREECustom Autograd FunctionsPyTorch Advanced

Problem Statement

Create a custom autograd function that implements a polynomial activation: f(x) = x² + 2x + 1.

Background

Sometimes you need activation functions not available in PyTorch. By subclassing torch.autograd.Function, you can define both the forward computation and its derivative for backpropagation.

Your Task

The starter code defines a PolyActivation class and test harness. Implement the forward method to compute f(x) = x² + 2x + 1, and the backward method to compute the correct derivative. Think about what the derivative of this polynomial is.

Output Format

The function returns a dictionary with "output" (activation values) and "grad" (input gradients).

Example

Input

None

Output

{'output': [4.0, 9.0, 16.0, 0.0], 'grad': [4.0, 6.0, 8.0, 0.0]}

Reasoning

Input: the harness builds a fixed 1-D tensor of four values, $1, 2, 3, -1$ , and reports two things per element: the activation value and the gradient flowing back to it.
Forward: $f(x) = x^2 + 2x + 1 = (x+1)^2$ , so $f(1) = 4$ , $f(2) = 9$ , $f(3) = 16$ , $f(-1) = 0$ , giving [4.0, 9.0, 16.0, 0.0].
Backward: the harness adds the four outputs into one scalar $L = \sum_i f(x_i)$ , so the upstream gradient at each output is $1$ . By the chain rule $\partial L / \partial x_i = 1 \cdot f'(x_i)$ with $f'(x) = 2x + 2$ : $4, 6, 8, 0$ , giving [4.0, 6.0, 8.0, 0.0].
Output: a dictionary with keys output and grad, each a list of floats rounded to 1 decimal place.
Notice that $x = -1$ is the parabola's minimum, so value and slope are both $0$ there. Also, the derivative must be evaluated at the original input $x$ , not at $f(x)$ : using $f(1) = 4$ would give $10$ instead of $4$ .

Constraints

Must use torch.autograd.Function
Forward: x^2 + 2x + 1
Backward: 2x + 2

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

The problem involves creating a custom autograd function in PyTorch, which is a fundamental concept in deep learning frameworks. Autograd is a system for automatically 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 functions are used to define custom operations that can be differentiated.

To understand this problem, you need to be familiar with the concept of activation functions in neural networks. Activation functions introduce non-linearity into the model, allowing it to learn more complex relationships between inputs and outputs. The given polynomial activation function f(x) = x² + 2x + 1 is a simple example of a non-linear function. You should also understand how to compute the derivative of this function, which is f'(x) = 2x + 2, as it will be used in the backward pass of the autograd function.

In PyTorch, custom autograd functions are defined by subclassing torch.autograd.Function and implementing the forward and backward methods. The forward method computes the output of the function, while the backward method computes the gradient of the output with respect to the input. This requires understanding of how to work with PyTorch tensors and how to use the torch.autograd module.

Algorithm/Approach

The general approach to solving this problem involves:

Defining a custom autograd function that implements the polynomial activation
Implementing the forward method to compute the output of the function
Implementing the backward method to compute the gradient of the output with respect to the input
Creating a tensor with requires_grad=True to track the gradient
Applying the custom autograd function to the tensor
Calling .sum().backward() to compute the gradient of the output with respect to the input

Step-by-Step Strategy

To implement the solution, follow these steps:

Define a custom autograd function by subclassing torch.autograd.Function.
Implement the forward method to compute the output of the polynomial activation function f(x) = x² + 2x + 1. Save the input tensor for use in the backward pass.
Implement the backward method to compute the gradient of the output with respect to the input. The gradient is given by grad_output * (2x + 2).
Create a tensor with requires_grad=True to track the gradient.
Apply the custom autograd function to the tensor.
Call .sum().backward() to compute the gradient of the output with respect to the input.
Return a dictionary with the output values and the gradient of the input.

Common Pitfalls

When implementing the solution, watch out for:

Forgetting to save the input tensor in the forward method for use in the backward pass
Incorrectly implementing the backward method, which can lead to incorrect gradients
Failing to set requires_grad=True when creating the input tensor, which will prevent the gradient from being computed

Time & Space Complexity

The time complexity of the solution is O(n), where n is the number of elements in the input tensor, since we need to iterate over all elements to compute the output and gradient. The space complexity is also O(n), since we need to store the input tensor and the output tensor. The custom autograd function will add a small overhead due to the additional computations required in the forward and backward methods.

solution.py

Python 3.13

Test results