CS58700 - Foundations of Deep Learning

Spring 2025

Instructor and Copyright: Bruno Ribeiro

---

Stochastic Optimization, Learning, and the Lottery Ticket Hypothesis

Reading:

The reading assigments for this lecture:

Reading assigment: Numerical Computation

• Chapter 4.1 (Overflow and Underflow)
  • Softmax issues
• Chapter 4.2 (Poor Conditioning)
• Chapter 4.3 (Gradient-Based Optimization)
  • Gradient descent, critical points, local minima and maxima, saddle points
  • Jacobian matrix

Reading assigment: Regularization for Deep Learning

• Chapter 7.1.1 (L2 Parameter Regularization = Weight Decay)
• Chapter 7.1.2 (L1 Regularization)
• Chapter 7.8 (Early Stopping)

Reading assigment: Optimization for Training Deep Models

• Chapter 8.1 (How Learning Differs from Pure Optimization)
• Chapter 8.2 (Challenges in Neural Network Optimization)
• Chapter 8.3 (Basic Algorithms)
  • Chapter 8.3.1 (Stochastic Gradient Descent)
  • Chapter 8.3.2 (Momentum)
  • Chapter 8.3.3 (Nesterov Momentum)
• Chapter 8.4 (Parameter Initialization Strategies)
• Chapter 8.5 (Algorithms with Adaptive Learning Rates)
  • Chapter 8.5.3 (Adam)
  • Chapter 8.5.4 (Choosing the Right Optimization Algorithm)
• Chapter 8.7.1 (Batch Normalization)

Optional Reading:

• Léon Bottou, Stochastic Gradient Descent Tricks, Neural Networks: Tricks of the Trade, 2012.
• Herbert Robbins and Sutton Monro. "A stochastic approximation method", The Annals of Mathematical Statistics, 1951.
• Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. ICLR, 2019.

Optimization $\neq$ Learning

In past last lectures we have used backpropagation to compute gradients of MLPs.

• MLP is just one of multiple neural network architectures.
• To learn the MLP neuron weights over the training data we used gradient descent
• It is, however, important to note that Optimization $\neq$ Learning

In what follows we will assume a general function with bounded (finite) gradients.

Task:

• Supervised learning: training dataset $D = \{x_i,y_i\}_{i=1}^N$, where $x_i \in \mathbb{R}^m$ and $y_i \in \mathbb{R}$ are the input and output of training example $i$, respectively.

• Goal: Minimize a loss $\sum_{i=1}^N L(x_i,y_i; \theta)$, $L:\mathbb{R}^m \times \mathbb{R} \times \mathbb{R}^d \to \mathbb{R}$, given the function's gradient $\sum_{i=1}^N \nabla_{\theta} L(x_i,y_i;\theta_t)$ at the gradient descent step $t \geq 0$.

• Gradient descent optimization is a procedure to find parameters that minimize a loss function over the training data.

• Machine learning is concerned with finding a model that, given a finite amount of training data, will generalize well to test data.
  • Training and generalization errors are generally different
  • It is always assumed that training and test distributions are the same
    • This will also be our assumption today

• It is important to not get these optimization and learning confused.
  • Minimizing the training loss is not the end-goal, it is just a tool among many other tools.
  • Other tools include sequential hypothesis testing and transient Markov chains (we will talk about this later)

1. Finding a Global Optimum is Not the Primary Goal

• Overparameterized networks can often fit training data perfectly (achieving zero or near-zero training loss)
• Multiple global optima exist, and they differ dramatically in generalization
• The optimization algorithm's trajectory matters: which optimum we reach depends on initialization, learning rate, batch size, etc.
• Takeaway: Machine learning $\neq$ optimization. We care about test performance, not just minimizing training loss.

2. Local Minima Are Rarely Problematic

Classical ML worry about "getting stuck" in poor local minima largely disappears in high dimensions:

• Saddle points dominate the landscape: In high-dimensional spaces, most critical points are saddle points, not local minima (Dauphin et al., NeurIPS 2014)

  • Intuition: For a critical point to be a local minimum, the Hessian must be positive definite—all eigenvalues positive. With millions of parameters, this becomes exponentially unlikely by chance.

• When local minima do exist, they tend to be "good": Empirical and theoretical work suggests that in overparameterized networks, local minima have loss values close to global minima (Choromanska et al., AISTATS 2015)

• The loss landscape is highly connected: Rather than isolated valleys, think of the geometry as broad, flat regions connected by paths of near-constant loss

3. What Does Pose Challenges?

• Saddle points can slow optimization (though momentum and adaptive methods help)
• Flat regions and plateaus create slow convergence
• Ill-conditioning: Large variation in curvature across directions
• Implicit regularization: Understanding which global optimum SGD finds

Gradient Descent Optimization

In what follows we will assume a general function with bounded (finite) gradients.

Task:

• Supervised learning: training dataset $D = \{x_i,y_i\}_{i=1}^N$, where $x_i \in \mathbb{R}^m$ and $y_i \in \mathbb{R}$ are the input and output of training example $i$, respectively.

• Goal: Minimize a loss $\sum_{i=1}^N L(x_i,y_i; \theta)$, $L:\mathbb{R}^m \times \mathbb{R} \times \mathbb{R}^d \to \mathbb{R}$, given the function's gradient $\sum_{i=1}^N \nabla_{\theta} L(x_i,y_i;\theta_t)$ at the gradient descent step $t \geq 0$.

Sampled Mini-batches for the Stochastic Gradient

Rather than computing the a gradient over the entire training dataset, we will compute stochastic gradients, which are computed over small subsets (mini-batches) of the training data:

• If the subset of the data is chosen uniformly at random, its gradient is called stochastic gradient
• We can choose between sampling this subset of the training data with or without replacement

We will prefer sample without replacement. This is because we want to avoid giving more importance to some examples over others in the optimization. Choosing with replacement means that at each $N$ training examples $$ P(\text{omit }x_i) = (1 - 1/N)^N \approx \exp(-1) = 0.37\dots $$

This high probability of omission of some training examples, leads to an increased variance and decreased data efficiency relative to sampling without replacement.

• Hence, we will use sampling without replacement for our mini-batches

Sampling without replacement procedure:

• An epoch is the act of going over the entire training data once

• We will make a vector of training data indices $(1,\ldots,N)$

• At each epoch we will shuffle the vector $\text{epochid} = \text{shuffle}( 1,\ldots,N)$

• Divide epoch_idx into equal-size mini-batches.
  • Say, mini-batch has size $k$, then
    • The first mini-batch uses examples $D_1 = \{(x_j,y_j) : j = \text{epochid}[0:k]\}$
    • The second mini-batch uses examples $D_2 = \{(x_j,y_j) : j = \text{epochid}[k:2k]\}$
    • ...

• If $N$ is not a multiple of $k$, we ignore the last mini-batch (whose size is smaller than $k$).

• Once all mini-batches are exhausted, we start a new epoch

The Stochastic Gradient is an Unbiased Gradient Estimate

The stochast gradient computed over a mini-batch is an unbiased estimate of the gradient computed over the entire training dataset

• The mini-batch training data $\mathcal{D}_k$ are i.i.d. samples (without replacement) from the training data $\mathcal{D}$

  • Each example sampled with probability $p(x,y)$.

• The loss functions is a linear function of the loss over each training example, i.e., $\mathcal{L}(\mathcal{D}_k;\theta) = \frac{1}{k} \sum_{i=1}^k L(x_i,y_i; \theta)$.

  • Hence the stochastic gradient at parameter value $\theta_t$ is $\nabla_\theta \mathcal{L}(\mathcal{D}_k;\theta_t) = \frac{1}{k} \sum_{i=1}^k \nabla_\theta L(X_i,Y_i; \theta_t)$.

• Then, the average shows $$\begin{align} E[\nabla_\theta \mathcal{L}(\mathcal{D}_k;\theta_t)] & = \frac{1}{k} \sum_{i=1}^k E\left[\nabla_\theta L(X_i,Y_i; \theta_t) \right] \\ & = \frac{1}{|\mathcal{D}|} \sum_{(x,y) \in \mathcal{D}} \nabla_\theta L(x,y; \theta_t) \end{align}$$ that the stochastic gradient is an unbiased estimate of the gradient over the entire training dataset.

• Note that while the average is the same, the stochastic gradient has a variance that depends on $k$.

-If we want to increase the learning rate but keep the same amount of gradient noise, we should also increase the mini-batch size (e.g., (Smith et al., ICLR 2018))

• Q: If we double the learning rate, by how much should we increase the batch size to keep the gradient variance the same?
  • (A) Half the batch size (1/2 x batch_size)
  • (B) Double the batch size (2 x batch_size)
  • (C) Keep batch size the same

Stochastic Gradient Descent (SGD) as a Markov Chain

Stochastic gradient descent can be seen as a Markov chain:

Starting from parameters $\theta_0$, the SGD update $$ \theta_1 = \theta_0 - \gamma_1 \frac{1}{k} \sum_{i=1}^k \nabla_\theta L(X_i,Y_i;\theta_0), $$ where $X_i$ and $Y_i$ are uniformly sampled training data examples. This will obtain new paramters $\theta_1$ as $$ \theta_1 \sim P(\theta | \theta_0), $$ due to the randomization created by the mini-batch sampling.

The next set of gradient step only depend on $\theta_1$, hence $$ \theta_2 \sim P(\theta | \theta_1). $$

More generally, at step $t > 0$, $$ \theta_t \sim P(\theta | \theta_{t-1}). $$

Early stopping:

• We should understand this Markov chain as generating model "proposals"
  • We will evaluate each proposal against the validation data, keeping only the best parameter proposal
  • Once the best parameter does not change for a few epochs, we can stop the optimization and output the best parameters
  • The above procedure is called early stopping
  • Hence, early stopping is a sequential hypothesis test procedure

Parameter initialization matters

Since we do early stopping, the transient of our Markov chain is very important.

The initialization of the neural network weights $\theta_0$ matters, since $$ \theta_t \sim P(\theta | \theta_{t-1})\cdots P(\theta_1 | \theta_0). $$

Types of Parameter Initialization

Glorot initialization (based on Feedforward networks) (Glorot AISTATS 2010):

