Mini-Batch SGD

Background Knowledge

Stochastic Gradient Descent (SGD) is a fundamental algorithm in Machine Learning used for training linear and non-linear models. It's an optimization technique that minimizes the loss function by iteratively adjusting the model parameters in the direction of the negative gradient of the loss. In the context of linear regression, the model is defined as $y = w \cdot x + b$, where $w$ is the weight, $b$ is the bias, $x$ is the input, and $y$ is the predicted output.

The Mean Squared Error (MSE) is a common loss function used for regression problems, which measures the average squared difference between predicted and actual values. The goal of SGD is to find the optimal values of $w$ and $b$ that minimize the MSE loss. The mini-batch variant of SGD splits the training data into smaller batches, computing the gradient of the loss for each batch, and updating the parameters accordingly. This approach helps to reduce the variance of the gradient estimates and improve the stability of the optimization process.

In the context of this problem, we need to implement one epoch of mini-batch SGD for a simple linear model. This involves splitting the training data into mini-batches, computing the mean gradient of the MSE loss for each batch, and updating the parameters $w$ and $b$ using the gradients and a learning rate $\alpha$. The process is repeated for all batches, and the final values of $w$ and $b$ are returned after one full epoch.