---
title: "Optimization - SGD"
subtitle: "Lecture 23"
author: "Dr. Colin Rundel"
footer: "Sta 663 - Spring 2026"
format:
  revealjs:
    theme: slides.scss
    transition: fade
    slide-number: true
    self-contained: true
execute: 
  echo: true
---

```{python}
#| include: false
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
import warnings
import time
from pprint import pprint


import jax
import jax.numpy as jnp

jax.config.update("jax_enable_x64", True)

import sklearn
sklearn.set_config(display="text")
from sklearn.model_selection import GridSearchCV
from sklearn.linear_model import LinearRegression, Ridge, Lasso

plt.rcParams['figure.dpi'] = 200

np.set_printoptions(
  edgeitems=30, linewidth=75,
  precision = 4, suppress=True
  #formatter=dict(float=lambda x: "%.5g" % x)
)

pd.set_option("display.width", 100)
pd.set_option("display.max_columns", 10)
pd.set_option("display.precision", 4)

from scipy import optimize
```

```{r setup, message=FALSE, warning=FALSE, include=FALSE}
{{< include _setup.R >}}
```

# From last time <br/> Maximum Likelihood example

## Normal MLE

:::: {.columns .xsmall}
::: {.column width='50%'}
```{python}
from scipy.stats import norm

n = norm(-3.2, 1.25)
x = n.rvs(size=100, random_state=1234)
{'µ': x.mean(), 'σ': x.std()}
```
```{python}
mle_norm = lambda θ: -np.sum(
  norm.logpdf(x, loc=θ[0], scale=θ[1])
)
```
:::

::: {.column width='50%' .fragment}
```{python}
mle_norm([0,1])
mle_norm([-3, 1])
mle_norm([-3.2, 1.25])
mle_norm([-3.3, 1.25])
```
:::
::::


## Minimizing

::: {.small}
```{python}
optimize.minimize(mle_norm, x0=[0,1], method="bfgs")
```
:::

## Adding constraints

A crude way of introducing constraints is to return `np.inf` (or another large value) for any parameter values that violate the constraint - here we require the standard deviation to be positive.

:::: {.columns .xsmall}
::: {.column width='50%'}
```{python}
def mle_norm2(θ): 
  if θ[1] <= 0:
    return np.inf
  else:
    return -np.sum(
      norm.logpdf(x, loc=θ[0], scale=θ[1])
    )
```
:::

::: {.column width='50%' .fragment}
```{python}
optimize.minimize(
  mle_norm2, x0=[0,1], method="bfgs"
)
```
:::
::::

. . .

::: {.xsmall}
```{python}
{'µ': x.mean(), 'σ': x.std()}
```
:::


## Reparameterization

An alternative approach is to transform the parameter so it is unconstrained - here we optimize over `log(σ)` instead of `σ`, using `exp()` to map back to the positive reals.

:::: {.columns .xsmall}
::: {.column width='50%'}
```{python}
def mle_norm3(θ):
  return -np.sum(
    norm.logpdf(
      x, loc=θ[0], scale=np.exp(θ[1])
    )
  )
```
:::

::: {.column width='50%' .fragment}
```{python}
res = optimize.minimize(
  mle_norm3, x0=[0, np.log(1)], method="bfgs"
)
res
```
:::
::::

. . .

::: {.xsmall}
```{python}
{'µ': res.x[0], 'σ': np.exp(res.x[1])}
```
:::


## Specifying Bounds

It is also possible to specify bounds via `bounds` but this is only available for certain optimization methods (i.e. Nelder-Mead & L-BFGS-B).

::: {.small}
```{python}
optimize.minimize(
  mle_norm, x0=[0,1], method="l-bfgs-b",
  bounds = [(-1e16, 1e16), (1e-16, 1e16)]
)
```
:::


## Exercise 1

Using `optimize.minimize()` recover the shape and scale parameters for these data using MLE.

::: {.small}
```{python}
from scipy.stats import gamma

g = gamma(a=2.0, scale=2.0)
x = g.rvs(size=100, random_state=1234)
x.round(2)
```
:::


# Stochastic Gradient Descent

## A regression example {.scrollable}