1. Initialize hidden layer biases to zero and output biases to optimal value if weights were zero (e.g., inverse sigmoid of mean target).
2. Initialize weight of neuron $i$ as $W_i \sim N(0, r)$, where $N()$ is the normal di or $W_i \sim \text{Uniform}(−\sqrt{r}, \sqrt{r})$, where for sigmoid activations $$ r = \frac{24}{\text{fan-in} + \text{fan-out}}, $$ where fan-in is the number of input neurons to neuron $i$ and fan-out is the number of outputs of neuron $i$. Note that the denominator is the number of weights in and out of neuron $i$. For tanh activations $$ r = \frac{6}{\text{fan-in} + \text{fan-out}}. $$ The constants are related to the activation thresholds of the different functions. For ReLUs, $$ r = \frac{1}{\text{fan-in} + \text{fan-out}} $$ should suffice.

PS: Make sure the initial gradients look OK. The neural network will first aggressively optimize the last layer.

• Why? The answer is simpler to explain with Normal initializations. Consider two independent Normal random variables [base source Wikipedia]: $$ X\sim N(\mu _{X},\sigma _{X}^{2}),$$ $$ Y\sim N(\mu _{Y},\sigma _{Y}^{2}),$$ and their sum $$Z=X+Y.$$ Then $$ Z\sim N(\mu _{X}+\mu _{Y},\sigma _{X}^{2}+\sigma _{Y}^{2}). $$

Proof of Glorot initialization: The distribution of X + Y is normal $N(\mu _{X}+\mu _{Y},\sigma _{X}^{2}+\sigma _{Y}^{2})$.

A very simplified version of Glorot's argument goes like this:

• (Fan-in) Assume all the layer $l$ inputs to be ${\bf 1}$. Consider a neuron $i$ at layer $l+1$ and initialize its parameters ${\bf w}^{(i)}$ with samples from $\text{Normal}(0,\sigma^2)$. The fan-in is the maximum number of Normally distributed values that we will sum: $({\bf w^{(i)}})^T {\bf 1}$. Thus, the input will be distributed as $N(0,(\text{fan-in})\times\sigma^2)$. Give too much variance to a sigmoid activation and it will saturate and return very small gradients.
  
• (Fan-out) Same analysis but thinking now of backpropagation to get $N(0,(\text{fan-out})\times\sigma^2)$.
  The division by $(\text{fan-in}+\text{fan-out})$ is a compromise to keep the variance OK both for the foward and backward passes.
  
• Glorot and Bengio's argument is more elaborate, see (Glorot AISTATS 2010).
  

Adaptive Optimizers

In case we prefer to generate model "proposals" using something a little more flexible than a Markov chain, we can use adaptive optimizers.

Adaptive optimizers generally have memory, and hence cannot be described as Markov chains.

There are a number of adaptive optimizers (see Chapter 8.5 of (Deep Learning) and Chapters 11.6 to 11.10 in (Dive into Deep Learning))

We will only cover two cases: Momentum and Adam

Momentum

SGD has dificulties navigating valleys, areas where the energy surface is much more steeply in one direction than in another direction. In valleys, SGD may oscillate across the slopes of the valley while making slow progress towards the bottom of the valley.

Momentum helps accelerate SGD in the most "common" direction, which is the bottom of the valley, dampening oscillations as the movement of going up and down the wall of the valley will tend to cancel out. It does so by making the following gradient update: $$ \theta_t = \theta_{t-1} - v_t, $$ where $$ v_t = \gamma v_{t-1} + \eta_t \nabla_\theta L(\theta_{t-1}). $$

Adam (Kingma and Ba, ICLR 2015)

Adaptive Moment Estimation (Adam) computes an adaptive learning rate for each parameter in the model. It also saves an exponentially decaying average of past squared gradients $v_t$, similar to momentum. Let $\beta_1,\beta_2 \in (0,1)$, then $$ m_t = \beta_1 m_{t-1} + (1-\beta_1) \nabla_{\theta} L(\theta_{t-1}) $$ $$ v_t = \beta_2 v_{t-1} + (1-\beta_2) (\nabla_{\theta} L(\theta_{t-1}))^2, $$ where $(\nabla_{\theta} L(\theta_{t-1}))^2$ is defined as the element-wise product $\nabla_{\theta} L(\theta_{t-1}) \odot \nabla_{\theta} L(\theta_{t-1})$. $m_t$ and $v_t$ are the (exponentially weighted) estimates of the average gradient and its second moment.

Adam bias correction: As $m_0$ and $v_0$ are initialized with 0's, there is a bias towards zero during the initial steps of the optimization. A partial correction for this bias is $$ \hat{m}_t = \frac{m_t}{1 -\beta^t_1} $$ and $$ \hat{v}_t = \frac{v_t}{1 -\beta^t_2}. $$

The gradient update is then $$ \theta_t = \theta_{t-1} - \frac{\eta_t}{\sqrt{\hat{v}_t} + \epsilon} \hat{m}_t, $$ which controls the variance of the gradient ($1/\sqrt{\hat{v}_t}$ is a element-wise square root) and gets a running estimate of its value $\hat{m}_t$. The value of $\epsilon$ is to avoid division by very small numbers when the variance is too small. (Balles and Hennig, ICML 2018) have a nice paper explaining how Adam connects with our broader understanding of stochastic optimization.

The original Adam paper proposes $\beta_1 = 0.9$, $\beta_2 = 0.99$ and $\epsilon = 10^{-8}$. The value of $\beta_1$ cannot be too close to 1 otherwise we will be treating the loss manifold as globally Euclidean, which is incorrect.

Memory Requirements of Adaptive Optimizers

Compared to vanilla SGD (which stores only the parameters), these methods require additional state:

Method	Extra Storage	Memory Overhead
Momentum	velocity $v_t$	+1× parameter count
Adam	first moment $m_t$ + second moment $v_t$	+2× parameter count

Momentum maintains one additional vector of the same dimensionality as θ to track the exponentially weighted velocity.

Adam maintains two additional vectors: $m_t$ for the running mean of gradients and $v_t$ for the running mean of squared gradients. Note that $\hat{m}_t$ and $\hat{v}_t$ don't require separate storage—they're computed on-the-fly from $m_t$ and $v_t$ using the scalar correction factors.

This becomes significant for large models. For instance, a 7B parameter model in bf16 requires ~14GB just for weights; Adam adds another ~28GB for optimizer state. This is one reason why memory-efficient variants (AdaFactor, 8-bit Adam, etc.) and techniques like optimizer state sharding across GPUs have become important for LLM training.

Why Stochastic Optimization Matters

• We will now try to understand why overparameterized neural networks can generalize well to test data
  • We do not fully understand it yet
  • A lot of what we will see is speculative but with mounting empirical evidence

Evidence of flat-region to flat-region optimization:

1. The neural network should be overparameterized (able to severely overfit the training data). In the small-parameter regime, the flat-regions are small and the optimization will get stuck at very small flat-regions.
   • Small flat-regions tend to generalize poorly while large ones have better generalization error, see (Basldassi et al., PNAS 2016), (Chaudhari et al. ICLR 2017), (Dziugaite & Roy, UAI 2017), (Achille & Soatto, JMLR 2018), (Smith & V.Le, ICLR 2018), and (Zhang et al. ICLR 2017).
   • While the quest for optimization objectives that find large flat-regions is old (Hinton & Camp, 1993) and (Hochreiter & Schmidhuber, 1997), only recently there has been some more concrete progress for deep learning models (Chaudhari et al. ICLR 2017). But because these methods don't seem to find better solutions than SGD, this line of research has received little attention recently
