import torch
import matplotlib.pyplot as plt
import numpy as np

What is Automatic Differentiation?

Training a neural network requires computing gradients (i.e., the derivatives of a loss function with respect to every model parameter or weight). For a model with millions of weights, doing this by hand is impractical. Automatic differentiation (autodiff, also known as autograd) solves this problem by letting the computer compute the derivatives for us automatically.

While finite differences, symbolic differentiation and automatic differentiation all compute derivatives, they do so in fundamentally different ways, as shown in the following table:

Method	How it works	Drawback
Finite differences	Approximates \frac{\partial f}{\partial x_i} \approx \frac{f(x_i + h) - f(x_i)}{h} for a small h	Requires a separate forward pass per parameter and could be very slow and numerically unstable.
Symbolic differentiation	Applies algebraic rules to produce an expression for the derivative (like Mathematica)	Expressions can blow up in size for complex functions, with many redundant sub-expressions.
Automatic differentiation	Records the sequence of elementary operations in a computational graph, then applies the chain rule at each node	Exact, efficient, and scales to arbitrarily complex programs.

Unlike finite differences or symbolic differentiation, automatic differentiation breaks any computation into a sequence of primitive operations (addition, multiplication, exponential, etc.) and records them in a directed graph. Then:

Forward pass: evaluate the graph left-to-right to compute the output.
Backward pass: walk the graph right-to-left, multiplying local derivatives at each node via the chain rule to accumulate the gradient of the output with respect to every input.

This “backward” flavour is called reverse-mode autodiff, and it is exactly what backpropagation implements. It computes the gradient of a single scalar loss with respect to all parameters in roughly the same time as one forward pass, which is why it scales so well to large neural networks.

Automatic Differentiation in PyTorch

PyTorch’s autograd engine is a reverse-mode autodiff system. The workflow is:

Create tensors with requires_grad=True so PyTorch records operations on them.
Run the forward computation as PyTorch silently builds the computational graph.
Call .backward() on the scalar loss, where PyTorch walks the graph in reverse, applying the chain rule at every node.
Read the accumulated gradients from each tensor’s .grad attribute.

The example below demonstrates this concept on the simplest possible case: y = x^2, where the derivative is \frac{dy}{dx} = 2x.

%%{init: {
    'theme': 'mc',
    'themeVariables': {
        'fontSize': '13px',
        'primaryColor': '#2d6a4f',
        'lineColor': '#52b788'},
        'flowchart': {'nodeSpacing': 30, 'rankSpacing': 50, 'curve': 'basis'}}}%%