::: {.xsmall}
```{python}
#| code-line-numbers: "|3"
from sklearn.datasets import make_regression
X, y, coef = make_regression(
  n_samples=200, n_features=20, n_informative=4, 
  bias=3, noise=1, random_state=1234, coef=True
)
```
:::

::: {.columns .xxsmall}
::: {.column width=15%}
```{python}
y
```
:::
::: {.column width=70%}
```{python}
X
```
:::
::: {.column width=15%}
```{python}
coef.reshape((-1,1))
```
:::
:::

## Minimalistic GD for LR

::: {.xsmall}
```{python}
#| code-line-numbers: "|11"
def grad_desc_lm(X, y, beta, step, max_step=50):
  X = jnp.c_[jnp.ones(X.shape[0]), X]
  n, k = X.shape

  f = lambda beta: jnp.sum((y - X @ beta)**2)
  grad = lambda beta: 2*X.T @ (X@beta - y)
  
  res = {"x": [beta], "loss": [f(beta).item()], "iter": [0]}
  
  for i in range(max_step):
    beta = beta - grad(beta) * step
    res["x"].append(beta)
    res["loss"].append(f(beta).item())
    res["iter"].append(res["iter"][-1]+1)
    
  return res
```
:::


## Linear regression

::: {.xsmall}
```{python}
lm = LinearRegression().fit(X,y)
np.r_[lm.intercept_, lm.coef_]
```
:::

. . .

<br/>

::: {.xsmall}
```{python}
t0 = time.perf_counter()
gd_lm = grad_desc_lm(
  X, y, np.zeros(X.shape[1]+1),
  step = 0.001, max_step=20
)
gd_time = time.perf_counter() - t0
gd_lm["x"][-1]
```
:::


##

```{python}
#| echo: false
plt.plot("iter", "loss", data=gd_lm, label = "GD")
plt.yscale("log")
plt.xlabel("iter")
plt.ylabel("Loss")
plt.legend()
plt.show()
```

## 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(s), we can define 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}})\\
%&= 
%\left[ 
%  \begin{matrix}
%    2 \boldsymbol{X}_{\cdot 1}^T(\boldsymbol{X}_{\cdot 1}\boldsymbol{\beta}_1 - \boldsymbol{y}) \\
%    2 \boldsymbol{X}_{\cdot 2}^T(\boldsymbol{X}_{\cdot 2}\boldsymbol{\beta}_2 - \boldsymbol{y}) \\
%    \vdots \\
%    2 \boldsymbol{X}_{\cdot k}^T(\boldsymbol{X}_{\cdot k}\boldsymbol{\beta}_k - \boldsymbol{y})
%  \end{matrix} 
%\right]
\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

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}
```{python}
#| code-line-numbers: "|11-15|4,17,18|"
def sto_grad_desc_lm(X, y, beta, step, max_step=50, seed=1234, replace=True):
  X = jnp.c_[jnp.ones(X.shape[0]), X]
  f = lambda beta: jnp.sum((y - X @ beta)**2)
  grad = lambda beta, i: 2*X[i,:] * (X[i,:]@beta - y[i])
  n, k = X.shape

  res = {"x": [beta], "loss": [f(beta).item()], "iter": [0]}
  rng = np.random.default_rng(seed)

  for i in range(max_step):
    if replace:
      js = rng.integers(0,n,n)
    else:
      js = np.array(range(n))
      rng.shuffle(js)

    for j in js:
      beta = beta - grad(beta, j) * step
      res["x"].append(beta)
      res["loss"].append(f(beta).item())
      res["iter"].append(res["iter"][-1]+1)
    
  return res

```
:::


## Fitting {.scrollable}

::: {.xxsmall}
```{python}
np.r_[lm.intercept_, lm.coef_]
```

```{python}
t0 = time.perf_counter()
sgd_lm_rep = sto_grad_desc_lm(
  X, y, np.zeros(X.shape[1]+1),
  step = 0.001, max_step=20, replace=True
)
sgd_rep_time = time.perf_counter() - t0
sgd_lm_rep["x"][-1]
```
```{python}
t0 = time.perf_counter()
sgd_lm_worep = sto_grad_desc_lm(
  X, y, np.zeros(X.shape[1]+1),
  step = 0.001, max_step=20, replace=False
)
sgd_worep_time = time.perf_counter() - t0
sgd_lm_worep["x"][-1]
```
:::


