---
title: "Stochastic Gradient Descent"
subtitle: "Lecture 26"
author: "Dr. Colin Rundel"
footer: "Sta 323 - Spring 2026"
format:
  revealjs:
    theme: slides.scss
    transition: fade
    slide-number: true
    self-contained: true
execute:
  echo: true
engine: knitr
---

```{r setup}
#| message: false
#| warning: false
#| include: false

knitr::opts_chunk$set(
  fig.align = "center", fig.retina = 2, dpi = 150,
  out.width = "100%"
)

library(tidyverse)
library(torch)
library(luz)
ggplot2::theme_set(ggplot2::theme_bw())

options(width = 70)
```


# Gradient Descent <br/> for Regression


## A regression example

::: {.mxsmall}
```{r}
set.seed(1234)
n = 200; p = 20
X = matrix(rnorm(n * p), n, p)
true_beta = rep(0, p)
true_beta[c(3, 7, 12, 18)] = c(5.2, -3.1, 8.7, -1.5)
y = 3 + X %*% true_beta + rnorm(n)
```
:::

:::: {.columns .mxsmall}
::: {.column width='50%'}
```{r}
head(y, 20)
```
:::

::: {.column width='50%'}
```{r}
true_beta
```
:::
::::


## Minimalistic GD for linear regression

::: {.xsmall}
```{r}
#| code-line-numbers: "|6,7|12|13-15"
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

::: {.xsmall}
```{r}
(lm_fit = lm(y ~ X)) |> coef() |> round(2) |> unname()
deviance(lm_fit) |> round(2)
```
:::

. . .

::: {.xsmall}
```{r}
gd_lm = grad_desc_lm(X, y, rep(0, p + 1), step = 0.001, max_step = 20)
gd_lm$x |> tail(1) |> unlist() |> round(2)
gd_lm$loss |> tail(1) |> round(2)
```
:::


##

```{r}
#| echo: false
ggplot(
  tibble(iter = gd_lm$iter, loss = gd_lm$loss),
  aes(iter, loss)
) +
  geom_line() + geom_point(size = 0.8) +
  scale_y_log10() +
  labs(x = "Iteration", y = "Loss")
```

## 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$?

::: {.aside}
*Hint* - $\underset{m \times n}{\boldsymbol{A}} \cdot \underset{n \times k}{\boldsymbol{B}}$ has complexity $O(mnk)$.
:::

. . .

::: {.columns}
::: {.column width=20%}
:::
::: {.column width=60%}
* Calculating the loss function costs ${O}(nk)$

* Calculating the gradient costs ${O}(nk)$
:::
::: {.column width=20%}
:::
:::



# 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

::: {.xsmall}
```{r}
#| code-line-numbers: "|5|11-15|10,17|18"
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 {.scrollable}

::: {.xxsmall}
```{r}
lm_fit |> coef() |> round(2) |> unname()
lm_fit |> deviance() |> round(2)
```
:::

. . .

::: {.xxsmall}
```{r}
sgd_lm_rep = sgd_lm(
  X, y, rep(0, p + 1),
  step = 0.001, max_step = 20, replace = TRUE
)
sgd_lm_rep$x |> tail(1) |> unlist() |> round(2)
sgd_lm_rep$loss |> tail(1) |> round(2)
```
:::

. . .

::: {.xxsmall}
```{r}
sgd_lm_worep = sgd_lm(
   X, y, rep(0, p + 1),
   step = 0.001, max_step = 20, replace = FALSE
)
sgd_lm_worep$x |> tail(1) |> unlist() |> round(2)
sgd_lm_worep$loss |> tail(1) |> round(2)
```
:::


## Learning by step

