Visualizing Optimization Trajectory of Neural Nets
By: Logan Yang
Tags: Machine Learning, Deep Learning, Mathematics, Optimization, Data Visualization

In the previous post, I showed some animated plots for the training process of linear regression and logistic regression. Developing a good “feel” of how they “learn” is helpful because they can be used as a baseline before applying more complex models. Although most deep neural networks also use gradient-based learning, similar intuition is much harder to come by. One reason is that the parameters are very high dimensional and there are a lot of nonlinearities involved, it’s very hard to picture in our heads what is going on during the optimization. Unlike computers, we are only programmed to perceive spaces less or equal to 3D.

In this post, I will show a lot more animated plots to offer a glimpse into those high dimensional spaces. The models I use are fully connected multilayer perceptrons with ReLU activations. We can then visually see how width and depth can affect the loss landscape and optimization trajectory.

Let’s consider this 2D dataset of 3 classes. We humans can easily see the pattern at a glance that it’s generated by some sort of spiral function. It takes us less than a second to recognize it and, at the same time, we automatically produce a way of extrapolation for unseen data. It is because we are trained on this sort of visual task since we were born (also don’t forget we can only do this in very low dimensional space). Here, I will train a few artificial neural networks to classify this shape and inspect the training process in the parameter space.

Before we go on, I’d like to raise an important question: how do we visualize the optimization trajectory in the high dimensional parameter space?

The most straightforward way is to find two directions to cut through the high dimensional space and visualize loss values over that plane. But which two directions to use? There are infinitely many potential directions to choose from. In the paper Visualizing Loss Landscape in Neural Nets, Li et. al., the authors discussed some options and adopted one using PCA for dimensionality reduction. Here’s a brief summary of the motivation:

2 random vectors in high dimensional space have a high probability of being orthogonal, and they can hardly capture any variation for the optimization path. The path’s projection onto the plane spanned by the 2 vectors will just look like random walk. If we pick one direction to be the vector pointing from the initial parameters to the final trained parameters, and another direction at random, the visualization will look like a straight line because the second direction doesn’t capture much variance compared to the first. If we use principal component analysis (PCA) on the optimization path and get the top 2 components, we can visualize the loss over the 2 orthogonal directions with the most variance.

Therefore, I am using the PCA approach for better-looking optimization paths. Keep in mind that it is not the “best” approach for path visualization because it may not work for other purposes. For example, if you aim to compare the paths taken by different optimizers, e.g. SGD vs. Adam, this PCA approach won’t work since the principal components come from the paths themselves. The fact that different optimizers have different paths and different PC directions, i.e. a different slice of the loss landscape, makes the comparison impossible. For that purpose, we should use two fixed directions.

The architectures here are produced by varying and combining the two properties below

Number of hidden layers: 1, 5, 10

Number of neurons in each hidden layer: 5, 20, 100

9 configurations in total.

Let’s take a look at the optimization paths alongside the decision areas/boundaries produced by different models. By showing the decision areas instead of validation/test accuracies which rely on a split on the generated dataset, we can get better intuitions about bad fits since we already have our prior — the “spiral” expectation in mind.

First, let me show you what a logistic regression does in this case,