Mean Squared Error Gradient

mediumFREECalculus & OptimizationCalculus, Machine LearningCV: Mathematical Foundations

Implement a function to compute the gradient of the Mean Squared Error (MSE) loss with respect to predictions, a crucial step in training models using optimization techniques. This task involves understanding the mathematical foundation of MSE and its derivative.

The Mean Squared Error is a measure of the average squared difference between predictions $\hat{y}$ and actual targets $y$ , given by the formula $\frac{1}{n}\sum_{i=1}^{n}(\hat{y}_i - y_i)^2$ . To minimize this loss, we need to calculate its gradient with respect to each prediction $\hat{y}_i$ .

Here are the steps to calculate the gradient:

Define the MSE loss function.
Apply the chain rule and the sum rule of calculus to differentiate the MSE loss with respect to each $\hat{y}_i$ .

\text{MSE} = \frac{1}{n}\sum_{i=1}^{n}(\hat{y}_i - y_i)^2

This technique is widely used in machine learning for training regression models.

Example

Input

mse_gradient([2.0, 4.0], [1.0, 3.0])

Output

[1.0, 1.0]

Reasoning

Inputs: predictions $\hat{y} = [2.0, 4.0]$ and targets $y = [1.0, 3.0]$ , two equal-length lists, so $n = 2$ . The loss here is $\text{MSE} = \frac{1}{2}\big[(2-1)^2 + (4-3)^2\big] = 1$ .
Math: only the $i$ -th term of the sum contains $\hat{y}_i$ , so every other term differentiates to $0$ . The chain rule on $(\hat{y}_i - y_i)^2$ brings down a factor of $2$ , and the $\frac{1}{n}$ in front stays, giving $\frac{\partial\,\text{MSE}}{\partial \hat{y}_i} = \frac{2}{n}(\hat{y}_i - y_i)$ .
Plugging in: both errors are $2 - 1 = 1$ and $4 - 3 = 1$ , so each component is $\frac{2}{2} \cdot 1 = 1$ .
Output: [1.0, 1.0], one partial derivative per prediction, rounded to 4 decimal places.
Notice: the sign follows the error. A prediction above its target gets a positive gradient (lowering it reduces the loss), one below gets a negative gradient, and a perfect prediction gets $0$ . Larger $n$ shrinks every component by the $\frac{2}{n}$ factor.

Constraints

Both predictions and targets have n elements where 1 ≤ n ≤ 1000
Return the gradient vector rounded to 4 decimal places

Learn this in

Mathematical FoundationsFoundations Linear RegressionML Study Plan

More from CV: Mathematical Foundations

Adam Optimizer Stephard Power Iteration for Dominant Eigenvectorhard SVD Image Compressionhard Batch Affine Transform with PyTorchmedium Covariance Matrix from Image Patcheshard

Concept primer

The gradient of MSE with respect to predictions comes from applying basic derivative rules (chain rule and power rule) to a common loss used in regression. MSE measures the average squared difference between predictions $\hat y$ and targets $y$ , and its gradient tells you how to change $\hat y$ (or the model parameters that produce $\hat y$ ) to reduce this error most quickly. Understanding this gradient is essential because it is what drives gradient descent and backpropagation in neural networks.

For each sample $i$ , the MSE includes a term $(\hat y_i - y_i)^2$ . The derivative of this term with respect to $\hat y_i$ is proportional to the error $(\hat y_i - y_i)$ itself, scaled by constants from the square and the averaging over $n$ . That is why the gradient looks like a scaled version of the prediction error and points in the direction that reduces the discrepancy between predictions and targets.

1. Background Knowledge

Mean Squared Error (MSE) MSE for predictions $\hat y \in \mathbb{R}^n$ and targets $y \in \mathbb{R}^n$ is

\text{MSE}(\hat y, y) = \frac{1}{n}\sum_{i=1}^n (\hat y_i - y_i)^2.

It is a scalar function of the vector $\hat y$ .

Partial derivatives and gradients The gradient of a scalar function $L(\hat y)$ with respect to the vector $\hat y$ is