## Learning by iteration

```{python}
#| echo: false
plt.figure()
plt.plot("iter", "loss", data=gd_lm, label="GD")
plt.plot("iter", "loss", data=sgd_lm_rep, label="SGD - replacement")
plt.plot("iter", "loss", data=sgd_lm_worep, label="SGD - w/o replacement")
plt.yscale("log")
plt.legend()
plt.xlabel("iter")
plt.ylabel("Loss")
plt.show()
```

## Learning by Epochs

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

```{python}
#| echo: false
plt.figure()
plt.plot("iter", "loss", data=gd_lm, label="GD")
plt.plot(
  "iter", "loss", 
  data=pd.DataFrame(sgd_lm_rep).assign(iter = lambda x: x.iter / X.shape[0]), 
  label="SGD - replacement"
)
plt.plot(
  "iter", "loss", 
  data=pd.DataFrame(sgd_lm_worep).assign(iter = lambda x: x.iter / X.shape[0]),
  label="SGD - w/o replacement"
)
plt.yscale("log")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend()
plt.show()
```

## Run times

```{python}
#| echo: false
from IPython.display import Markdown
Markdown(
  "| Method | Time / Epoch |\n"
  "|--------|-----------:|\n"
  f"| GD | {gd_time/20:.3f}s |\n"
  f"| SGD (replacement) | {sgd_rep_time/20:.3f}s |\n"
  f"| SGD (w/o replacement) | {sgd_worep_time/20:.3f}s |\n"
)
```


## A bigger example

::: {.xsmall}
```{python}
#| code-line-numbers: "|3"
from sklearn.datasets import make_regression
X, y, coef = make_regression(
  n_samples=10000, n_features=20, n_informative=4, 
  bias=3, noise=1, random_state=1234, coef=True
)
```
:::

. . .

::: {.xsmall}
```{python}
lm = LinearRegression().fit(X,y)
np.r_[lm.intercept_, lm.coef_]
```
:::

## Fitting {.scrollable}

::: {.xxsmall}
```{python}
t0 = time.perf_counter()
gd_lm = grad_desc_lm(
  X, y, np.zeros(X.shape[1]+1),
  step = 0.00005, max_step=3
)
gd_time = time.perf_counter() - t0
gd_lm["x"][-1]
```
```{python}
t0 = time.perf_counter()
sgd_lm_rep = sto_grad_desc_lm(
  X, y, np.zeros(X.shape[1]+1),
  step = 0.001, max_step=3, replace=True
)
sgd_rep_time = time.perf_counter() - t0
sgd_lm_rep["x"][-1]
```
```{python}
t0 = time.perf_counter()
sgd_lm_worep = sto_grad_desc_lm(
  X, y, np.zeros(X.shape[1]+1),
  step = 0.001, max_step=3, replace=False
)
sgd_worep_time = time.perf_counter() - t0
sgd_lm_worep["x"][-1]
```
```{python}
#| echo: false
from IPython.display import Markdown
Markdown(
  "| Method | Time / Epoch |\n"
  "|--------|-----------:|\n"
  f"| GD | {gd_time/3:.3f}s |\n"
  f"| SGD (replacement) | {sgd_rep_time/3:.3f}s |\n"
  f"| SGD (w/o replacement) | {sgd_worep_time/3:.3f}s |\n"
)
```
:::


## Results

::: {.panel-tabset} 

### Full

```{python}
#| echo: false
plt.figure()
plt.plot("iter", "loss", data=gd_lm, label="GD")
plt.plot(
  "iter", "loss", 
  data=pd.DataFrame(sgd_lm_rep).assign(iter = lambda x: x.iter / X.shape[0]), 
  label="SGD - replacement"
)
plt.plot(
  "iter", "loss", 
  data=pd.DataFrame(sgd_lm_worep).assign(iter = lambda x: x.iter / X.shape[0]), 
  label="SGD - w/o replacement"
)
plt.yscale("log")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend()
plt.show()
```

### Zoom