```{r}
#| echo: false
bind_rows(
  tibble(iter = gd_lm$iter, loss = gd_lm$loss, method = "GD"),
  tibble(iter = sgd_lm_rep$iter, loss = sgd_lm_rep$loss, method = "SGD - replacement"),
  tibble(iter = sgd_lm_worep$iter, loss = sgd_lm_worep$loss, method = "SGD - w/o replacement")
) |>
  ggplot(aes(iter, loss, color = method)) +
  geom_line(alpha = 0.7) +
  scale_y_log10() +
  labs(x = "Steps", y = "Loss", color = "")
```

## Learning by Epochs

Generally, rather than thinking in iterations we use epochs instead - an epoch is one complete pass through the data.

```{r}
#| echo: false
bind_rows(
  tibble(iter = gd_lm$iter, loss = gd_lm$loss, method = "GD"),
  tibble(iter = sgd_lm_rep$iter / n, loss = sgd_lm_rep$loss, method = "SGD - replacement"),
  tibble(iter = sgd_lm_worep$iter / n, loss = sgd_lm_worep$loss, method = "SGD - w/o replacement")
) |>
  ggplot(aes(iter, loss, color = method)) +
  geom_line(alpha = 0.7) +
  scale_y_log10() +
  labs(x = "Epoch", y = "Loss", color = "")
```

## Run times

::: {.xsmall}
```{r}
gd_time = system.time(
  grad_desc_lm(X, y, rep(0, p + 1), step = 0.001, max_step = 20)
)
sgd_rep_time = system.time(
  sgd_lm(X, y, rep(0, p + 1), step = 0.001, max_step = 20, replace = TRUE)
)
sgd_worep_time = system.time(
  sgd_lm(X, y, rep(0, p + 1), step = 0.001, max_step = 20, replace = FALSE)
)
```
:::

```{r}
#| echo: false
tibble(
  Method = c("GD", "SGD (replacement)", "SGD (w/o replacement)"),
  `Run time` = c(
    sprintf("%.3fs", gd_time["elapsed"]),
    sprintf("%.3fs", sgd_rep_time["elapsed"]),
    sprintf("%.3fs", sgd_worep_time["elapsed"])
  ),
  `Time / Epoch` = c(
    sprintf("%.3fs", gd_time["elapsed"] / 20),
    sprintf("%.3fs", sgd_rep_time["elapsed"] / 20),
    sprintf("%.3fs", sgd_worep_time["elapsed"] / 20)
  )
) |>
  knitr::kable()
```


## A bigger example

Remember that our costs scale as $O(nk)$, so let's see what happens when we increase the data size.

::: {.xsmall}
```{r}
#| code-line-numbers: "|2"
set.seed(1234)
n = 10000; p = 20
X = matrix(rnorm(n * p), n, p)
true_beta = rep(0, p)
true_beta[c(3, 7, 12, 18)] = c(5.2, -3.1, 8.7, -1.5)
y = 3 + X %*% true_beta + rnorm(n)
```
:::

. . .

::: {.xsmall}
```{r}
lm_fit = lm(y ~ X)
lm_fit |> coef() |> round(2) |> unname()
```
:::

## Fitting

::: {.xsmall}
```{r}
lm_fit |> coef() |> round(2) |> unname()
lm_fit |> deviance() |> round(2)
```

```{r}
gd_lm = grad_desc_lm(
  X, y, rep(0, p + 1),
  step = 0.00005, max_step = 3
)
gd_lm$x |> tail(1) |> unlist() |> round(2)
gd_lm$loss |> tail(1) |> round(2)
```
:::

##

::: {.xsmall}
```{r}
sgd_lm_rep = sgd_lm(
  X, y, rep(0, p + 1),
  step = 0.001, max_step = 3, replace = TRUE
)
sgd_lm_rep$x |> tail(1) |> unlist() |> round(2)
sgd_lm_rep$loss |> tail(1) |> round(2)
```
```{r}
sgd_lm_worep = sgd_lm(
  X, y, rep(0, p + 1),
  step = 0.001, max_step = 3, replace = FALSE
)
sgd_lm_worep$x |> tail(1) |> unlist() |> round(2)
sgd_lm_worep$loss |> tail(1) |> round(2)
```
:::