2. (Goodfellow et al., 2014) observes that the training loss of a straight between the initialization $\theta_0$ and the final model parameter $\theta_\text{final}$ tends to decay monotonically.
3. (Nakkira et al. 2019) shows that overparameterized neural networks generalize better (although the result was known before) (Author's figure)

Stochastic optimization will help us:

1. Find large flat minima regions
2. Avoid getting stuck on a set of mid-size flat-regions

Hypothesis: Stochastic Gradient Descent (SGD) preference for large flat minima regions

• (Standard gradient descent) Consider a full-batch gradient at step $t-1$ using the entire dataset

  • This is a low-noise gradient, pointing to the lowest minima in the region of paramter $\theta_{t-1}$
  • Suppose there is a low small flat-region near $\theta_{t-1}$.
  • Then, the gradient will prefer the small flat-region

• (SGD) Consider a stochastic gradient at step $t-1$ using a mini-batch

  • This is a noisy gradient, pointing to various lower-loss regions in the vicinity of paramter $\theta_{t-1}$
  • With SGD the larger flat-regions have an advantage over small flat-regions, due to their size

Hypothesis: Optimizing from flat-minima to flat-minima

Consider a MLP for handwritten digit classification. Let $\theta_0$ be the initial (random) parameters of this network and let $\theta_\text{final}$ be the final parameters after training with SGD until we obtain very good accuracy. (Goodfellow et al., ICLR 2015) makes the following experiment with this neural network.

• Let $L(\theta)$ be the loss over the data (training or validation) for parameter $\theta$.
  • We start with parameters $\theta_0$ chosen randomly from a standard Gaussian
  • Let $\theta_\text{final}$ be the paramters after a large number of gradient descent steps

• Let $\theta(\alpha) = (1-\alpha)\theta_0 + \alpha \theta_\text{final}$ be a straight line between the initial parameters and the final parameters.

• Plot the loss function $L(\theta(\alpha))$ as a function of $\alpha$ (figure adapted from Goodfellow et al.).
  • This shows how the loss function behaves if we were following a straight line from the initial $\theta_0$ to the final $\theta_\text{final}$.
  • At the end we do this interpolation with a Convolutional Neural Network over the MNIST dataset

Learning rates and mini-batch sizes

Learning rates and mini-batch sizes are optimization hyperparameters that should be considered modeling decisions. The initial learning rate and the SGD batch size are likely the two single most important hyperparameters for training a deep neural network.

• Learning rate: we always need to make sure to try different learning rates, as model performance depends tremendously on this hyperparameter. Typical values need to be in the range $(10^{-6},1)$ for FP32 and $(10^{-10},1)$ for FP64. Default values of $10^{-1}$, $10^{-2}$, and $10^{-3}$ typically work in most architectures.
  • Learning rate schedulers: Decreasing or adapting the learning rate schedule is necessary for some architectures. The default action is to keep the learning rate constant over training iterations. Often the benefit of choosing a Robbins-Monro learning rate scheduler is small (except in some Reinforcement Learning tasks). A good learning rate scheduler is $O(1/t^\alpha)$ with plateaus (as we have seen in previous classes), $\alpha \in (0,1)$.
    • Bach and Moulines (2011) argue for small values of $\alpha$ when using a gradient averaging or momentum-based techniques.

• Mini-batch sizes: The mini-batch size hyperparameter $B$ can significantly impact training. Typical mini-batch size values are $B = 32$ and $B = 64$.
  • Performance-wise, $B$ impacts the amount of memory and time used in backpropagation.
  • Solution-wise, mini-batch sizes have enourmous impacts on accuracy (both training and test accuracies).
    • Mini-batches are the main source of gradient noise, which is very important to find models with better generalization (as we see next).

Impact of mini-batch sizes

Example (MNIST digit classification CNN (we will cover CNNs later)):
How "flat" are the solutions with large and small mini-batch sizes? We now reproduce the (Goodfellow et al., ICLR 2015) experiment.

• Let $L(\theta)$ be the loss over validation data for parameter $\theta$.
• Let $\theta_\text{final}$ be the final trained parameter and $\theta_0$ be the initial (random) parameter.
• Let $\theta(\alpha) = (1-\alpha)\theta_0 + \alpha \theta_\text{final}$. This function denotes a straight line passing through the initial and final parameters.
• Plot the loss function $L(\theta(\alpha))$ as a function of $\alpha$.
• This is a rough measure of curvature of the manifold SGD sees in the optimization function, the "flatness" or "sharpness" around a solution.

• Mini-batch size $B = 32$, epochs = 10, learning rate = 0.05 (training time: 1 minute GTX 1050 2GB):

     Train Epoch: 10. Training Loss: 0.000354
     Validation set: Average loss: 0.0326, Accuracy: 9910/10000 (99%)
  
     alpha = 0.0 has validation error 93.4%
     alpha = 0.1666666716337204 has validation error 36.49%
     alpha = 0.3333333432674408 has validation error 3.010000000000005%
     alpha = 0.5 has validation error 1.3400000000000034%
     alpha = 0.6666666865348816 has validation error 1.0600000000000023%
     alpha = 0.8333333730697632 has validation error 0.9399999999999977%
     alpha = 1.0 has validation error 0.9000000000000057%
     alpha = 1.1666667461395264 has validation error 0.8700000000000045%
     alpha = 1.3333333730697632 has validation error 0.9099999999999966%
     alpha = 1.5 has validation error 0.9699999999999989%

• Why are learning rate and batch sizes so curcial for model accuracy?
  • Consider the difference in probabilities of hitting a "not-so-flat" solution with a noisy gradient against the probability of hitting a "not-so-flat" solution with an accurate gradient.

Effect of Larger Gradient Steps

A number of recent empirical results show that the right amount of gradient variance in SGD is essential to train some types of deep neural networks.

• E.g., (Goyal et al., NeurIPS 2017) trainins ImageNet in 1 hour by adjusting the learning rates as a function of minibatch size to keep the gradient noise constant
  • Goyal et al. argues: When the minibatch size is multiplied by k, multiply the learning rate by k.
  • It makes more sense, theoretically, to say that increasing the mini-batch size $k$-fold should increase the learning rate by $\sqrt{k}$.
    • In a mini-batch, increasing the mini-batch size $k$-fold will divide its variance by $1/k$ and its standard deviation by $1/\sqrt{k}$
    • To keep the standard deviation (variance) the same, we must multiply the learning rate by $\sqrt{k}$

Weight Space Symmetries and Invariances

An important but often overlooked aspect of neural network optimization is the structure of the weight space—the high-dimensional space of all possible parameter configurations. Understanding the symmetries and invariances in this space provides fundamental insights into training dynamics, loss landscape geometry, and modern techniques like model merging and parameter-efficient fine-tuning.

What Are Weight Space Symmetries?

Definition: A weight space symmetry is a transformation $T \in G$ of the network parameters $\theta$ that preserves the network's input-output function for all inputs: $$ f(x; \theta) = f(x; T \circ \theta) \quad \forall x, \forall T \in G. $$

The most important weight space symmetry in neural networks is permutation symmetry:

Permutation Symmetry in MLPs

In a multilayer perceptron, the neurons within a hidden layer have no inherent ordering. Permuting the neurons in hidden layer $k$—while consistently updating the weights connecting to and from that layer—produces a functionally identical network.

Formally: For a network with $d-1$ hidden layers with $n_k$ neurons in layer $k$, there are $$ \prod_{k=1}^{d-1} n_k! $$ equivalent parameter configurations that implement the exact same function (Hecht-Nielsen, 1990).

Example: Consider a two-layer network with 100 hidden neurons. There are $100! \approx 9.3 \times 10^{157}$ equivalent global minima arising from permutation symmetry alone!

Implications for the Loss Landscape

Weight space symmetries have profound consequences for the structure of the loss landscape:

1. Multiple Equivalent Global Minima

Due to permutation symmetry, any global minimum has $(n_1! \times n_2! \times \cdots \times n_{d-1}!)$ equivalent partner minima, all with identical loss values. This creates a highly structured loss landscape.

2. Permutation Points and Saddle Points

(Brea et al., 2019) showed that the paths connecting equivalent global minima pass through critical points called permutation points, where neurons in a layer have identical weight configurations. Key findings include:

• Permutation points are critical points with at least $n_{k+1}$ vanishing Hessian eigenvalues (flat directions)
• These points transition into flat high-dimensional plateaus enabling all $n_k!$ permutations at the same loss
• Many permutation points correspond to saddle points ("permutation saddles")

This provides a geometric explanation for the prevalence of plateaus and saddles in neural network training.

3. Linear Mode Connectivity

A striking recent discovery is Linear Mode Connectivity (LMC): independently trained networks can often be connected by low-loss linear paths in parameter space (Frankle et al., ICML 2020), which can be improved if we account for permutation symmetries (Ainsworth et al., ICML 2022).

Figure from (Ainsworth et al., ICML 2022) .

Re-basin Methods: Algorithms like Git Re-Basin (Ainsworth et al., ICML 2022) align independently trained networks by finding the optimal permutation, revealing that what appear to be separate minima are often the same minimum under permutation symmetry.

Recent Extensions: (Singh et al., 2024) provide a "mountainside and ridge" perspective on when LMC occurs, showing that the loss landscape topography determines connectivity patterns. (Zhou et al., NeurIPS 2024) extend this to show Layerwise Linear Feature Connectivity (LLFC): not just the loss but also the internal features interpolate linearly!

Additional Symmetries in Modern Architectures

Beyond permutation symmetry, modern architectures exhibit additional invariances:

Scale Invariance

For ReLU networks, rescaling weights in one layer and inversely rescaling in the next preserves the function: $$ W^{(l)} \to \alpha W^{(l)}, \quad W^{(l+1)} \to \frac{1}{\alpha} W^{(l+1)} $$

Orthogonal Symmetries in Transformers

(Ashkboos et al., 2024) discovered that Transformers with LayerNorm or RMSNorm exhibit orthogonal symmetry along the residual path, a continuous symmetry group much richer than discrete permutations. This has implications for understanding Transformer loss landscapes and enables new model merging techniques.

Why Weight Space Symmetries Matter

Understanding weight space symmetries is crucial for:

1. Bayesian Neural Networks: Symmetries create multimodal posteriors that challenge variational inference (Gelberg et al., 2024)
2. Model Merging: Techniques like model soups and task arithmetic depend on aligning models across symmetry equivalence classes
3. Understanding Generalization: The structure of equivalent minima and their connections may explain implicit regularization in SGD
4. Neural Functional Networks: Designing networks that output functions (other networks' weights) requires respecting weight space symmetries (Navon et al., NeurIPS 2023), (Zhou et al., NeurIPS 2024)

Small Models: Want a smal (sparse) neural net? Train a large model then prune it

Making a model large, with a lot of paramters, allows SGD to more easily find "flat" solutions in the loss manifold. However, once a model is learned, it can be significantly compressed into a much smaller model.

These ideas (prune or compress after learning) have found their way into state-of-the-art image recognition performance with no loss of accuracy in (Han et al., NeurIPS 2015):

• The method prunes redundant connections using a three-step approach.

1. First, we train the network to learn which connections are important.
2. Next, we prune the unimportant connections.
3. Finally, we will retrain the network to fine tune the weights of the remaining connections. Note that this requires to train the original (large model). Similarly, the Optimal Brain Damage (Le Cun et al., NIPS 1990) and Optimal Brain Surgeon (Hassibi and Stork, NeurIPS 1993) prune networks to reduce the number of connections based on the Hessian of the loss function and suggest that such pruning is more accurate than magnitude-based pruning such as weight decay (weigth decay is just $L_2$ regularization).

Quantization and approximations are also used, for instance, (Denton et al., NeurIPS 2014) exploits the linear structure of the neural network by finding an appropriate low-rank approximation of the parameters and keeping the accuracy within 1% of the original model.

Model sparsification

• Unsurprisingly, large flat-regions have many nearly identical models.
• Among these models, there are very sparse models
  • A sparse model is a model whose weight parameters are mostly zero
  • A trained (optimized) neural network with parameters $\theta_\text{final}$ often has a simpler neural network $M'$ that closely matches the output of $M$ and with significantly fewer parameters $\theta_\text{final}'$
  • The process of obtaining this sparser network is often called pruning

Compressing Neural Networks

There are a few ways to compress neural network models. Often, more than one approach is combined:

Method	Action	Works in
Parameter pruning and sharing	Remove redundant parameters (same accuracy)	Convolutional and fully connected layers
Low-rank factorization	Using matrix decomposition to represent parameters	Convolutional and fully connected layers
Knowledge distillation	Training a compact neural network with distilled knowledge from a large model	Convolutional and fully connected layers
Quantization	Quantize the neural network to allow parameter sharing	Convolutional and fully connected layers

Parameter Pruning

Network pruning consists of removing some parameters (weights) of a pretrained network $M$ (by setting them to zero) and fine tuning the remaining parameters.

• The original dense network is transformed into a sparse one preserving the most important connections.
  • A simple criteria for parameter selection is by how much setting parameter to zero affects the output
    • Often, will zero the parameters that are close to zero
  • The new model is tuned (further optimized for a few gradient steps) by fixing the removed weights to 0.

• The pruning process is applied multiple times until the remaining weights have cannot be pruned further without affecting performance (how $M'$ mimics $M$)
  • Han et al. (2015) reduces the number of parameters of VGG16 (Simonyan et al. (2014)) by 1/13.

Pruning behavior

Here is an illustration of the effectiveness of pruning:

• Seems wasteful. We trained a large neural network to end up with a small one.

• Can we directly train the small one using the same optimization procedure?

Unfortunately, no (Frankle and Carbin, ICLR 2019):

Effect of preserving initialization of pruned network

(Frankle and Carbin, ICLR 2019) discovered that by preserving the original weight initializations $\theta_0$ from the unpruned network, you can train a network with the topology of the pruned network and achieve the same or better test accuracy within the same number of gradient steps.

• The reasons for this behavior are not yet fully clear
  • But there are some hypotheses

Lottery Ticket Hypothesis

• Assume our goal is to find a winning lottery ticket.

  • To increase our chances we will buy a billion tickets
  • Very likely one or more tickets are winning tickets

• Same goes with overparameterized neural networks.

  • Winning the lottery is equivalent to training a neural network to find a large flat-region that has high accuracy on our task.
  • Flat-regions are large and have sparse networks in them
  • There are many such large flat-regions, hence, we are likely to find one in the optimization

LoRA: Low-Rank Adaptation

LoRA Not only model weights are sparse, they also tend to be low rank. (Hu et al., ICLR 2022) uses the low rank nature of the model weights to constrain weight updates to a low-rank structure.

Core Idea

For a pre-trained weight matrix $W_0 \in \mathbb{R}^{d \times k}$, instead of updating the full matrix, LoRA represents the update as a low-rank decomposition: $$ W = W_0 + \Delta W = W_0 + BA $$ where $B \in \mathbb{R}^{d \times r}$ and $A \in \mathbb{R}^{r \times k}$ with $r \ll \min(d, k)$. $W_0$ is considered a fixed input, not learnable parameters.

Key Properties

Property	Description
Parameter Reduction	For GPT-3 175B, reduces trainable parameters by 10,000$\times$
Memory Efficiency	Reduces GPU memory requirement by 3$\times$
No Inference Latency	Merged weights $W_0 + BA$ can be computed once after training
Task Switching	Different LoRA adapters can be swapped without reloading the base model

Initialization: $A$ is initialized with random Gaussian entries, $B$ is initialized to zero, so $\Delta W = 0$ at the start of training.

Why does this work? The key insight from (Aghajanyan et al., 2021) is that fine-tuning updates have low "intrinsic rank": the weight changes needed for task adaptation concentrate in a low-dimensional subspace.

Other Tricks of the Trade

Keep derivatives under control.

Construct a model whose derivatives are easy to calculate and easy to use for optimization.

• Make judicious use of softmax and tanh activations in hidden layers, avoid too small or too large gradients.

• Gradient Clipping. Gradient norm clipping (Pascanu et al., 2013) uses the following algorithm, for some norm $\Vert \cdot \Vert$,
  1. $\hat{g} \leftarrow \nabla_{\boldsymbol \theta} L(\text{Data};{\boldsymbol \theta})$
  2. If $\Vert \hat{g} \Vert \geq \text{threshold}$
     Then $$\hat{g} \leftarrow \frac{\text{threshold}}{\Vert \hat{g} \Vert} \hat{g}.$$

• Batch Normalization (Ioffe and Szegedy, ICML 2015).
  Batch normalization reduces the amount that hidden unit values can vary within a minibatch. The primary reason is believed to be "covariance shift" but there isn't yet enough evidence to support this claim (see (Santurkar et al., ArXiv 2018) for some arguments why it cannot be covariate shift).

Regardless of the reason, Batch Normalization significantly help control vanishing and exploding gradients.

• Consider a mini-batch of $m$ training examples

• Let $$\mathcal{B} = \{{\bf h}_1^{(l-1)},\dots,{\bf h}^{(l-1)}_m\}$$ be the inputs to layer $(l)$ after we forwardpass the mini-batch.

  • Note that ${\bf h}_1^{(l-1)}$ is a vector containing the post-activation values of the neurons that are inputs to layer $l$.
  • For simplicity we will drop the superscript $^{(l-1)}$.
  • We want to replace the inputs $B$ with the batch-normalized values $\{{\bf h}'_1,\ldots,{\bf h}'_m\}$ using the batch normalization transform that we will soon describe.
  • Batch normalization is applied to each layer independently.

• Batch normalization has two parameter vectors per layer: parameter vectors $\boldsymbol{\gamma}^{(l-1)}$ and $\boldsymbol{\beta}^{(l-1)}$.

  • The insight: these parameters allow the neural network to undo the normalization.
  • They are optimized just like any other parameters: using gradient descent using derivatives computed via backpropagation

We refer to the following procedure $\text{BN}_{\boldsymbol{\gamma},\boldsymbol{\beta}} : ({\bf h}_1,\ldots,{\bf h}_m) \to ({\bf h}'_1,\ldots,{\bf h}_m)$ as the Batch Normalizing Transform.

The algorithm is as follows:

1. Step 1: $$\boldsymbol{\mu}_\mathcal{B} \leftarrow \frac{1}{m} \sum_{i=1}^m {\bf h}_i$$
2. Step 2: $$\boldsymbol{\sigma}^2_\mathcal{B} \leftarrow \frac{1}{m} \sum_{i=1}^m ({\bf h}_i - \boldsymbol{\mu}_\mathcal{B})^2$$
3. Step 3: $$\hat{\bf h}_i \leftarrow \frac{{\bf h}_i - \boldsymbol{\mu}_\mathcal{B}}{\boldsymbol{\sigma}_\mathcal{B} + \epsilon}, \qquad i=1,\ldots,m$$
4. Step 4: $${\bf h}'_i \leftarrow \boldsymbol{\gamma} \odot \hat{\bf h}_i + \boldsymbol{\beta}, \qquad i=1,\ldots,m$$

# Batch normalization code
def batchnorm_forward(h, gamma, beta, eps=1e-4):

    # x is the data (batch)
    # N will be the number of examples
    # D is the dimension of the data
    N, D = h.shape

    # Step 1: calculate mean of batch
    mu = np.mean(h, axis = 0)

    # center the data by subtracting mean vector of every trainings example
    h_mu = h - mu

    # calculate the centralized squares
    sq = h_mu ** 2

    # Step 2: calculate variance
    var = np.sum(sq, axis = 0)/N

    #  add eps for numerical stability, then sqrt
    sqrtvar = np.sqrt(var + eps)

    # Step 3: execute normalization
    h_hat = h_mu/sqrtvar

    # Step 4: Nor the two transformation steps
    output_value = gamma * h_hat + beta

    return output_value

A Note on Parameter Regularization for Overparameterized Neural Networks

Reducing overfitting is not our main goal.

• Thus, parameter regularization may or may not be relevant in for your model.
  • We often find that $L_1$ regularized models, where the penalty $\lambda \Vert {\boldsymbol \theta}\Vert_1$ is added to the negative log-likelihood does not lead to significantly sparser models (unlike on Logistic Regression, where it has a huge effect).
  • Similarly, $L_2$ regularization, where the penalty $\lambda \Vert {\boldsymbol \theta}\Vert_2^2$ is added to the negative log-likelihood also may not lead to models that generalize better.
  • There is a clean Bayesian justification for such a regularization term, where the regularization is defined as $-\log(p({\boldsymbol \theta}))$ and $p({\boldsymbol \theta})$ is a prior probability over the parameters ${\boldsymbol \theta}$.
  • Less clear why the Bayesian prior should improve our ability to find Flat-regions

Parallel GPU Data Transfers

First, important notes:

• Host to GPU copies are much faster when they originate from pinned (page-locked) memory.
  • In pytorch's DataLoader you can just set pin_memory=True
• Make the largest batch size you can make (keep in mind the corresponding (larger) learning rate)
  • You can also perform multiple backward passes before one optimization step (without calling zero_grad(), your gradients accumulate)

Running I/O operations in parallel may mitigate the costly I/O speed bottleneck

import threading
import queue
# Producer function that places data on the Queue and pushes it to the GPU memory asynchronously 

def producer(id,device,my_queue,loader):
    #print(f'obj={id} device={device} queue={my_queue} loader={loader}')
    # Place our data on the Queue
    for batchIdx, batchTrainData in enumerate(loader):
        # non_blocking=True makes the operation asynchronous in Pytorch 
        # Any computation using batchTrainData[0], batchTrainData[1]
        batchTrainData[0] = batchTrainData[0].to(device, non_blocking=True)
        batchTrainData[1] = batchTrainData[1].to(device, non_blocking=True)
        # If queue is full, it will block
        my_queue.put( (batchIdx, batchTrainData[0], batchTrainData[1]) , block=True)

    # Tell consumer we are done
    my_queue.put((batchIdx,None,None))

# Function that trains the model
def train_model(model, criterion, optimizer, num_epochs):

    for epoch in range(num_epochs):
        print(f'Epoch {epoch + 1}/{num_epochs}')
        print('-' * 10)

        # Start producer at each epoch
        my_queue = queue.Queue(10) # Careful with GPU memory. This Queue(10) requires 10x more GPU memory than the traditional approach
        dataproducer = threading.Thread(target=producer,args=(0,device,my_queue,data_loaders['train']))
        dataproducer.start()
        
        while True:
            batchIdx, inputs, labels = my_queue.get(block=True)
            if inputs == None:
                break
            
            # zero the gradients of "g"
            optimizer.zero_grad()

            outputs = model(inputs)
            
            loss = criterion(outputs, labels) # loss negative log-likelihood loss using the output "probabilities" 
                                                # actually, the CrossEntropyLoss() does
                                                # the softmax for us
                    
            loss.backward() # Compute gradient
            optimizer.step() # One optimization step

            _, preds = torch.max(outputs, 1) # predicted label (we will use later)

    return model

Multi-GPU Optimization in Pytorch (Data Parallelelism)

Pytorch code example: https://pytorch.org/tutorials/beginner/former_torchies/parallelism_tutorial.html

Maximum Likelihood Estimation of Large Datasets

• Let $D = \{x_1, \ldots, x_n\}$ be a set of observations in a dataset, $x_i \in \mathbb{R}^d$.
• Note that $D$ can be seen as a matrix with $n$ rows and $d$ columns.
• Assume the data D examples are independently sampled from the same distribution: $P(X|{\bf W})$.
• The maximum likelihood estimate of $D$ via gradient descent (over negative log-likelihood) is: $$ \hat{{\bf W}} = \arg\min_{\bf W} -\log P(x_1,\ldots,x_n | {\bf W})= \arg\min_{\bf W} -\log \prod_{i=1}^n P(x_i|{\bf W}) = \arg\min_{\bf W} \sum_{i=1}^n -\log P(x_i|{\bf W}) $$

Traditional Gradient Descent:

To find $\hat{{\bf W}}$, we will use gradient descent as a greedy search algorithm.

Let gradient(-log(P(x[i]|theta))) be $\frac{\partial}{\partial \theta} -\log P(x_i|\theta) = -\frac{1}{P(x_i|\theta)} \frac{\partial}{\partial \theta} P(x_i|\theta)$:

eta = <small constant>
epochs = <some constant greater than one>

W = random.Normal(0,I)

for epoch in 1 to epochs:
    total_gradient = 0
    for i in 1 to n:
        total_gradient += gradient(-log(P(x[i]|W)))
    W = W - eta * total_gradient

Multi-GPU, Multi-CPU: Data Parallel Training of a Machine Learning Models

Multi-GPU data parallel training of deep learningh models is similar to map-reduce procedure on Spark

• First step is to scatter the data (distribute the slices of data to each worker CPU or GPU)

• Second step is to broadcast the most up-to-date model

• Third step is the map, where each GPU worker node computes the gradient of its data slice and writes the output (gradient) to its GPU memory

• Fourth step is the reduce step, where we add the gradients

• Final step is to update the model

Step 1: Scatter

Scatter distributes the data to each (corresponding) worker GPU (or each CPU worker, or each machine worker)

• Must pay attention to GPU-CPU bandwidth

CPU-GPU Bandwidth in detail

• Current Intel CPUs only have 48 PCIe lanes = 3x PCIe x16 running at full speed (PCIe x16 $\approx$ 11Gbps each)
• With 8x GPUs and 2x high-speed network cards, two GPUs must share a single PCIe x16 to the CPU
  • At full load, this means 5.5Gbps for each GPU

Scatter in Distributed Gradient Computation

The gradient computations can be performed on multiple machines (distributed)

• In a platform like Spark and Hadoop, the data is likely already scattered
• Note that network communication is an even bigger bottleneck than any PCIe bottleneck

Step 2: Broadcast

Broadcast distributes the model parameters ${\bf W}$ to all workers

• Still must pay attention to
  • GPU-CPU bandwidth (if computing on GPUs)
  • CPU-CPU bandwidth (if computing on multiple CPUs)
  • Machine-machine network bandwidth (if computing on multiple machines)

Step 3: Map: Local Gradient Computation

Map computes the gradient of model ${\bf W}$ on training data example $X_i$ on GPU worker $i$

Step 4: Reduce: Compute Gradient Sum

Reduce computes the sum of the gradients of model ${\bf W}$ on training data example $[X_1,\ldots,X_4]$

Step 5: Update Model

Update model uses the final gradient to update the model ${\bf W}$ using learning rate $\eta$

Model Parallelism

Model parallelism happens when we break our model (neural network) into smaller models that live on different GPUs

• Useful when the model does not fit in a single GPU
• Models must be constructed with multiple GPUs in mind (not convered in this class)

Mixed-Precision Optimization

Quantization for Scalable Machine Learning

Partial credit to Christopher De Sa and NVidia

Comparison:

• Scientific computing models: 64-bit (double-precision) floating point arithmetic
• Traditional models: 32-bit (single-precision) floating point arithmetic
• Quantized models: 16bit (half-precision) floating point arithmetic, 8-bit and 4-bit integer arithmetic
  • Quantized models are becoming standard practice in industrial applications of Machine Learning
  • Quantized models also help with the representation compression issue
• Extreme quantization: 1-bit neuron weights? Hubara et al., 2018

How are floating point numbers represented?

For 32-bits:

If the above bitstring is a vector ${\bf b}$:

The represented numbers are given in base 2: $$ \text{represented\_number} = (-1)^\text{sign} 2^{\text{exponent}-126} 1.\text{fractional}(b_{22}b_{21}\cdots b_{0}) ,$$ where fractional is such that if we want to represent number 0.7 we set exponent=$01111110_2$=$126_{10}$ and $b_{22}b_{21}\cdots b_{0} = 0110011001100110011001100110011_2$ such that the number is represented as

because $$ 0.7 = 1/2 + 0*(1/4) + 1*(1/8) + 1*(1/16) + 0 * (1/32) + 0*(1/64) + 1*(1/128) + 1*(1/256) + 0*(1/512) + 0*(1/1024) + 1*(1/2048) + 1*(1/4096) + 0*(1/2^{13}) + 0*(1/2^{14}) + 1*(1/2^{15}) + 1*(1/2^{16}) + \ldots $$

Example

number = 0.1 + 0.2

print(f"Result {number}")

Result 0.30000000000000004

number = 1/2  + 1/8 + 1/16 + 1/128 + 1/256 + 1/2048 + 1/4096\
            + 1/(2**15) + 1/(2**16) + 1/(2**19) + 1/(2**20) + 1/(2**23) + 1/(2**24) 

print(f'Floating point 32 approximation of 0.7 is the number = {number}')

Floating point 32 approximation of 0.7 is the number = 0.699999988079071

# Exact floating point representation of fractions is an illusion:
number = 1 + 1/3
number *= 3

print(f'Floating point 32 approximation pretending 1/3 is not an approximation, number = {number}')

number = 1 + 1/3
number += 0.001
number *= 3
number -= 0.003

print(f'Floating point 32 can no longer pretend 1/3 is not an approximation, number = {number}')

Floating point 32 approximation pretending 1/3 is not an approximation, number = 4.0
Floating point 32 can no longer pretend 1/3 is not an approximation, number = 3.999999999999999

Single-precision (FP32) Special Cases

Above is how floating point numbers are represented, except for three special cases:

1. When the exponent = 0, which some hardware will also set the mantissa to zero to avoid numerical problems
   • denormal numbers are numbers with exponent = 0 and mantissa non-zero. They are sometimes treated specially to try to extend their range.
2. If exponent = 255, and
   1. The mantissa is zero. Then sign will determine if the number is $-\infty$ or $+\infty$.
   2. The mantissa is nonzero. Then the floating point value represents something that is not a number, called a NaN value.

Floating point operations and precision

Floating-point operations are usually equivalent to doing the following.

• Interpret the ordinary floating-point number inputs as real numbers.
• Perform the desired operation on those real numbers, producing a real-number result.
• Find the floating-point number that is closest to that result, and return that floating-point number.

Floating point precision

In practice this is done with a circuit that is equivalent to this procedure, since obviously we can’t actually compute with real numbers. This last rounding step has bounded relative error, such that for some constant $\epsilon_\text{machine}$ (called the machine epsilon precision) $$ |\text{FloatingPointOp}(x, y) - \text{RealNumberOp}(x, y)| \leq |\text{RealNumberOp}(x, y)| \cdot \epsilon_\text{machine}. $$

• The machine epsilon $\epsilon_\text{machine}$ measures the numerical error caused by the floating point format. For single-precision floats, $\epsilon_\text{machine} \approx 1.2 × 10^{−7}$.

Floating point operations (Special Cases)

There are a few special cases of floating point operations:

• (overflow) If the real-number result of the operation is larger in magnitude than the largest non-infinite representable floating point number (about $\pm 3.4 \times 10^{38}$ for single-precision floats), then the operation returns positive or negative infinity, as appropriate.

• (underflow) If the real-number result of the operation is smaller in magnitude than 1/2 the smallest non-zero representable floating point number (about $\pm 1.4 \times 10{−45}$ for single-precision floats) then the result of the operation is zero.

• (NaN) If the operation does not make sense, such as division by zero or something like ($+\infty$) − ($+\infty$), then the result is NaN.

Other Floating Point Types

What is the cost of 32-bit floating-point computation?

• Need to store 32 bits in memory for each number used in the algorithm.
• Need to have specialized hardware to support the 32-bit FP computation.

New floating-point types reduce storage and computation costs (source Nvidia):

Industry shift to half-precision (FP16), TF32 & BFLOAT16:

NVidia GPU speeds (source Anandtech):

New Trend: Half-precision (FP16) speed-ups in hardware

Matrix Multiply and Accumulate is a common operation in neural networks.

• Matrix Multiply and Accumulate (MMA) is now supported in hardware

• Specialized hardware is much less complex if matrix multiplication in FP16

Hardware-based Matrix Multiply and Accumulate (FP16 $\times$ FP16 + FP32)

NVidia marketing material:

• NVidia Volta, Turing, Ampere GPUs can perform Matrix Multiply and Accumulate in a single clock cycle.

Machine Learning motivation for FP16

• Machine learning computations are already noisy

  • Due to random sampling of examples in SGD
  • Due to noise in data and labels (from the real world)

• As long as the numerical error from a reduced-bit-count representation is small compared with the noise already observable in SGD, we can expect performance impact to be small

Half-precision

For 16-bit floating point numbers (half-precision, FP16):

• IEEE 754 standard defines the following 16-bit half-precision floating point format: 1 sign bit, 5 exponent bits, and 10 fractional bits.

$$ \text{represented_number} = (-1)^\text{sign} 2^{\text{exponent}-15} 1.\text{fractional}(b_{9}b_{8}\cdots b_{0}) $$

In comparison to single-precision floating point numbers, half-precision floats have:

• a larger machine epsilon (meaning more numerical errors due to rounding), $\epsilon_\text{machine} \approx 9.8 \times 10^{-4}$
• a smaller overflow threshold (overflows more common) of about $6.5 \times 10^4$
• a larger underflow threshold of about $6.0 \times 10^{−8}$

IMPORTANT: As long as our optimization stay far away from these thresholds, half-precision floating point numbers will be OK in approximating real numbers.

Pros of low-precision computing:

• Can fit more numbers (and therefore more training examples, activations, etc.) in memory
• Can store more numbers (and therefore larger models) in the cache
• Can transmit more numbers per second
• Can compute faster by extracting more parallelism in hardware
• Fater training = uses less energy (GPU cards can consume up to 350W each!)

Cons of low-precision computing:

• Limits the range of numbers we can represent (easy to overflow and underflow)
• Needs specialized support from the hardware (more expensive hardware)
• Introduces quantization error (FP32 to FP16 will introduce quantization effects)

In FP16, We Must Monitor Gradients

During model development, always verify if gradients are using FP16 well.

Half precision format leads to the following dynamic range and precision:

• Normalized values
  • $2^{-14}$ to $2^{15}$, 10+1 bits of significand (including sign)
• Denormal values
  • $2^{-24}$ to $2^{-15}$, significand bits decrease as the exponent gets smaller. Exponent k in [-24, -15] range results in (25 - k) bits of significand precision.
• Some example magnitudes:
  • Maximum normalized: 65,504
  • Minimum normalized: $2^{-14} \approx \text{6.10e-5}$
  • Minimum denormal: $2^{-24} \approx \text{5.96e-8}$

Image source: NVidia

• Much of the FP16 representable range above can be left unused by the gradient values.
  • Therefore, if we shift the gradient values to occupy more of that range, we can preserve many values that are otherwise lost to 0s.
  • This shift is called scaling.

Mixed-Precision Arithmetic

In mixed-precision arithmetic we will:

• Maintain a master copy of neural network parameters ${\bf W}$ in FP32

• For each iteration:

  • Make an FP16 copy of the parameters (this copy may negate all benefits)
  • Multiply the loss by a scaling factor $\alpha$
  • Compute the gradient $\frac{\partial}{\partial {\bf W}} \text{loss}$ in FP16
  • Multiply the weight gradient by $1/\alpha$ when converting it back to FP32
  • If the gradient did not suffer overflow
    • Complete the weight update in FP32
    • If suffered overflow, ignore update
      • Will cause gradient bias

(Future) TF32 & BFLOAT16 Arithmetic

• Moving models from FP32 to TF32 or BFLOAT16 may actually give huge benefits over FP16

• No mixed-precision neeed. No master copy of neural network parameters

  • Saves space (fewer bits to represent parameters)
  • Saves FP32 $\to$ FP16 conversion time

• Nvidia Ampere GPUs (TF32): PyTorch 1.7+ has a flag called allow_tf32 which defaults to true.

  • This flag controls whether PyTorch is allowed to use the TensorFloat32 (TF32) tensor cores, available on new NVIDIA GPUs since Ampere, internally to compute matmul (matrix multiplies and batched matrix multiplies) and convolutions.

# The flag below controls whether to allow TF32 on matmul. This flag defaults to True.
torch.backends.cuda.matmul.allow_tf32 = True

# The flag below controls whether to allow TF32 on cuDNN. This flag defaults to True.
torch.backends.cudnn.allow_tf32 = True

Pytorch's Automatic Mixed Precisions (AMP)

Python's automatic mixed precision training is a new library that simplified mixed precision training of neural networks

• This example REQUIRES modern Nvidia GPUs

• We will do the following steps in order:

  • Load and normalizing the CIFAR10 training and test datasets using torchvision
  • Define an MLP
  • Define a loss function and the Stochastic Gradient Descent optimizer
  • Train the network on the training data
  • Test the network on the test data

# You can download statistical_header.py from 
#     https://www.cs.purdue.edu/homes/ribeirob/courses/Spring2026/lectures/05Optimization/statistical_header.py
from statistical_header import *
import torch
import torch.nn.functional as F
import torch.nn as nn
import torch.utils.data as data_utils
from timeit import default_timer as timer

import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)

### NEW for MIXED PRECISION ###
## Mixed Precision Libraries. Needs pytorch 1.6+ and CUDA 10.2+
## It only works on Nvidia GPUs that support CUDA 10.2+ (Volta, Turin, Ampere,... architectures)
from torch.cuda.amp import GradScaler 
from torch.cuda.amp import autocast 


class MLPNet(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(MLPNet, self).__init__()
        """
        In the constructor we instantiate three nn.Linear modules and assign them as
        member variables.

        first parameter: input dimension
        second parameter: dimension of the output
        """
        self.fc1 = nn.Linear(input_size, hidden_size) 
        self.fc2 = nn.Linear(hidden_size, hidden_size) 
        self.fc3 = nn.Linear(hidden_size, 2)  

    def forward(self, x):
        # First fully connected layer,  out = out^T W_1 + b_1
        out = self.fc1(x)
        # Activation function (sigmoid, relu, tanh, ...)
        out = torch.relu(out)
        # Second fully connected layer out = out^T W_2 + b_2
        out = self.fc2(out)
        # Activation function (sigmoid, relu, tanh, ...)
        out = torch.relu(out)
        # Third fully connected layer out = out^T W_3 + b_3
        out = self.fc3(out)
        # Final activation function (log of the sigmoid)
        out = F.log_softmax(out, dim=1)
        return out

train_file = "Bank_Data_Train.csv"
validation_file = "Bank_Data_Validation.csv"

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

data_train, target_train = data_processing(train_file)
data_validation, target_validation = data_processing(validation_file)

# replace categorical strings with one-hot encoding and add a small amount of Gaussian noise so it follows Gaussian model assumption

data_train, data_validation = add_categorical(train=data_train,validation=data_validation,feature_str='FICO Range')
data_train, data_validation = add_categorical(train=data_train,validation=data_validation,feature_str='Loan Purpose')

# Create the neural network classifier object in GPU or CPU memory (depending on device)
# Creates the model in the standard FP32 precision
model = MLPNet(input_size=data_train.shape[1],hidden_size=10000).to(device=device)

# Loss and optimizer
# nn.NLLLoss() is the negative log likelihood. It takes the log_softmax as input.
#  ... if using nn.CrossEntropyLoss(), it will compute softmax internally.
criterion = nn.NLLLoss()
# Optimizer is the standard stochastic gradient descent
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)  


# Decide if we are using mixed precision
mixed_precision = False



# Convert data from pandas DataFrame to Pytorch's Tensor data and move to "device"
if mixed_precision:
    input_batch = torch.from_numpy(np.array(data_train).astype('float16')).to(device=device)
else:
    input_batch = torch.from_numpy(np.array(data_train).astype('float32')).to(device=device)
    
labels_batch = torch.from_numpy(np.array(target_train).astype('int64')).to(device=device)

train_tensor = data_utils.TensorDataset(input_batch, labels_batch) 
train_loader = data_utils.DataLoader(dataset = train_tensor, batch_size = 50, shuffle = True) 
                                  # No need for pin_memory=True in DataLoader since data already on device

if mixed_precision:
    # Creates a GradScaler once at the beginning of training.
    scaler = GradScaler()
    print(f"Using GPU mixed precision arithmetic: FP16 & FP32 using device {device}")
else:
    print(f"Using FP32 arithmetic using device {device}")

start = timer()

# Loop training over epochs. 
for epoch in range(30):
    print(".",end="")
    for inputs, labels in train_loader:
        # Need to zero the gradient for this batch, since Pytorch gradients accumulate
        optimizer.zero_grad()

        if mixed_precision:
            # Runs the forward pass with autocasting.
            with autocast():
                # Get model outputs
                outputs = model(inputs)
                # Compute loss
                loss = criterion(outputs, labels)
        else:
            # Get model outputs
            outputs = model(inputs)
            # Compute loss
            loss = criterion(outputs, labels)


        # Backpropagation and optimize
        if mixed_precision:
            # Scales loss.  Calls backward() on scaled loss to create scaled gradients.
            # Backward passes under autocast are not recommended.
            # Backward ops run in the same dtype autocast chose for corresponding forward ops.
            scaler.scale(loss).backward()
        else:
            loss.backward()

        if mixed_precision:
            # scaler.step() first unscales the gradients of the optimizer's assigned params.
            # If these gradients do not contain infs or NaNs, optimizer.step() is then called,
            # otherwise, optimizer.step() is skipped.
            scaler.step(optimizer)
        else:
            optimizer.step()

        if mixed_precision:
            # Updates the scale for next iteration.
            scaler.update()
        else:
            # Does nothing... FP32 optimization does generally need to scale gradients
            pass            
            
end = timer()
print(f'\nOptimization time {end - start}s') 

Using FP32 arithmetic using device cuda
..............................
Optimization time 7.61641539700031s

Appendix

The exemplar code below shows the interpolation of initial parameters and final parameters for a Convolutional Neural Network

• While we have not yet seen Convolutional Neural Networks, the following example shows that the flatness of the loss manifold is not a characteristic of MLPs only.
• The majority of the architectures have the flatness observed by (Goodfellow et al., ICLR 2015)

import matplotlib.pyplot as plt
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.fc1 = nn.Linear(320, 100)
        self.fc2 = nn.Linear(100, 10)

    def forward(self, x):
        # 2D convolutional layer with max pooling layer and reLU activations
        x = F.relu(F.max_pool2d(self.conv1(x), 2)) 
        # 2nd layer of 2D convolutional layer with max pooling layer and reLU activations
        x = F.relu(F.max_pool2d(self.conv2(x), 2))
        # rearrange output from a 2d array (matrix) to a vector
        x = x.view(-1, 320)
        # fully connected layer with ReLU activation
        x = F.relu(self.fc1(x)) 
        # fully connected layer
        x = self.fc2(x)
        # log(softmax(x)) activation
        return F.log_softmax(x, dim=1)
    

def train(model, device, train_loader, epoch, printout=False):
    # Initializes model for training
    model.train()
    
    for batch_idx, (data, target) in enumerate(train_loader):
        # move data from CPU to GPU if GPU is available
        data, target = data.to(device), target.to(device) 
        
        # Forward pass. Gets output *before* softmax
        output = model(data)
        
        # Computes the negative log-likelihood over the training data target. 
        #     log(softmax(.)) avoids computing the normalization factor. 
        #     Note that target does not need to be 1-hot encoded (pytorch will 1-hot encode for you)
        loss = F.nll_loss(output, target)
        
        # **Very important,** need to zero the gradient buffer before we can use the computed gradients
        model.zero_grad()     # zero the gradient buffers of all parameters

        if printout:
            print('conv1.bias.grad before backward')
            print(model.conv1.bias.grad)

        # Backpropagate loss for all paramters over all examples
        loss.backward()

        if printout:
            print('conv1.bias.grad after backward')
            print(model.conv1.bias.grad)        
        
        # Performs one step of SGD with a fixed learning rate (not a Robbins-Monro procedure)
        learning_rate = 0.05

        if printout:
            print("")
            print("conv1.bias before SGD step")
            print(model.conv1.bias)

        # iterate over all model parameters
        for f in model.parameters():
            # why subtract? Are we minimizing or maximizing?
            f.data.sub_(f.grad.data * learning_rate)

        if printout:
            print("conv1.bias after SGD step")
            print(model.conv1.bias)
        
        if printout:
            print("")
            print("... next gradient step")
        
        if batch_idx % 100 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))

def validation(model, device, validation_loader, quiet = False):
    
    # Some operations (such as Dropout and BatchNorm) will behave differently between training and testing.
    #   In our model it does not matter, but it is generally good practice to explain what you intend to do to pytorch
    model.eval()
    
    # Initialize local variables
    validation_loss = 0
    correct = 0
    
    # no_grad() tells pytorch not to worry about keeping intermediate resutls around as we are not looking to 
    #   backpropagate this forward pass
    with torch.no_grad():
        for data, target in validation_loader:
            data, target = data.to(device), target.to(device)
            
            # Forward pass
            output = model(data)
            
            # sums up the loss over the entire batch
            validation_loss += F.nll_loss(output, target, reduction='sum').item()
            
            # get the index of the max log-probability
            #    why does getting the max of the log-probability suffices for prediction? Don't we need to compute the softmax probability?
            pred = output.max(1, keepdim=True)[1] 
            
            # Computes accuracy
            correct += pred.eq(target.view_as(pred)).sum().item()

    validation_loss /= len(validation_loader.dataset)
    if not quiet: 
        print('\nValidation set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
            validation_loss, correct, len(validation_loader.dataset),
            100. * correct / len(validation_loader.dataset)))
    
    return (100. * correct / len(validation_loader.dataset))