```{python}
#| echo: false
plt.figure()
plt.plot("iter", "loss", data=gd_lm, label="GD")
plt.plot(
  "iter", "loss", 
  data=pd.DataFrame(sgd_lm_rep).assign(iter = lambda x: x.iter / X.shape[0]), 
  label="SGD - replacement"
)
plt.plot(
  "iter", "loss", 
  data=pd.DataFrame(sgd_lm_worep).assign(iter = lambda x: x.iter / X.shape[0]), 
  label="SGD - w/o replacement"
)
plt.yscale("log")
plt.xlabel("Epoch")
plt.ylabel("Loss")
lim = plt.xlim(0,0.3)
plt.legend()
plt.show()
```

### Timings

```{python}
#| echo: false
from IPython.display import Markdown
Markdown(
  "| Method | Time / Epoch |\n"
  "|--------|-----------:|\n"
  f"| GD | {gd_time/3:.3f}s |\n"
  f"| SGD (replacement) | {sgd_rep_time/3:.3f}s |\n"
  f"| SGD (w/o replacement) | {sgd_worep_time/3:.3f}s |\n"
)
```

:::


## 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}
```{python}
#| code-line-numbers: "|4|11-15|17|18,4"
def mb_grad_desc_lm(X, y, beta, step, batch_size = 10, max_step=50, seed=1234, replace=True):
  X = jnp.c_[jnp.ones(X.shape[0]), X]
  f = lambda beta: jnp.sum((y - X @ beta)**2)
  grad = lambda beta, i: 2*X[i,:].T @ (X[i,:]@beta - y[i])
  n, k = X.shape

  res = {"x": [beta], "loss": [f(beta).item()], "iter": [0]}
  rng = np.random.default_rng(seed)

  for i in range(max_step):
    if replace:
      js = rng.integers(0,n,n)
    else:
      js = np.array(range(n))
      rng.shuffle(js)

    for j in js.reshape(-1, batch_size):
      beta = beta - grad(beta, j) * step
      res["x"].append(beta)
      res["loss"].append(f(beta).item())
      res["iter"].append(res["iter"][-1]+1)
    
  return res

```
:::

## Fitting {.scrollable}

::: {.xxsmall}
```{python}
lm = LinearRegression().fit(X,y)
np.r_[lm.intercept_, lm.coef_]
```
```{python}
sizes = [10,50,100]
mbgd = {}
mbgd_time = {}
for size in sizes:
  t0 = time.perf_counter()
  mbgd[size] = mb_grad_desc_lm(
    X, y, np.zeros(X.shape[1]+1), batch_size=size,
    step = 0.001, max_step=3, replace=False
  )
  mbgd_time[size] = time.perf_counter() - t0
```

```{python}
#| echo: false
for size in sizes:
  print(f"Batch size: {size} ({mbgd_time[size]/3:.3f}s / epoch)")
  print(mbgd[size]["x"][-1])
  print("")

```

:::


## Results

::: {.panel-tabset} 

### Full

```{python}
#| echo: false
plt.figure()
plt.plot("iter", "loss", data=gd_lm, label="GD")
plt.plot(
  "iter", "loss", 
  data=pd.DataFrame(sgd_lm_rep).assign(iter = lambda x: x.iter / X.shape[0]), 
  label="SGD - replacement"
)
plt.plot(
  "iter", "loss", 
  data=pd.DataFrame(sgd_lm_worep).assign(iter = lambda x: x.iter / X.shape[0]), 
  label="SGD - w/o replacement"
)
for size in sizes:
  plt.plot(
    "iter", "loss", 
    data=pd.DataFrame(mbgd[size]).assign(iter = lambda x: x.iter / (X.shape[0]/size)), 
    label=f"MBGD ({size})"
  )
plt.yscale("log")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend()
plt.show()
```

### Zoom

```{python}
#| echo: false
plt.figure()
plt.plot("iter", "loss", data=gd_lm, label="GD")
plt.plot(
  "iter", "loss", 
  data=pd.DataFrame(sgd_lm_rep).assign(iter = lambda x: x.iter / X.shape[0]), 
  label="SGD - replacement"
)
plt.plot(
  "iter", "loss", 
  data=pd.DataFrame(sgd_lm_worep).assign(iter = lambda x: x.iter / X.shape[0]), 
  label="SGD - w/o replacement"
)
for size in sizes:
  plt.plot(
    "iter", "loss", 
    data=pd.DataFrame(mbgd[size]).assign(iter = lambda x: x.iter / (X.shape[0]/size)), 
    label=f"MBGD ({size})"
  )
plt.yscale("log")
plt.xlabel("Epoch")
plt.ylabel("Loss")
lim = plt.xlim(0,0.15)
plt.legend()
plt.show()
```