## Results

::: {.panel-tabset}

### Full

```{r}
#| echo: false
bind_rows(
  tibble(iter = gd_lm$iter, loss = gd_lm$loss, method = "GD"),
  tibble(iter = sgd_lm_rep$iter / n, loss = sgd_lm_rep$loss, method = "SGD - replacement"),
  tibble(iter = sgd_lm_worep$iter / n, loss = sgd_lm_worep$loss, method = "SGD - w/o replacement")
) |>
  ggplot(aes(iter, loss, color = method)) +
  geom_line(alpha = 0.7) +
  scale_y_log10() +
  labs(x = "Epoch", y = "Loss", color = "")
```

### Zoom

```{r}
#| echo: false
bind_rows(
  tibble(iter = gd_lm$iter, loss = gd_lm$loss, method = "GD"),
  tibble(iter = sgd_lm_rep$iter / n, loss = sgd_lm_rep$loss, method = "SGD - replacement"),
  tibble(iter = sgd_lm_worep$iter / n, loss = sgd_lm_worep$loss, method = "SGD - w/o replacement")
) |>
  ggplot(aes(iter, loss, color = method)) +
  geom_line(alpha = 0.7) +
  scale_y_log10() +
  coord_cartesian(xlim = c(0, 0.3)) +
  labs(x = "Epoch", y = "Loss", color = "")
```

### Timings

::: {.xsmall}
```{r}
gd_time = system.time(
  grad_desc_lm(X, y, rep(0, p + 1), step = 0.00005, max_step = 3)
)
sgd_rep_time = system.time(
  sgd_lm(X, y, rep(0, p + 1), step = 0.001, max_step = 3, replace = TRUE)
)
sgd_worep_time = system.time(
  sgd_lm(X, y, rep(0, p + 1), step = 0.001, max_step = 3, replace = FALSE)
)
```
:::

```{r}
#| echo: false
tibble(
  Method = c("GD", "SGD (replacement)", "SGD (w/o replacement)"),
  `Run time` = c(
    sprintf("%.3fs", gd_time["elapsed"]),
    sprintf("%.3fs", sgd_rep_time["elapsed"]),
    sprintf("%.3fs", sgd_worep_time["elapsed"])
  ),
  `Time / Epoch` = c(
    sprintf("%.3fs", gd_time["elapsed"] / 3),
    sprintf("%.3fs", sgd_rep_time["elapsed"] / 3),
    sprintf("%.3fs", sgd_worep_time["elapsed"] / 3)
  ),
  `Time to convergence` = c(
    sprintf("%.3fs", gd_time["elapsed"] / 3 * 3),
    sprintf("%.3fs", sgd_rep_time["elapsed"] / 3 * 0.2),
    sprintf("%.3fs", sgd_worep_time["elapsed"] / 3 * 0.2)
  )
) |>
  knitr::kable()
```

:::


# 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

::: {.xsmall}
```{r}
#| code-line-numbers: "|5|11-15|17|19"
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 {.scrollable}

::: {.xxsmall}
```{r}
lm_fit |> coef() |> round(2) |> unname()
lm_fit |> deviance() |> round(2)
```

```{r}
sizes = c(10, 50, 100)
mbgd = list()
for (size in sizes) {
  mbgd[[as.character(size)]] = mb_grad_desc_lm(
    X, y, rep(0, p + 1), batch_size = size,
    step = 0.001, max_step = 3, replace = FALSE
  )
}
```

```{r}
mbgd[["10"]]$x |> tail(1) |> unlist() |> round(2)
mbgd[["10"]]$loss |> tail(1) |> round(2)
```
```{r}
mbgd[["50"]]$x |> tail(1) |> unlist() |> round(2)
mbgd[["50"]]$loss |> tail(1) |> round(2)
```
```{r}
mbgd[["100"]]$x |> tail(1) |> unlist() |> round(2)
mbgd[["100"]]$loss |> tail(1) |> round(2)
```
:::


## Results

::: {.panel-tabset}

### Full

```{r}
#| echo: false
df_all = bind_rows(
  tibble(iter = gd_lm$iter, loss = gd_lm$loss, method = "GD"),
  tibble(iter = sgd_lm_rep$iter / n, loss = sgd_lm_rep$loss, method = "SGD - replacement"),
  tibble(iter = sgd_lm_worep$iter / n, loss = sgd_lm_worep$loss, method = "SGD - w/o replacement"),
  purrr::map_dfr(sizes, \(size) {
    res = mbgd[[as.character(size)]]
    tibble(iter = res$iter / (n / size), loss = res$loss, method = paste0("MBGD (", size, ")"))
  })
)