# Train batch size
train_batch_size = 32

# How many times we will go over the entire dataset (keeps same number of gradient steps)
no_epochs = 10

# Validation batch size (as large as GPU can support)
validation_batch_size = 8192

# Do we have GPUs available?py
use_cuda = torch.cuda.is_available()
print(f'Is cuda available? {use_cuda}')

# Handy way to use GPUs if there are any
device = torch.device("cuda" if use_cuda else "cpu")

# Data loader parameters
kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}

# Data loader for training data (will download MNIST data from the Web if data is not cached in the system)
train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('../../data', train=True, download=True,
                   transform=transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,))
                   ])),
                      batch_size=train_batch_size, shuffle=True, **kwargs)

# Data loader for test data (will download MNIST data from the Web if data is not cached in the system)
validation_loader = torch.utils.data.DataLoader(
    datasets.MNIST('../../data', train=False, transform=transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,))
                   ])),
                      batch_size=validation_batch_size, shuffle=True, **kwargs)

# Instantiates the model. If devide="cuda", the model parameters will be created in the GPU memory.
model = Net().to(device)


# Copy initial model parameters
paramsinitial = model.named_parameters()
initial_params = {}
for name, param in paramsinitial:
    initial_params[name] = param.data.clone()

##
## Train Model
##
for epoch in range(1, no_epochs + 1):
    # Train for one epoch 
    train(model, device, train_loader, epoch,printout=False)
    # Validation of model accuracy so far 
    validation(model, device, validation_loader)
