Stochastic Gradient Descent
Lecture 26
Gradient Descent
for Regression
A regression example
[,1]
[1,] 4.54630945
[2,] -0.27563573
[3,] -8.92201672
[4,] 3.81019348
[5,] -8.01772803
[6,] -5.65360256
[7,] -9.45140764
[8,] 4.92135279
[9,] -4.13507430
[10,] 6.77822004
[11,] -1.16028786
[12,] 20.83784033
[13,] 4.07084356
[14,] 9.58649951
[15,] 1.41834659
[16,] 13.08886089
[17,] 8.68524259
[18,] -1.72856652
[19,] 0.97934767
[20,] 0.03342649Minimalistic GD for linear regression
grad_desc_lm = function(X, y, beta, step, max_step = 50) {
X = cbind(1, X)
n = nrow(X)
k = ncol(X)
f = \(beta) sum((y - X %*% beta)^2)
grad = \(beta) 2 * t(X) %*% (X %*% beta - y)
res = list(x = list(beta), loss = f(beta), iter = 0)
for (i in seq_len(max_step)) {
beta = beta - grad(beta) * step
res$x = c(res$x, list(beta))
res$loss = c(res$loss, f(beta))
res$iter = c(res$iter, i)
}
res
}Linear regression
A quick analysis
Let’s take a quick look at the linear regression loss function and gradient descent and think a bit about its cost.
We’ve defined the loss function and its gradient as follows:
\[ \begin{aligned} f(\underset{k\times 1}{\boldsymbol{\beta}}) &= (\underset{n \times 1}{y} - \underset{n\times k}{\boldsymbol{X}} \, \underset{k \times 1}{\boldsymbol{\beta}})^T (\underset{n \times 1}{y} - \underset{n\times k}{\boldsymbol{X}} \, \underset{k \times 1}{\boldsymbol{\beta}}) \\ \\ \nabla f(\underset{k\times 1}{\boldsymbol{\beta}}) &= \underset{k\times n}{2 \boldsymbol{X}^T}(\underset{n \times k}{\boldsymbol{X}}\,\underset{k\times 1}{\boldsymbol{\beta}} - \underset{n \times 1}{\boldsymbol{y}}) \end{aligned} \]
What are the costs of calculating the loss function and gradient respectively in terms of \(n\) and \(k\)?
Calculating the loss function costs \({O}(nk)\)
Calculating the gradient costs \({O}(nk)\)
Stochastic Gradient Descent
Stochastic Gradient Descent
This is a variant of gradient descent where rather than using all \(n\) data points we randomly sample one at a time and use that single point to make our gradient step.
Sampling of observations can be done with or without replacement
Will take more steps to converge but each step is now cheaper to compute
SGD has slower asymptotic convergence than GD, but is often faster in practice in terms of runtime
Generally requires the learning rate to shrink as a function of iteration to guarantee convergence
SGD - Linear Regression
sgd_lm = function(X, y, beta, step, max_step = 50, seed = 1234, replace = TRUE) {
X = cbind(1, X)
n = nrow(X); k = ncol(X)
f = \(beta) sum((y - X %*% beta)^2)
grad_i = \(beta, i) 2 * X[i, ] * (X[i, ] %*% beta - y[i])
res = list(x = list(beta), loss = f(beta), iter = 0)
set.seed(seed)
for (ep in seq_len(max_step)) {
if (replace) {
js = sample.int(n, n, replace = TRUE)
} else {
js = sample.int(n)
}
for (j in js) {
beta = beta - grad_i(beta, j) * step
res$x = c(res$x, list(beta))
res$loss = c(res$loss, f(beta))
res$iter = c(res$iter, tail(res$iter, 1) + 1)
}
}
res
}Fitting
[1] 3.11 -0.07 0.00 5.38 0.06 0.07 0.05 -3.04 -0.06 0.09 0.07
[12] 0.13 8.65 0.01 0.12 0.01 0.14 -0.06 -1.45 0.03 -0.08[1] 173.31 [1] 3.11 -0.08 0.00 5.34 0.04 0.07 0.02 -3.06 -0.04 0.10 0.04
[12] 0.11 8.65 -0.01 0.08 -0.01 0.10 -0.01 -1.43 0.07 -0.07[1] 170.65Learning by step
Learning by Epochs
Generally, rather than thinking in iterations we use epochs instead - an epoch is one complete pass through the data.
Run times
| Method | Run time | Time / Epoch |
|---|---|---|
| GD | 0.000s | 0.000s |
| SGD (replacement) | 0.357s | 0.018s |
| SGD (w/o replacement) | 0.358s | 0.018s |
A bigger example
Remember that our costs scale as \(O(nk)\), so let’s see what happens when we increase the data size.
Fitting
[1] 3.03 0.01 0.00 5.22 -0.01 0.00 0.00 -3.10 -0.01 0.00 -0.01
[12] -0.01 8.69 -0.02 -0.01 0.02 0.02 0.01 -1.48 0.00 -0.01[1] 9859.22 [1] 3.00 0.00 0.01 5.24 -0.01 0.00 -0.03 -3.14 -0.03 -0.02 -0.01
[12] -0.03 8.69 -0.01 0.01 0.06 -0.01 -0.04 -1.51 0.01 -0.02[1] 9980.55 [1] 2.97 -0.04 0.02 5.15 0.03 -0.02 0.04 -3.09 0.03 -0.03 0.00
[12] -0.05 8.72 -0.05 -0.01 -0.01 0.01 0.00 -1.47 0.03 0.04[1] 10085.99Results
| Method | Run time | Time / Epoch | Time to convergence |
|---|---|---|---|
| GD | 0.002s | 0.001s | 0.002s |
| SGD (replacement) | 9.395s | 3.132s | 0.626s |
| SGD (w/o replacement) | 9.659s | 3.220s | 0.644s |
Mini-batch Gradient Descent
Mini-batch gradient descent
This is a variant of stochastic gradient descent where a subset of \(m\) data points is selected for each gradient update.
The idea is to find a balance between the cost of increasing the data size vs the speed-up of vectorized calculations.
More updates per epoch than GD, but less than SGD
Mini batch composition can be constructed by sampling data with or without replacement
MBGD - Linear Regression
mb_grad_desc_lm = function(X, y, beta, step, batch_size = 10, max_step = 50, seed = 1234, replace = TRUE) {
X = cbind(1, X)
n = nrow(X); k = ncol(X)
f = \(beta) sum((y - X %*% beta)^2)
grad_b = \(beta, i) 2 * t(X[i,]) %*% (X[i,] %*% beta - y[i])
res = list(x = list(beta), loss = f(beta), iter = 0)
set.seed(seed)
for (ep in seq_len(max_step)) {
if (replace) {
js = sample.int(n, n, replace = TRUE)
} else {
js = sample.int(n)
}
batches = split(js, ceiling(seq_along(js) / batch_size))
for (batch in batches) {
beta = beta - grad_b(beta, batch) * step
res$x = c(res$x, list(beta))
res$loss = c(res$loss, f(beta))
res$iter = c(res$iter, tail(res$iter, 1) + 1)
}
}
res
}Fitting
[1] 3.03 0.01 0.00 5.22 -0.01 0.00 0.00 -3.10 -0.01 0.00 -0.01
[12] -0.01 8.69 -0.02 -0.01 0.02 0.02 0.01 -1.48 0.00 -0.01[1] 9859.22 [1] 2.97 -0.04 0.02 5.15 0.03 -0.02 0.04 -3.09 0.03 -0.03 0.00
[12] -0.05 8.72 -0.05 -0.01 -0.01 0.01 0.00 -1.47 0.03 0.04[1] 10088.01 [1] 2.97 -0.04 0.02 5.15 0.03 -0.02 0.04 -3.09 0.03 -0.03 0.00
[12] -0.05 8.72 -0.05 -0.01 -0.01 0.01 0.00 -1.47 0.03 0.05[1] 10105.2Results
A bit of theory
We’ve talked a bit about the computational side of things, but why do these approaches work at all?
In statistics and machine learning many of our problems have a form that looks like,
\[ \underset{\theta}{\text{arg min}} \; \ell(\boldsymbol{X}, \theta) = \underset{\theta}{\text{arg min}} \; \frac{1}{n} \sum_{i=1}^n \ell(\boldsymbol{X}_i, \theta) \]
which means that the gradient of the loss function is given by
\[ \nabla \ell(\boldsymbol{X}, \theta) = \frac{1}{n} \sum_{i=1}^n \nabla \ell(\boldsymbol{X}_i, \theta) \]
\[ \nabla \ell(\boldsymbol{X}, \theta) \approx \frac{1}{|B|} \sum_{i \in B} \nabla \ell(\boldsymbol{X}_i, \theta) \]
SGD estimator
Because we are sampling \(B\) randomly, then our SGD and mini batch GD approximations are unbiased estimates of the full gradient,
\[ E\left[ \frac{1}{|B|} \sum_{i \in B} \nabla \ell(\boldsymbol{X}_i, \theta) \right] = \frac{1}{n} \sum_{i=1}^n \nabla \ell(\boldsymbol{X}_i, \theta) = \nabla \ell(\boldsymbol{X}, \theta) \]
Each update can be viewed as a noisy gradient descent step (gradient + zero mean noise).
- The difference between mini batch and stochastic gradient descent is that by increasing the computation cost per step we are reducing the noise variance for that step
Limitations
As mentioned previously we need to be a bit careful with learning rates and convergence for both of these methods. So far, our approach has been naive and runs for a fixed number of epochs.
If we want to use a convergence criterion we need to keep the following in mind:
Let \(\theta^*\) be a global / local minimizer of our loss function \(\ell(\boldsymbol{X},\theta)\), then by definition \(\nabla \ell(\boldsymbol{X},\theta^*) = 0\)
The issue is that our gradient approximation, \[ \frac{1}{|B|} \sum_{i \in B} \nabla \ell(\boldsymbol{X}_i, \theta) \ne 0 \] as \(B\) is a subset of the data, therefore our algorithm is likely to never converge.
Solution
The practical solution to this is to implement a learning rate schedule which generally shrinks the learning rate / step size over time to ensure convergence.
The choice of the exact learning schedule is problem specific, and is usually about finding the right balance.
Some common examples:
Piecewise constant - \(\eta_t = \eta_i \text{ if } t_i \leq t \leq t_{i+1}\)
Exponential decay - \(\eta_t = \eta_0 e^{-\lambda t}\)
Polynomial decay - \(\eta_t = \eta_0 (\beta t+1)^{-\alpha}\)
There are many more approaches including more exotic techniques that allow the learning rate to increase and decrease to help the optimizer better explore the objective function and in some cases escape local optima.
torch
What is torch?
The torch package is a native R port of PyTorch, built directly on top of libtorch (the C++ backend that powers PyTorch). There is no Python dependency.
It provides:
Multi-dimensional arrays (tensors) with GPU support
Automatic differentiation (autograd)
Neural network building blocks (
nn_module,nn_linear,nn_relu, …)A growing ecosystem -
luz(high-level training),torchvision,torchaudio
Tensors
A torch_tensor is the basic data type - a multi-dimensional array with a specific dtype and device.
From R to torch
Converting R objects to tensors uses torch_tensor():
From torch to R
Going the other direction uses as_array() (or as.matrix() / as.numeric() for 2d / 1d tensors):
Reshaping tensors
Tensors can be reshaped with $view() or $reshape(), transposed with $t():
Adding and removing dimensions
$squeeze() drops dimensions of size 1, $unsqueeze(dim) inserts a new dimension of size 1:
Autograd
This is where torch starts to do real work for us - any tensor created with requires_grad = TRUE will have its operations recorded so that gradients can be computed automatically:
This is the same gradient we would get by hand from \(\frac{df}{da} = 3a^2 + 2 = 14\) at \(a = 2\). Importantly, we did not have to write the gradient ourselves - we wrote the grad function by hand earlier in grad_desc_lm(), but autograd will do that for us.
Computational graph
Behind the scenes, every operation on a tensor with requires_grad = TRUE adds a node to a dynamic computational graph. The $grad_fn attribute points at the operation that produced the tensor:
When we call f$backward(), torch walks this graph from f back to a, applies the chain rule, and accumulates the result into a$grad.
Linear regression
with torch
Tensors from our regression data
We will reuse the regression problem from earlier in this lecture (X is \(n \times p\) with \(n = 10000\), \(p = 20\), y is the response). We need to convert these to tensors and prepend a column of 1s for the intercept:
One step by hand
Notice that we never wrote down the gradient ourselves - autograd derived it from the forward pass.
A torch training loop
[1] 3.03 0.01 0.00 5.22 -0.01 0.00 0.00 -3.10 -0.01 0.00 -0.01
[12] -0.01 8.69 -0.02 -0.01 0.02 0.02 0.01 -1.48 0.00 -0.01 [1] 3.03 0.01 0.00 5.22 -0.01 0.00 0.00 -3.10 -0.01 0.00 -0.01
[12] -0.01 8.69 -0.02 -0.01 0.02 0.02 0.01 -1.48 0.00 -0.01nn_module
What is nn_module?
nn_module() creates a model class with two key methods:
initialize()- builds the parameters and submodulesforward()- defines how an input is transformed into an output
Any tensor wrapped in nn_parameter() (or any nn_* submodule like nn_linear) is automatically registered as a parameter of the model and will show up in model$parameters.
A first linear regression model
Training loop with an optimizer
Rather than updating beta by hand, we can hand the model’s parameters to a torch optimizer and let it do the work:
Using nn_linear
In practice we rarely build a nn_parameter by hand - torch ships with layer modules that handle parameter creation, weight initialization, and the bias term for us. nn_linear(in_features, out_features) implements \(y = XW^T + b\):
Linear parameters
Constructing the model registers a weight matrix (shape (out_features, in_features)) and a bias vector (shape (out_features)) as parameters:
$linear.weight
torch_tensor
Columns 1 to 10 0.0777 -0.0754 0.1269 0.0282 0.1229 0.1435 -0.0987 0.0813 -0.0967 0.0701
Columns 11 to 20-0.1168 0.1034 0.0453 -0.0875 -0.1097 0.0579 0.2086 0.1073 -0.0216 -0.0109
[ CPUFloatType{1,20} ][ requires_grad = TRUE ]
$linear.bias
torch_tensor
0.1271
[ CPUFloatType{1} ][ requires_grad = TRUE ]Fitting
Loss curve
Sta 323 - Spring 2026