### Timings

```{python}
#| echo: false
from IPython.display import Markdown
rows = (
  "| Method | Time / Epoch |\n"
  "|--------|-----------:|\n"
  f"| GD | {gd_time/3:.3f}s |\n"
  f"| SGD (replacement) | {sgd_rep_time/3:.3f}s |\n"
  f"| SGD (w/o replacement) | {sgd_worep_time/3:.3f}s |\n"
)
for size in sizes:
  rows += f"| MBGD ({size}) | {mbgd_time[size]/3:.3f}s |\n"
Markdown(rows)
```

:::


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


# Adaptive updates & Momentum 

## AdaGrad

This approach was proposed by Duchi, Hazan, & Singer in 2011 and is based on the idea of scaling the learning rates for the current step by the sum of the square gradients of previous steps - this has the effect of shrinking the step size of dimensions with large previous gradients.

$$
\begin{aligned}
\boldsymbol{\theta}_{t+1} &= \boldsymbol{\theta}_t - \eta_t \frac{1}{\sqrt{\boldsymbol{s}_t + \epsilon}} \odot \nabla \ell(\boldsymbol{X}_{B_t}, \boldsymbol{\theta}_t)\\
\boldsymbol{s}_t &= \sum_{i=1}^t \left(\nabla \ell(\boldsymbol{X}_{B_i}, \boldsymbol{\theta}_i)\right)^2
\end{aligned}
$$

here $\epsilon$ is a small constant (i.e. $10^{-7}$) to avoid division by zero.


## Implementation

::: {.xsmall}
```{python}
#| code-line-numbers: "|10|20-23"
def adagrad_lm(X, y, beta, step, batch_size = 10, max_step=50, seed=1234, replace=True, eps=1e-8):
  X = jnp.c_[jnp.ones(X.shape[0]), X]
  f = lambda beta: jnp.sum((y - X @ beta)**2)
  grad = lambda beta, i: 2*X[i,:].T @ (X[i,:]@beta - y[i])
  n, k = X.shape

  res = {"x": [beta], "loss": [f(beta).item()], "iter": [0]}
  rng = np.random.default_rng(seed)

  S = np.zeros(k)

  for i in range(max_step):
    if replace:
      js = rng.integers(0,n,n)
    else:
      js = np.array(range(n))
      rng.shuffle(js)

    for j in js.reshape(-1, batch_size):
      G = grad(beta, j)
      S += G**2
      
      beta = beta - step * (1/np.sqrt(S + eps)) * G
      
      res["x"].append(beta)
      res["loss"].append(f(beta).item())
      res["iter"].append(res["iter"][-1]+1)
    
  return res

```
:::


## A medium example

::: {.xsmall}
```{python}
#| code-line-numbers: "|3"
from sklearn.datasets import make_regression
X, y, coef = make_regression(
  n_samples=1000, n_features=20, n_informative=4, 
  bias=3, noise=1, random_state=1234, coef=True
)
```
:::

. . .

::: {.xsmall}
```{python}
lm = LinearRegression().fit(X,y)
np.r_[lm.intercept_, lm.coef_]
```
:::

## Fitting

::: {.xxsmall}
```{python}
sizes = [1, 25, 50, 1000]
lrs = [10] * 4
algos = ["AdaGrad - SGD", "AdaGrad - MBGD (25)", "AdaGrad - MBGD (50)", "AdaGrad - GD"]

adagrad = { size: adagrad_lm(
                    X, y, np.zeros(X.shape[1]+1), batch_size=size,
                    step = lr, max_step=7, replace=True, eps=1e-8
                  )
            for size, lr in zip(sizes,lrs) }
```

```{python}
#| echo: false
for size, algo in zip(sizes, algos):
  print(algo)
  print(adagrad[size]["x"][-1])
  print("")

```
:::

## Results