##
## End Train Model
##


# Save final model parameters
paramsfinal = model.named_parameters()
final_params = {}
for name, param in paramsfinal:
    final_params[name] = param.data.clone()

print("The End.")

Is cuda available? True

100.0%
100.0%
100.0%
100.0%

Train Epoch: 1 [0/60000 (0%)]	Loss: 2.341128
Train Epoch: 1 [3200/60000 (5%)]	Loss: 0.394493
Train Epoch: 1 [6400/60000 (11%)]	Loss: 0.544606
Train Epoch: 1 [9600/60000 (16%)]	Loss: 0.118881
Train Epoch: 1 [12800/60000 (21%)]	Loss: 0.175817
Train Epoch: 1 [16000/60000 (27%)]	Loss: 0.276667
Train Epoch: 1 [19200/60000 (32%)]	Loss: 0.118975
Train Epoch: 1 [22400/60000 (37%)]	Loss: 0.032347
Train Epoch: 1 [25600/60000 (43%)]	Loss: 0.106753
Train Epoch: 1 [28800/60000 (48%)]	Loss: 0.025911
Train Epoch: 1 [32000/60000 (53%)]	Loss: 0.083347
Train Epoch: 1 [35200/60000 (59%)]	Loss: 0.006960
Train Epoch: 1 [38400/60000 (64%)]	Loss: 0.367526
Train Epoch: 1 [41600/60000 (69%)]	Loss: 0.005039
Train Epoch: 1 [44800/60000 (75%)]	Loss: 0.370557
Train Epoch: 1 [48000/60000 (80%)]	Loss: 0.042653
Train Epoch: 1 [51200/60000 (85%)]	Loss: 0.037402
Train Epoch: 1 [54400/60000 (91%)]	Loss: 0.014302
Train Epoch: 1 [57600/60000 (96%)]	Loss: 0.063911