ggplot(df_all, aes(iter, loss, color = method)) +
  geom_line(alpha = 0.7) +
  scale_y_log10() +
  labs(x = "Epoch", y = "Loss", color = "")
```

### Zoom

```{r}
#| echo: false
ggplot(df_all, aes(iter, loss, color = method)) +
  geom_line(alpha = 0.7) +
  scale_y_log10() +
  coord_cartesian(xlim = c(0, 0.25)) +
  labs(x = "Epoch", y = "Loss", color = "")
```

### Timings

::: {.xsmall}
```{r}
#| include: false
mbgd_time = list()
for (size in sizes) {
  mbgd_time[[as.character(size)]] = system.time(
    mb_grad_desc_lm(
      X, y, rep(0, p + 1), batch_size = size,
      step = 0.001, max_step = 3, replace = FALSE
    )
  )
}
```
:::

```{r}
#| echo: false
tibble(
  Method = c("GD", "SGD (replacement)", "SGD (w/o replacement)",
             paste0("MBGD (", sizes, ")")),
  `Run time` = c(
    sprintf("%.3fs", gd_time["elapsed"]),
    sprintf("%.3fs", sgd_rep_time["elapsed"]),
    sprintf("%.3fs", sgd_worep_time["elapsed"]),
    sapply(sizes, \(s) sprintf("%.3fs", mbgd_time[[as.character(s)]]["elapsed"]))
  ),
  `Time / Epoch` = c(
    sprintf("%.3fs", gd_time["elapsed"] / 3),
    sprintf("%.3fs", sgd_rep_time["elapsed"] / 3),
    sprintf("%.3fs", sgd_worep_time["elapsed"] / 3),
    sapply(sizes, \(s) sprintf("%.3fs", mbgd_time[[as.character(s)]]["elapsed"] / 3))
  ),
  `Time to convergence` = c(
    sprintf("%.3fs", gd_time["elapsed"] / 3 * 2),
    sprintf("%.3fs", sgd_rep_time["elapsed"] / 3 * 0.2),
    sprintf("%.3fs", sgd_worep_time["elapsed"] / 3 * 0.2),
    sapply(sizes, \(s) sprintf("%.3fs", mbgd_time[[as.character(s)]]["elapsed"] / 3 * 0.2))
  )
) |>
  knitr::kable()