```{python}
#| echo: false
plt.figure()
for size, algo in zip(sizes, algos):
  plt.plot(
    "iter", "loss", 
    data=pd.DataFrame(adagrad[size]).assign(iter = lambda x: x.iter / (X.shape[0]/size)), 
    label=algo
  )
plt.yscale("log")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend()
plt.show()
```


## RMSProp

With AdaGrad the denominator involving $\boldsymbol{s}_t$ gets larger as $t$ increases, but in some cases it gets too large too fast to effectively explore the loss function. An alternative is to use a moving average of the past squared gradients instead.

RMSProp replaces AdaGrad's $\boldsymbol{s}_t$ with the following,

$$
\boldsymbol{s}_t = \beta \, \boldsymbol{s}_{t-1} + (1-\beta) \, (\nabla \ell(\boldsymbol{X},\boldsymbol{\theta}_t))^2
$$

$$
\boldsymbol{s}_0 = \boldsymbol{0}
$$

in practice a value of $\beta \approx 0.9$ is often used.


## Implementation

::: {.xsmall}
```{python}
#| code-line-numbers: "|21"
def rmsprop_lm(X, y, beta, step, batch_size = 10, max_step=50, seed=1234, replace=True, eps=1e-8, b=0.9):
  X = jnp.c_[jnp.ones(X.shape[0]), X]
  f = lambda beta: jnp.sum((y - X @ beta)**2)
  grad = lambda beta, i: 2*X[i,:].T @ (X[i,:]@beta - y[i])
  n, k = X.shape
  
  res = {"x": [beta], "loss": [f(beta).item()], "iter": [0]}
  rng = np.random.default_rng(seed)
  
  S = np.zeros(k)
  
  for i in range(max_step):
    if replace:
      js = rng.integers(0,n,n)
    else:
      js = np.array(range(n))
      rng.shuffle(js)
    
    for j in js.reshape(-1, batch_size):
      G = grad(beta, j)
      S = b*S + (1-b) * G**2
      
      beta = beta - step * (1/np.sqrt(S + eps)) * G
      
      res["x"].append(beta)
      res["loss"].append(f(beta).item())
      res["iter"].append(res["iter"][-1]+1)
    
  return res

```
:::

## Fitting

::: {.xxsmall}
```{python}
sizes = [1, 25, 50, 1000]
lrs = [0.01, 0.1, 0.25, 1]
algos = ["RMSProp - SGD", "RMSProp - MBGD (25)", "RMSProp - MBGD (50)", "RMSProp - GD"]

rmsprop = { size: rmsprop_lm(
                    X, y, np.zeros(X.shape[1]+1), batch_size=size,
                    step = lr, max_step=25, replace=True
                  )
            for size, lr in zip(sizes,lrs) }
```

```{python}
#| echo: false
for size, algo in zip(sizes, algos):
  print(algo)
  print(rmsprop[size]["x"][-1])
  print("")

```
:::

## Results

```{python}
#| echo: false
plt.figure()
for size, algo in zip(sizes, algos):
  plt.plot(
    "iter", "loss", 
    data=pd.DataFrame(rmsprop[size]).assign(iter = lambda x: x.iter / (X.shape[0]/size)), 
    label=algo
  )
plt.yscale("log")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend()
plt.show()
```

## Momentum

Rather than just using the gradient information at our current location it may be beneficial to use information from our previous steps as well. A general setup for this type approach looks like,

::: {.small}
$$
\boldsymbol{\theta}_{t+1} = \boldsymbol{\theta}_t - \eta \, \boldsymbol{m}_t
$$

$$
\boldsymbol{m}_t = \beta \, \boldsymbol{m}_{t-1} + (1-\beta) \, \nabla \ell(\boldsymbol{X}, \boldsymbol{\theta}_t)
$$
:::

where $\eta$ is our step size and $\beta$ determines the weighting of the current gradient and the previous gradients.

If you have taken a course on time series, this has a flavor that looks a lot like moving average models,

::: {.small}
$$
\boldsymbol{m}_t = (1-\beta) \, \nabla \ell(\boldsymbol{X}, \boldsymbol{\theta}_t) + \beta(1-\beta) \, \, \nabla \ell(\boldsymbol{X}, \boldsymbol{\theta}_{t-1}) + \beta^2(1-\beta) \, \, \nabla \ell(\boldsymbol{X}, \boldsymbol{\theta}_{t-2}) + \cdots
$$
:::