Validation set: Average loss: 0.0752, Accuracy: 9762/10000 (98%)

Train Epoch: 2 [0/60000 (0%)]	Loss: 0.055846
Train Epoch: 2 [3200/60000 (5%)]	Loss: 0.200267
Train Epoch: 2 [6400/60000 (11%)]	Loss: 0.000455
Train Epoch: 2 [9600/60000 (16%)]	Loss: 0.035182
Train Epoch: 2 [12800/60000 (21%)]	Loss: 0.151589
Train Epoch: 2 [16000/60000 (27%)]	Loss: 0.009371
Train Epoch: 2 [19200/60000 (32%)]	Loss: 0.037208
Train Epoch: 2 [22400/60000 (37%)]	Loss: 0.116903
Train Epoch: 2 [25600/60000 (43%)]	Loss: 0.085956
Train Epoch: 2 [28800/60000 (48%)]	Loss: 0.138026
Train Epoch: 2 [32000/60000 (53%)]	Loss: 0.071198
Train Epoch: 2 [35200/60000 (59%)]	Loss: 0.009628
Train Epoch: 2 [38400/60000 (64%)]	Loss: 0.074091
Train Epoch: 2 [41600/60000 (69%)]	Loss: 0.076068
Train Epoch: 2 [44800/60000 (75%)]	Loss: 0.005357
Train Epoch: 2 [48000/60000 (80%)]	Loss: 0.007064
Train Epoch: 2 [51200/60000 (85%)]	Loss: 0.049709
Train Epoch: 2 [54400/60000 (91%)]	Loss: 0.030474
Train Epoch: 2 [57600/60000 (96%)]	Loss: 0.030441