\nabla_{\hat y} L = \begin{bmatrix} \frac{\partial L}{\partial \hat y_1} \\ \vdots \\ \frac{\partial L}{\partial \hat y_n} \end{bmatrix}.

Each component is found by treating all other $\hat y_j$ as constants and differentiating with respect to $\hat y_i$ .

Derivative rules you need
Power rule: $\frac{d}{dx}(x^2) = 2x$ .
Chain rule: for $f(g(x))$ , $\frac{d}{dx} f(g(x)) = f'(g(x)) g'(x)$ .
Linearity: derivative of a sum is the sum of derivatives; constants factor out:

\frac{d}{dx}\left(\frac{1}{n}\sum_i f_i(x)\right) = \frac{1}{n}\sum_i f_i'(x).

2. Algorithm / Approach Pattern

To get the gradient of a loss like MSE with respect to predictions:

Write the loss explicitly as a sum over individual sample terms.
Pick a single prediction component $\hat y_i$ and compute $\frac{$ \partial $}{$ \partial $\hat y_i}$ of the loss.
Use linearity: constants and terms not involving $\hat y_i$ disappear or factor out.
Differentiate the remaining term using the power rule and a simple chain rule.
Vectorize the result to write the gradient in compact form for all $i$ .

This is the standard pattern for deriving gradients of averaged per-sample losses.

3. Step-by-Step Strategy

Start from the definition

\text{MSE} = \frac{1}{n}\sum_{k=1}^{n}(\hat{y}_k - y_k)^2.

Focus on one component $\hat y_i$ Compute

\frac{\partial \text{MSE}}{\partial \hat y_i} = \frac{\partial}{\partial \hat y_i} \left( \frac{1}{n}\sum_{k=1}^n(\hat{y}_k - y_k)^2 \right).

Use linearity of derivative

For $k \neq i$ , $(\hat y_k - y_k)^2$ does not depend on $\hat y_i$ , so its derivative is 0.
Only the $k = i$ term remains:

\frac{\partial \text{MSE}}{\partial \hat y_i} = \frac{1}{n} \cdot \frac{\partial}{\partial \hat y_i} (\hat y_i - y_i)^2.

Apply power + chain rule

\frac{\partial}{\partial \hat y_i} (\hat y_i - y_i)^2 = 2(\hat y_i - y_i)\cdot \frac{\partial}{\partial \hat y_i} (\hat y_i - y_i) = 2(\hat y_i - y_i)\cdot 1 = 2(\hat y_i - y_i).

Combine

\frac{\partial \text{MSE}}{\partial \hat y_i} = \frac{1}{n} \cdot 2(\hat y_i - y_i) = \frac{2}{n}(\hat y_i - y_i).

Vector form (implementation-friendly) If $\hat y$ and $y$ are vectors:

\nabla_{\hat y} \text{MSE} = \frac{2}{n}(\hat y - y).

In code (NumPy-style):

def mse_gradient(y_pred, y_true):
 n = y_true.shape
 return (2.0 / n) * (y_pred - y_true)

4. Common Pitfalls

Missing the $1/n$ factor Forgetting the averaging leads to gradient $\;2(\hat y_i - y_i)\;$ instead of $\;$ \frac{2}{n} $(\hat y_i - y_i)$ .
Sign error The gradient w.r.t. predictions is proportional to $\hat y - y$ , not $y - \hat y$ . If you see an extra minus sign, re-check which variable you’re differentiating with respect to.
Confusing gradient w.r.t. predictions vs. parameters Here you only differentiate with respect to $\hat y$ . When using backprop, this gradient becomes the “incoming gradient” for the layer that produced $\hat y$ .
Shape mismatches in vectorized code Make sure $y$ and $\hat y$ have the same shape and that any division by $n$ is broadcast correctly (e.g., avoid integer division in some languages).

5. Time & Space Complexity

Assuming you already have $\hat y$ and $y$ :

Time complexity:
Computing the gradient requires one subtraction and one scalar multiplication per element.
Overall: $O(n)$ .
Space complexity:
You need space to store the gradient, typically same shape as $\hat y$ .
Overall: $O(n)$ additional space.

solution.py

Python 3.13

Test results