```

:::


## 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`](https://torch.mlverse.org) package is a native R port of [PyTorch](https://pytorch.org), 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`](https://mlverse.github.io/luz/) (high-level training), [`torchvision`](https://torchvision.mlverse.org), [`torchaudio`](https://curso-r.github.io/torchaudio/)

::: {.aside}
The torch + luz combination is roughly the R analogue of PyTorch + PyTorch Lightning.
:::


## Tensors

A `torch_tensor` is the basic data type - a multi-dimensional array with a specific dtype and device.

::: {.columns .xsmall}
::: {.column}
```{r}
torch_zeros(3)
torch_ones(3, 2)
```
:::
::: {.column}
```{r}
torch_manual_seed(1234)
torch_rand(2, 2, 2)
```

:::
:::


## From R to torch

Converting R objects to tensors uses `torch_tensor()`:

::: {.mxsmall}
```{r}
torch_tensor(diag(3))
torch_tensor(matrix(1:6, 2, 3))
```
:::

. . .

Note that R integer vectors and matrices become `Long` tensors. Most `nn` layers expect `Float` (float32) so you usually want to convert explicitly:

::: {.mxsmall}
```{r}
torch_tensor(matrix(1:6, 2, 3), dtype = torch_float())
```
:::


## From torch to R

Going the other direction uses `as_array()` (or `as.matrix()` / `as.numeric()` for 2d / 1d tensors):

::: {.xsmall}
```{r}
t = torch_tensor(matrix(1:6, 2, 3), dtype = torch_float())
as_array(t)
as.matrix(t)
```
:::



::: {.aside}

If a tensor is tracking gradients you must `$detach()` it first:

::: {.small}
```{r}
t = torch_tensor(c(1, 2, 3), requires_grad = TRUE)
as_array(t$detach())
```
:::

:::


## Reshaping tensors

Tensors can be reshaped with `$view()` or `$reshape()`, transposed with `$t()`:

::: {.mxsmall}
```{r}
a = torch_arange(1, 12)
a$view(c(3, 4))
```
:::

. . .

::: {.mxsmall}
```{r}
a$reshape(c(2, 6))
a$view(c(3, 4))$t()
```
:::

::: {.aside}
`$view()` requires the tensor to be contiguous in memory; `$reshape()` will copy if needed.
:::


## Adding and removing dimensions

`$squeeze()` drops dimensions of size 1, `$unsqueeze(dim)` inserts a new dimension of size 1:

::: {.xsmall}
```{r}
b = torch_zeros(2, 1, 3)
b$shape
b$squeeze()$shape
```
:::

. . .

::: {.xsmall}
```{r}
c = torch_zeros(2, 3)
c$unsqueeze(1)$shape
c$unsqueeze(2)$shape
```
:::

::: {.aside}
Adding and removing dimensions is typically done to conform tensor shapes so that they can be combined via [broadcasting](https://pytorch.org/docs/stable/notes/broadcasting.html).
:::



## Autograd

::: {.medium}
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:
:::

::: {.xsmall}
```{r}
a = torch_tensor(2, requires_grad = TRUE)
f = a^3 + 2*a
f
```
:::

. . .

::: {.xsmall}
```{r}
f$backward()
a$grad
```
:::

. . .

::: {.medium}
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:

::: {.xsmall}
```{r}
a = torch_tensor(2, requires_grad = TRUE)
f = a^3 + 2*a
f$grad_fn
```
:::

. . .

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`.

::: {.aside}
The graph is rebuilt on every forward pass - this is what makes it "dynamic" and is one of the things people like about PyTorch (and torch).
:::


# Linear regression <br/> 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:

::: {.xsmall}
```{r}
Xt = torch_tensor(cbind(1, X), dtype = torch_float())
yt = torch_tensor(as.numeric(y), dtype = torch_float())

beta = torch_zeros(p + 1, dtype = torch_float(), requires_grad = TRUE)
```
:::

. . .

::: {.xsmall}
```{r}
Xt$shape
yt$shape
beta$shape
```
:::



## One step by hand

::: {.xsmall}
```{r}
loss = ((Xt %*% beta - yt)^2)$sum()
loss
```
:::

. . .

::: {.xsmall}
```{r}
loss$backward()
beta$grad
```
:::

. . .

Notice that we never wrote down the gradient ourselves - autograd derived it from the forward pass.


## A torch training loop

::: {.mxsmall}
```{r}
#| code-line-numbers: "|3|4|5-8"
beta = torch_zeros(p + 1, dtype = torch_float(), requires_grad = TRUE)
for (i in 1:200) {
  loss = ((Xt %*% beta - yt)^2)$sum()
  loss$backward()
  with_no_grad({
    beta$sub_(1e-5 * beta$grad)
    beta$grad$zero_()
  })
}
```
:::

. . .

::: {.mxsmall}
```{r}
loss$item()
deviance(lm_fit)
```
:::

. . .

::: {.mxsmall}
```{r}
beta$detach() |> as_array() |> round(2)
lm_fit |> coef() |> round(2) |> unname()
```
:::

::: {.aside}
`with_no_grad({...})` tells autograd "don't track these operations" - we need it because the parameter update is *not* part of the loss, and we don't want it on the computational graph.
:::


# `nn_module`

## What is `nn_module`?

`nn_module()` creates a model class with two key methods:

* `initialize()` - builds the parameters and submodules

* `forward()` - 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`.

::: {.aside}
This is a direct analogue of `class Model(torch.nn.Module)` in PyTorch.
:::


## A first linear regression model

::: {.xsmall}
```{r}
#| code-line-numbers: "|2-4|5-7"
linear_model = nn_module(
  initialize = function(k) {
    self$beta = nn_parameter(torch_zeros(k))
  },
  forward = function(X) {
    X %*% self$beta
  }
)
```
:::

. . .

::: {.xsmall}
```{r}
(m = linear_model(p + 1))
```
:::

. . .

::: {.xsmall}
```{r}
m$parameters
```
:::


## 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:

::: {.xsmall}
```{r}
#| code-line-numbers: "|2|5|6|7|8"
m = linear_model(p + 1)
opt = optim_sgd(m$parameters, lr = 1e-5)

for (i in 1:200) {
  opt$zero_grad()
  loss = ((m(Xt) - yt)^2)$sum()
  loss$backward()
  opt$step()
}
```
:::

. . .

::: {.xsmall}
```{r}
loss$item()
round(as_array(m$beta$detach()), 2)
```
:::


## 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$:

::: {.xsmall}
```{r}
#| code-line-numbers: "|3|6"
linear_model = nn_module(
  initialize = function(in_features) {
    self$linear = nn_linear(in_features, 1)
  },
  forward = function(X) {
    self$linear(X)
  }
)
```
:::

. . .

Because `nn_linear` already includes a bias term, we drop the column of 1s and feed in `X` directly. We do need to reshape `y` to `(n, 1)` so it matches the shape of the model's output:

::: {.xsmall}
```{r}
Xt = torch_tensor(X, dtype = torch_float())
yt = torch_tensor(matrix(y, ncol = 1), dtype = torch_float())
```
:::

## Linear parameters

Constructing the model registers a `weight` matrix (shape `(out_features, in_features)`) and a `bias` vector (shape `(out_features)`) as parameters:

::: {.xsmall}
```{r}
(m = linear_model(p))
```
:::

. . .

::: {.xsmall}
```{r}
m$parameters
```
:::


## Fitting

::: {.xsmall}
```{r}
#| code-line-numbers: "|2|5"
m = linear_model(p)
opt = optim_sgd(m$parameters, lr = 0.05)

losses = numeric(100)
for (i in 1:100) {
  opt$zero_grad()
  loss = nnf_mse_loss(m(Xt), yt)
  loss$backward()
  opt$step()
  losses[i] = loss$item()
}
```
:::

. . .

::: {.xsmall}
```{r}
c(torch = loss$item(), lm = deviance(lm_fit) / n)
c(torch_intercept = as.numeric(m$linear$bias), lm_intercept = unname(coef(lm_fit)[1]))
```
:::

::: {.aside}
`nnf_mse_loss()` defaults to mean reduction (versus the manual sum we used before), which is why the learning rate is larger here.
:::


## Loss curve

```{r}
#| echo: false
ggplot(tibble(iter = seq_along(losses), loss = losses), aes(iter, loss)) +
  geom_line() +
  scale_y_log10() +
  labs(x = "Epoch", y = "Loss (MSE)")
```