Validation set: Average loss: 0.0459, Accuracy: 9849/10000 (98%)

Train Epoch: 3 [0/60000 (0%)]	Loss: 0.014080
Train Epoch: 3 [3200/60000 (5%)]	Loss: 0.069974
Train Epoch: 3 [6400/60000 (11%)]	Loss: 0.098497
Train Epoch: 3 [9600/60000 (16%)]	Loss: 0.068064
Train Epoch: 3 [12800/60000 (21%)]	Loss: 0.050928
Train Epoch: 3 [16000/60000 (27%)]	Loss: 0.103381
Train Epoch: 3 [19200/60000 (32%)]	Loss: 0.003418
Train Epoch: 3 [22400/60000 (37%)]	Loss: 0.036549
Train Epoch: 3 [25600/60000 (43%)]	Loss: 0.001735
Train Epoch: 3 [28800/60000 (48%)]	Loss: 0.046959
Train Epoch: 3 [32000/60000 (53%)]	Loss: 0.025560
Train Epoch: 3 [35200/60000 (59%)]	Loss: 0.002048
Train Epoch: 3 [38400/60000 (64%)]	Loss: 0.111647
Train Epoch: 3 [41600/60000 (69%)]	Loss: 0.028877
Train Epoch: 3 [44800/60000 (75%)]	Loss: 0.003346
Train Epoch: 3 [48000/60000 (80%)]	Loss: 0.119741
Train Epoch: 3 [51200/60000 (85%)]	Loss: 0.002772
Train Epoch: 3 [54400/60000 (91%)]	Loss: 0.013800
Train Epoch: 3 [57600/60000 (96%)]	Loss: 0.024684

Validation set: Average loss: 0.0388, Accuracy: 9889/10000 (99%)

Train Epoch: 4 [0/60000 (0%)]	Loss: 0.014513
Train Epoch: 4 [3200/60000 (5%)]	Loss: 0.025604
Train Epoch: 4 [6400/60000 (11%)]	Loss: 0.006497
Train Epoch: 4 [9600/60000 (16%)]	Loss: 0.002374
Train Epoch: 4 [12800/60000 (21%)]	Loss: 0.022659
Train Epoch: 4 [16000/60000 (27%)]	Loss: 0.006287
Train Epoch: 4 [19200/60000 (32%)]	Loss: 0.021101
Train Epoch: 4 [22400/60000 (37%)]	Loss: 0.067158
Train Epoch: 4 [25600/60000 (43%)]	Loss: 0.015008
Train Epoch: 4 [28800/60000 (48%)]	Loss: 0.147569
Train Epoch: 4 [32000/60000 (53%)]	Loss: 0.014231
Train Epoch: 4 [35200/60000 (59%)]	Loss: 0.003611
Train Epoch: 4 [38400/60000 (64%)]	Loss: 0.008196
Train Epoch: 4 [41600/60000 (69%)]	Loss: 0.029074
Train Epoch: 4 [44800/60000 (75%)]	Loss: 0.117044
Train Epoch: 4 [48000/60000 (80%)]	Loss: 0.123418
Train Epoch: 4 [51200/60000 (85%)]	Loss: 0.037239
Train Epoch: 4 [54400/60000 (91%)]	Loss: 0.204570
Train Epoch: 4 [57600/60000 (96%)]	Loss: 0.010661

Validation set: Average loss: 0.0363, Accuracy: 9881/10000 (99%)

Train Epoch: 5 [0/60000 (0%)]	Loss: 0.002227
Train Epoch: 5 [3200/60000 (5%)]	Loss: 0.008240
Train Epoch: 5 [6400/60000 (11%)]	Loss: 0.004109
Train Epoch: 5 [9600/60000 (16%)]	Loss: 0.005075
Train Epoch: 5 [12800/60000 (21%)]	Loss: 0.009389
Train Epoch: 5 [16000/60000 (27%)]	Loss: 0.159250
Train Epoch: 5 [19200/60000 (32%)]	Loss: 0.052160
Train Epoch: 5 [22400/60000 (37%)]	Loss: 0.030101
Train Epoch: 5 [25600/60000 (43%)]	Loss: 0.012411
Train Epoch: 5 [28800/60000 (48%)]	Loss: 0.013003
Train Epoch: 5 [32000/60000 (53%)]	Loss: 0.003172
Train Epoch: 5 [35200/60000 (59%)]	Loss: 0.015160
Train Epoch: 5 [38400/60000 (64%)]	Loss: 0.008633
Train Epoch: 5 [41600/60000 (69%)]	Loss: 0.002586
Train Epoch: 5 [44800/60000 (75%)]	Loss: 0.005363
Train Epoch: 5 [48000/60000 (80%)]	Loss: 0.006879
Train Epoch: 5 [51200/60000 (85%)]	Loss: 0.000544
Train Epoch: 5 [54400/60000 (91%)]	Loss: 0.001782
Train Epoch: 5 [57600/60000 (96%)]	Loss: 0.024111