flowchart RL
    y(["y"]) -->|"∂y/∂x = 2x"| xsq(["x²"])

    xsq~~~ df1

    df1["x = 18.9
    ∂y/∂x = 2x = 37.8"]

    style df1 fill:#e6e8e7,stroke:#e6e8e7,color:#0d0900
    style y fill:#1b4332,stroke:#52b788,color:#fff
    style xsq fill:#1b4332,stroke:#52b788,color:#fff

An example of computational graph for y = x^2 at x = 18.9

How Does It Work?

x = torch.tensor(18.9, requires_grad=True)
y = x * x
y.backward()
print("x =", x)
print("y =", y)
print("x.grad =", x.grad)

x = tensor(18.9000, requires_grad=True)
y = tensor(357.2100, grad_fn=<MulBackward0>)
x.grad = tensor(37.8000)

x = torch.tensor(18.9, requires_grad=True): Creates a scalar tensor x = 18.9 with gradient tracking enabled. The requires_grad=True flag tells PyTorch to record all operations on x so it can later compute derivatives.
y = x * x: Computes y = x^2. PyTorch builds a computational graph behind the scenes, recording that y was derived from x via multiplication.
y.backward(): Triggers backpropagation: PyTorch walks the computational graph in reverse and computes \frac{dy}{dx}. Since y = x^2, the derivative is y' = 2x.
x.grad: After .backward() is called, the gradient \frac{dy}{dx} is stored in x.grad. With x = 18.9, this gives 2 \times 18.9 = 37.8.

Using Automatic Differentiation for Polynomial Regression

Theoretical Background

Automatic differentiation is not just for neural networks. It can be used to compute gradients for any differentiable function, including simple polynomial regression. For example, we know that the polynomial y = x^2 + 5x + 3 generates our data, but imagine we don’t know the coefficients, and we only have the samples (X, Y). Can we recover [1, 5, 3] automatically?

This is where automatic differentiation can be used to compute gradients and recover unknown coefficients. The idea:

Model: Assume y = w_1 x^2 + w_2 x + w_3. We stack the inputs into a matrix \mathbf{X} = [x^2 \;\; x \;\; 1] so the prediction is simply \hat{y} = \mathbf{X} \mathbf{w}.
Loss: Define a loss function. Here, we are using mean squared error (MSE) that measures how far the predictions \hat{y} are from the true y: \text{MSE} = \frac{1}{N}\sum_{i=1}^{N}(y_i - \hat{y}_i)^2
Autograd: Because w is created with requires_grad=True, PyTorch tracks every operation involving w. When we call mse.backward(), it automatically computes \frac{\partial \, \text{MSE}}{\partial \mathbf{w}}, the gradient of the loss with respect to each coefficient.
Optimizer: The optimizer (here, NAdam) uses those gradients to update w in the direction that reduces the loss. Repeating this for 1000 iterations would push w toward the true coefficients [1, 5, 3].

Without autograd, we would need to derive and code the gradient formula by hand. Autograd does this for us, no matter how complex the computation graph becomes, which is exactly why autograd is the backbone of neural network training.

Computational Graph for Polynomial Regression

Automatic differentiation breaks a complex expression into elementary operations and chains their derivatives together. Below is the computational graph for our polynomial y = x^2 + 5x + 3, evaluated at x = 2.

%%{init: {
    'theme': 'mc',
    'themeVariables': {
        'fontSize': '13px',
        'primaryColor': '#2d6a4f',
        'lineColor': '#52b788'},
        'flowchart': {'nodeSpacing': 30, 'rankSpacing': 50, 'curve': 'basis'}}}%%
flowchart RL
    y(["y = v₃ + 3 = 17"]) -->|"∂y/∂v₃ = 1"| v3(["v₃ = v₁ + v₂ = 14"])
    v3 -->|"∂v₃/∂v₁ = 1"| v1(["v₁ = x² = 4"])
    v3 -->|"∂v₃/∂v₂ = 1"| v2(["v₂ = 5x = 10"])
    v1 -->|"∂v₁/∂x = 2x = 4"| x(["x = 2"])
    v2 -->|"∂v₂/∂x = 5"| x

    x~~~ df1

    df1["x = 2
    dy/dx = 2x + 5 = 9"]

    style df1 fill:#e6e8e7,stroke:#e6e8e7,color:#0d0900
    style y fill:#1b4332,stroke:#52b788,color:#fff
    style v3 fill:#1b4332,stroke:#52b788,color:#fff
    style v1 fill:#1b4332,stroke:#52b788,color:#fff
    style v2 fill:#1b4332,stroke:#52b788,color:#fff
    style x fill:#2d6a4f,stroke:#52b788,color:#fff

An example of computational graph for y = x^2 + 5x + 3 at x = 2

Forward pass (left → right) compute the value:

Node	Operation	Value
x	input	2
v_1	x^2	4
v_2	5x	10
v_3	v_1 + v_2	14
y	v_3 + 3	17

Backward pass (right → left) compute the gradient \dfrac{dy}{dx} via the chain rule:

Path	Chain rule	Value
y \to v_3	\dfrac{\partial y}{\partial v_3} = 1	1
v_3 \to v_1 \to x	\dfrac{\partial y}{\partial v_3} \cdot \dfrac{\partial v_3}{\partial v_1} \cdot \dfrac{\partial v_1}{\partial x} = 1 \cdot 1 \cdot 2x	4
v_3 \to v_2 \to x	\dfrac{\partial y}{\partial v_3} \cdot \dfrac{\partial v_3}{\partial v_2} \cdot \dfrac{\partial v_2}{\partial x} = 1 \cdot 1 \cdot 5	5
Total	Sum both paths	4 + 5 = \mathbf{9}

Since both paths from y back to x contribute, the total gradient is \dfrac{dy}{dx} = 4 + 5 = 9. This matches the analytical derivative y' = 2x + 5 = 2(2) + 5 = 9.

Note

PyTorch’s autograd does exactly this, where it records the forward operations into a graph, then walks backward through it applying the chain rule at every node. In our polynomial regression, this same mechanism computes \frac{\partial\,\text{MSE}}{\partial\mathbf{w}} automatically, no matter how many operations are chained together.

Step-1: Generating the Training Data

The code below creates a synthetic dataset of (x, y) pairs from a known polynomial, simulating a scenario where we have observed data and want to recover the underlying relationship.

polynomial = np.poly1d([1, 5, 3]): Defines the ground-truth polynomial y = x^2 + 5x + 3. The coefficients [1, 5, 3] are listed in descending power order (x^2, x^1, x^0).
N = 20: The number of data samples to generate.
X = np.random.randn(N,1) * 5: Draws 20 random x-values from a normal distribution (mean 0, std 5), giving a spread roughly in [-15, +15]. The shape (N, 1) makes X a column vector.
Y = polynomial(X): Evaluates the polynomial at each x-value to produce the corresponding y-values. These are the “labels” our model will try to predict.

# creates a NumPy polynomial object from 
# its coefficients [1, 5, 3] 
# in descending power order:
polynomial = np.poly1d([1, 5, 3])
print(polynomial)

# number of samples
N = 20

# The samples are generated from a standard normal distribution 
# (mean 0, std 1), multiplying by 5 scales the standard 
# deviation from 1 to 5, 
# so the samples are spread roughly in the range 
# [−15,+15] (since ~99.7% of values fall within ±3σ = ±15).
X = np.random.randn(N,1) * 5
Y = polynomial(X)

for i in range(N):
    print(f"X[{i}] = {X[i][0]:.2f}, Y[{i}] = {Y[i][0]:.2f}")

w_true = torch.tensor([1, 5, 3], dtype=torch.float32)

   2
1 x + 5 x + 3
X[0] = -4.94, Y[0] = 2.70
X[1] = -2.86, Y[1] = -3.12
X[2] = 1.57, Y[2] = 13.28
X[3] = 6.93, Y[3] = 85.71
X[4] = -1.05, Y[4] = -1.13
X[5] = 1.17, Y[5] = 10.22
X[6] = -1.86, Y[6] = -2.84
X[7] = -1.64, Y[7] = -2.51
X[8] = -5.17, Y[8] = 3.87
X[9] = -2.85, Y[9] = -3.13
X[10] = -1.37, Y[10] = -1.96
X[11] = -2.25, Y[11] = -3.19
X[12] = -1.60, Y[12] = -2.43
X[13] = 5.09, Y[13] = 54.43
X[14] = 5.14, Y[14] = 55.14
X[15] = 1.94, Y[15] = 16.44
X[16] = 3.41, Y[16] = 31.74
X[17] = 1.41, Y[17] = 12.06
X[18] = -5.17, Y[18] = 3.87
X[19] = 2.04, Y[19] = 17.39

Step-2: Preparing Inputs, Weights, and Optimiser

This code block sets up everything needed before the training loop begins.

XX = np.hstack([X*X, X, np.ones_like(X)]): Builds the design matrix of shape (N, 3). Each row for a sample x_i becomes [x_i^2 \;\; x_i \;\; 1]. This lets us express the polynomial w_1 x^2 + w_2 x + w_3 as a single matrix multiplication \hat{y} = \mathbf{X}\mathbf{w}.
w_pred = torch.randn(3, 1, requires_grad=True): Initialises the three unknown coefficients with random values. The flag requires_grad=True tells PyTorch to track every operation on w_pred so gradients can be computed later.
x = torch.tensor(XX, dtype=torch.float32) and y = torch.tensor(Y, dtype=torch.float32): Converts the NumPy arrays to PyTorch tensors so they can be used in autograd computations.
optimizer = torch.optim.NAdam([w_pred], lr=0.01): Creates an NAdam optimiser (a variant of Adam with Nesterov momentum) that will update w_pred using the gradients computed by autograd. The learning rate lr=0.01 controls the step size of each update.

# Prepare input as an array of shape (N,3)
# To express the polynomial w_1*x^2 + w_2*x + w_3 as a 
# single matrix multiplication
XX = np.hstack([X*X, X, np.ones_like(X)])

# Prepare tensors
# We initialize the coefficients w randomly and 
# set requires_grad=True to enable gradient tracking for optimization.
w_pred = torch.randn(3, 1, requires_grad=True)

# Convert the input and output samples to PyTorch tensors
x = torch.tensor(XX, dtype=torch.float32)
print("Input tensor x:")
print(x)
print("-"*79)

y = torch.tensor(Y, dtype=torch.float32)
print("Output tensor y:")
print(y)
print("-"*79)

# We will use the NAdam optimizer to update the 
# coefficients w during training.
optimizer = torch.optim.NAdam([w_pred], lr=0.01)
print("Coefficients w before training:")
print(w_pred)
print("-"*79)

Input tensor x:
tensor([[24.3900, -4.9386,  1.0000],
        [ 8.1567, -2.8560,  1.0000],
        [ 2.4528,  1.5662,  1.0000],
        [48.0512,  6.9319,  1.0000],
        [ 1.0933, -1.0456,  1.0000],
        [ 1.3708,  1.1708,  1.0000],
        [ 3.4626, -1.8608,  1.0000],
        [ 2.6892, -1.6399,  1.0000],
        [26.7200, -5.1691,  1.0000],
        [ 8.1422, -2.8535,  1.0000],
        [ 1.8668, -1.3663,  1.0000],
        [ 5.0798, -2.2538,  1.0000],
        [ 2.5442, -1.5950,  1.0000],
        [25.9542,  5.0945,  1.0000],
        [26.4319,  5.1412,  1.0000],
        [ 3.7525,  1.9371,  1.0000],
        [11.6621,  3.4150,  1.0000],
        [ 1.9977,  1.4134,  1.0000],
        [26.7135, -5.1685,  1.0000],
        [ 4.1740,  2.0430,  1.0000]])
-------------------------------------------------------------------------------
Output tensor y:
tensor([[ 2.6969],
        [-3.1233],
        [13.2836],
        [85.7107],
        [-1.1347],
        [10.2248],
        [-2.8414],
        [-2.5102],
        [ 3.8743],
        [-3.1251],
        [-1.9647],
        [-3.1894],
        [-2.4311],
        [54.4269],
        [55.1379],
        [16.4382],
        [31.7370],
        [12.0648],
        [ 3.8710],
        [17.3893]])
-------------------------------------------------------------------------------
Coefficients w before training:
tensor([[-0.1128],
        [-0.7510],
        [-1.3068]], requires_grad=True)
-------------------------------------------------------------------------------

Step-3: Run the Optimization for 5000 Iterations

The code below runs the training loop for 5000 iterations. Each iteration performs four actions:

optimizer.zero_grad(): Clears the gradients from the previous iteration. PyTorch accumulates gradients by default, so without this reset the gradients would keep growing and the optimiser would take incorrect steps.
y_pred = x @ w_pred: Computes the predicted \hat{y} values via matrix multiplication \mathbf{X}\mathbf{w}. This is the forward pass.
mse = torch.mean(torch.square(y - y_pred)): Calculates the Mean Squared Error loss between the true y and the predictions. This single number measures how far off the current weights are.
mse.backward(): Triggers backpropagation: PyTorch walks the computation graph in reverse and fills w_pred.grad with \frac{\partial\,\text{MSE}}{\partial\mathbf{w}}.
optimizer.step(): The NAdam optimiser uses those gradients to guide w_pred in the direction that reduces the loss.

After 5000 such updates, w_pred should converge close to the true coefficients [1, 5, 3], and the final lines print both sets of coefficients plus the remaining error.

# Run optimizer
for _ in range(5000):
    optimizer.zero_grad()
    y_pred = x @ w_pred
    mse = torch.mean(torch.square(y - y_pred))
    mse.backward()
    optimizer.step()

print("Coefficients w after training:")
print(w_pred)
print("-"*79)

print("True coefficients w:")
print(w_true)
print("-"*79)

print("Errors in coefficients:")
print(w_true - w_pred.detach().squeeze())

Coefficients w after training:
tensor([[1.0000],
        [5.0000],
        [3.0000]], requires_grad=True)
-------------------------------------------------------------------------------
True coefficients w:
tensor([1., 5., 3.])
-------------------------------------------------------------------------------
Errors in coefficients:
tensor([-1.5497e-06,  1.0490e-05,  2.5034e-05])

Step-4: Plotting the Results

The plot below shows the 20 training samples alongside the true polynomial curve and the curve recovered by gradient descent. A close match confirms that autograd successfully recovered the coefficients.

# Smooth x-values for plotting the curves
x_plot = np.linspace(X.min() - 1, X.max() + 1, 200).reshape(-1, 1)

# True curve
y_true_curve = polynomial(x_plot)

# Predicted curve using learned coefficients
w_np = w_pred.detach().numpy().flatten()
y_pred_curve = w_np[0] * x_plot**2 + w_np[1] * x_plot + w_np[2]

fig, ax = plt.subplots(figsize=(8, 5))

# Plot training samples
ax.scatter(
    X, Y, color="#2d6a4f", edgecolors="#1b4332",
    s=60, zorder=3,
    label="Training samples")

# Plot true curve with known coefficients
ax.plot(
    x_plot, y_true_curve, color="black", linewidth=2,
    label=f"True: $y = x^2 + 5x + 3$")

# Plot learned curve with recovered coefficients
ax.plot(
    x_plot, y_pred_curve, color="gold", linewidth=2, linestyle="--",
    label=f"Learned: $y = {w_np[0]:.2f}x^2 + {w_np[1]:.2f}x + {w_np[2]:.2f}$")

ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_title("Polynomial Regression via Automatic Differentiation")
ax.legend()
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()