## Adam

The "adaptive moment estimation" algorithm is a combination of momentum with RMSProp,

$$
\begin{aligned}
\theta_{t+1} &= \theta_t - \eta_t \frac{\boldsymbol{m_t}}{\sqrt{\boldsymbol{s}_t + \epsilon}} \\
\boldsymbol{m}_t &= \beta_1 \, \boldsymbol{m}_{t-1} + (1-\beta_1) \, \nabla \ell(\boldsymbol{X}, \theta_t) \\
\boldsymbol{s}_t &= \beta_2 \, \boldsymbol{s}_{t-1} + (1-\beta_2) \, (\nabla \ell(\boldsymbol{X},\boldsymbol{\theta}_t))^2 \\
\end{aligned}
$$

Note that RMSProp is a special case of Adam when $\beta_1 = 0$.

Adam is widely used in practice and is commonly available within tools like Torch for fitting NN models. 

In typical use $\beta_1=0.9$, $\beta_2=0.999$, $\epsilon=10^{-6}$, and $\eta_t=0.001$ are used. As the learning rate is not guaranteed to decrease over time, the algorithm is not guaranteed to converge.

## Bias corrections

One small alteration that was suggested by the original authors and is commonly used is to correct for the bias towards small values in the initial estimates of $\boldsymbol{m}_t$ and $\boldsymbol{s}_t$. In which case they are replaced with,

$$
\begin{aligned}
{\hat{\boldsymbol{m}}}_t &= \boldsymbol{m}_t / (1-{\beta_1}^t) \\
{\hat{\boldsymbol{s}}}_t &=\boldsymbol{s}_t / (1-{\beta_2}^t) \\
\end{aligned}
$$

## Implementation

::: {.xsmall}
```{python}
#| code-line-numbers: "|10-12|22-25|27,28|30"
def adam_lm(X, y, beta, step=0.001, batch_size = 10, max_step=50, seed=1234, replace=True, eps=1e-6, b1=0.9, b2=0.999):
  X = jnp.c_[jnp.ones(X.shape[0]), X]
  f = lambda beta: jnp.sum((y - X @ beta)**2)
  grad = lambda beta, i: 2*X[i,:].T @ (X[i,:]@beta - y[i])
  n, k = X.shape
  
  res = {"x": [beta], "loss": [f(beta).item()], "iter": [0]}
  rng = np.random.default_rng(seed)
  
  S = np.zeros(k)
  M = np.zeros(k)
  t = 0
  
  for i in range(max_step):
    if replace:
      js = rng.integers(0,n,n)
    else:
      js = np.array(range(n))
      rng.shuffle(js)
    
    for j in js.reshape(-1, batch_size):
      t += 1
      G = grad(beta, j)
      S = b2*S + (1-b2) * G**2
      M = b1*M + (1-b1) * G
      
      M_hat = M / (1-b1**t)
      S_hat = S / (1-b2**t)
      
      beta = beta - step * (M_hat / np.sqrt(S_hat + eps))
      
      res["x"].append(beta)
      res["loss"].append(f(beta).item())
      res["iter"].append(t)
    
  return res

```
:::

## Fitting

::: {.xxsmall}
```{python}
sizes = [1, 25, 50, 1000]
lrs = [0.01, 0.5, 0.75, 1]
algos = ["Adam - SGD", "Adam - MBGD (25)", "Adam - MBGD (50)", "Adam - GD"]

adam = { size: adam_lm(
                 X, y, np.zeros(X.shape[1]+1), batch_size=size,
                 step=lr, max_step=25, replace=True
               )
         for size, lr in zip(sizes,lrs) }
```

```{python}
#| echo: false
for size, algo in zip(sizes, algos):
  print(algo)
  print(adam[size]["x"][-1])
  print("")

```
:::

## Results

```{python}
#| echo: false
plt.figure()
for size, algo in zip(sizes, algos):
  plt.plot(
    "iter", "loss", 
    data=pd.DataFrame(adam[size]).assign(iter = lambda x: x.iter / (X.shape[0]/size)), 
    label=algo
  )
plt.yscale("log")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend()
plt.show()
```