Validation set: Average loss: 0.0344, Accuracy: 9889/10000 (99%)

Train Epoch: 6 [0/60000 (0%)]	Loss: 0.053980
Train Epoch: 6 [3200/60000 (5%)]	Loss: 0.001280
Train Epoch: 6 [6400/60000 (11%)]	Loss: 0.000244
Train Epoch: 6 [9600/60000 (16%)]	Loss: 0.003609
Train Epoch: 6 [12800/60000 (21%)]	Loss: 0.020801
Train Epoch: 6 [16000/60000 (27%)]	Loss: 0.179235
Train Epoch: 6 [19200/60000 (32%)]	Loss: 0.001885
Train Epoch: 6 [22400/60000 (37%)]	Loss: 0.005691
Train Epoch: 6 [25600/60000 (43%)]	Loss: 0.000702
Train Epoch: 6 [28800/60000 (48%)]	Loss: 0.000815
Train Epoch: 6 [32000/60000 (53%)]	Loss: 0.000060
Train Epoch: 6 [35200/60000 (59%)]	Loss: 0.005580
Train Epoch: 6 [38400/60000 (64%)]	Loss: 0.000957
Train Epoch: 6 [41600/60000 (69%)]	Loss: 0.015330
Train Epoch: 6 [44800/60000 (75%)]	Loss: 0.005295
Train Epoch: 6 [48000/60000 (80%)]	Loss: 0.024392
Train Epoch: 6 [51200/60000 (85%)]	Loss: 0.000298
Train Epoch: 6 [54400/60000 (91%)]	Loss: 0.004544
Train Epoch: 6 [57600/60000 (96%)]	Loss: 0.310109

Validation set: Average loss: 0.0369, Accuracy: 9887/10000 (99%)

Train Epoch: 7 [0/60000 (0%)]	Loss: 0.001513
Train Epoch: 7 [3200/60000 (5%)]	Loss: 0.000615
Train Epoch: 7 [6400/60000 (11%)]	Loss: 0.000089
Train Epoch: 7 [9600/60000 (16%)]	Loss: 0.000304
Train Epoch: 7 [12800/60000 (21%)]	Loss: 0.000211
Train Epoch: 7 [16000/60000 (27%)]	Loss: 0.000789
Train Epoch: 7 [19200/60000 (32%)]	Loss: 0.000775
Train Epoch: 7 [22400/60000 (37%)]	Loss: 0.005375
Train Epoch: 7 [25600/60000 (43%)]	Loss: 0.000204
Train Epoch: 7 [28800/60000 (48%)]	Loss: 0.000847
Train Epoch: 7 [32000/60000 (53%)]	Loss: 0.000616
Train Epoch: 7 [35200/60000 (59%)]	Loss: 0.072994
Train Epoch: 7 [38400/60000 (64%)]	Loss: 0.001726
Train Epoch: 7 [41600/60000 (69%)]	Loss: 0.000518
Train Epoch: 7 [44800/60000 (75%)]	Loss: 0.007243
Train Epoch: 7 [48000/60000 (80%)]	Loss: 0.062051
Train Epoch: 7 [51200/60000 (85%)]	Loss: 0.000205
Train Epoch: 7 [54400/60000 (91%)]	Loss: 0.018500
Train Epoch: 7 [57600/60000 (96%)]	Loss: 0.026744

Validation set: Average loss: 0.0314, Accuracy: 9899/10000 (99%)

Train Epoch: 8 [0/60000 (0%)]	Loss: 0.000584
Train Epoch: 8 [3200/60000 (5%)]	Loss: 0.004015
Train Epoch: 8 [6400/60000 (11%)]	Loss: 0.004231
Train Epoch: 8 [9600/60000 (16%)]	Loss: 0.003181
Train Epoch: 8 [12800/60000 (21%)]	Loss: 0.003478
Train Epoch: 8 [16000/60000 (27%)]	Loss: 0.000253
Train Epoch: 8 [19200/60000 (32%)]	Loss: 0.016351
Train Epoch: 8 [22400/60000 (37%)]	Loss: 0.000854
Train Epoch: 8 [25600/60000 (43%)]	Loss: 0.004954
Train Epoch: 8 [28800/60000 (48%)]	Loss: 0.001034
Train Epoch: 8 [32000/60000 (53%)]	Loss: 0.012457
Train Epoch: 8 [35200/60000 (59%)]	Loss: 0.000117
Train Epoch: 8 [38400/60000 (64%)]	Loss: 0.034235
Train Epoch: 8 [41600/60000 (69%)]	Loss: 0.000424
Train Epoch: 8 [44800/60000 (75%)]	Loss: 0.007374
Train Epoch: 8 [48000/60000 (80%)]	Loss: 0.000778
Train Epoch: 8 [51200/60000 (85%)]	Loss: 0.000863
Train Epoch: 8 [54400/60000 (91%)]	Loss: 0.003457
Train Epoch: 8 [57600/60000 (96%)]	Loss: 0.000038

Validation set: Average loss: 0.0324, Accuracy: 9906/10000 (99%)

Train Epoch: 9 [0/60000 (0%)]	Loss: 0.000718
Train Epoch: 9 [3200/60000 (5%)]	Loss: 0.001542
Train Epoch: 9 [6400/60000 (11%)]	Loss: 0.009722
Train Epoch: 9 [9600/60000 (16%)]	Loss: 0.002582
Train Epoch: 9 [12800/60000 (21%)]	Loss: 0.000063
Train Epoch: 9 [16000/60000 (27%)]	Loss: 0.004710
Train Epoch: 9 [19200/60000 (32%)]	Loss: 0.000513
Train Epoch: 9 [22400/60000 (37%)]	Loss: 0.000668
Train Epoch: 9 [25600/60000 (43%)]	Loss: 0.000115
Train Epoch: 9 [28800/60000 (48%)]	Loss: 0.005287
Train Epoch: 9 [32000/60000 (53%)]	Loss: 0.001097
Train Epoch: 9 [35200/60000 (59%)]	Loss: 0.018604
Train Epoch: 9 [38400/60000 (64%)]	Loss: 0.030686
Train Epoch: 9 [41600/60000 (69%)]	Loss: 0.000104
Train Epoch: 9 [44800/60000 (75%)]	Loss: 0.004387
Train Epoch: 9 [48000/60000 (80%)]	Loss: 0.002285
Train Epoch: 9 [51200/60000 (85%)]	Loss: 0.019901
Train Epoch: 9 [54400/60000 (91%)]	Loss: 0.003204
Train Epoch: 9 [57600/60000 (96%)]	Loss: 0.000679

Validation set: Average loss: 0.0332, Accuracy: 9909/10000 (99%)

Train Epoch: 10 [0/60000 (0%)]	Loss: 0.084954
Train Epoch: 10 [3200/60000 (5%)]	Loss: 0.000407
Train Epoch: 10 [6400/60000 (11%)]	Loss: 0.000224
Train Epoch: 10 [9600/60000 (16%)]	Loss: 0.042110
Train Epoch: 10 [12800/60000 (21%)]	Loss: 0.000902
Train Epoch: 10 [16000/60000 (27%)]	Loss: 0.000451
Train Epoch: 10 [19200/60000 (32%)]	Loss: 0.010534
Train Epoch: 10 [22400/60000 (37%)]	Loss: 0.000252
Train Epoch: 10 [25600/60000 (43%)]	Loss: 0.000281
Train Epoch: 10 [28800/60000 (48%)]	Loss: 0.004495
Train Epoch: 10 [32000/60000 (53%)]	Loss: 0.005983
Train Epoch: 10 [35200/60000 (59%)]	Loss: 0.065250
Train Epoch: 10 [38400/60000 (64%)]	Loss: 0.000037
Train Epoch: 10 [41600/60000 (69%)]	Loss: 0.003491
Train Epoch: 10 [44800/60000 (75%)]	Loss: 0.000108
Train Epoch: 10 [48000/60000 (80%)]	Loss: 0.000362
Train Epoch: 10 [51200/60000 (85%)]	Loss: 0.002694
Train Epoch: 10 [54400/60000 (91%)]	Loss: 0.009262
Train Epoch: 10 [57600/60000 (96%)]	Loss: 0.001658

Validation set: Average loss: 0.0320, Accuracy: 9907/10000 (99%)

The End.

%matplotlib inline
import matplotlib.pyplot as plt

v_err = []
alpha_v = []

# Interpolate model parameters
for alpha in torch.linspace(0, 1.5, steps=10): 
    alpha = alpha.to(device)
    for name, param in model.named_parameters():
        param.data = (1. - alpha)*initial_params[name].data + alpha *final_params[name].data
    v_err.append(100. - validation(model, device, validation_loader,quiet=True))
    alpha_v.append(alpha.cpu())
    print(f' alpha = {alpha} has validation error {v_err[-1]}%')

plt.rcParams.update({'font.size': 22})
plt.xlabel("Alpha")
plt.ylabel("Validation Error (%)")
plt.semilogy(alpha_v,v_err,linestyle='-', marker='o', color='b')
plt.show()

 alpha = 0.0 has validation error 91.41%
 alpha = 0.1666666716337204 has validation error 26.060000000000002%
 alpha = 0.3333333432674408 has validation error 3.1299999999999955%
 alpha = 0.5 has validation error 1.4099999999999966%
 alpha = 0.6666666865348816 has validation error 1.0900000000000034%
 alpha = 0.8333333134651184 has validation error 0.9500000000000028%
 alpha = 1.0 has validation error 0.9300000000000068%
 alpha = 1.1666666269302368 has validation error 0.9599999999999937%
 alpha = 1.3333333730697632 has validation error 0.9699999999999989%
 alpha = 1.5 has validation error 1.0%