Visualizing Loss Functions as Topological Landscape Profiles

TL;DR

A new representation based on topological data analysis that enables the visualization of higher-dimensional loss landscapes is introduced and it is found that the topology of the loss landscape is simpler for better-performing models; and the shape of loss landscapes near transitions from low to high model performance is observed.

Abstract

In machine learning, a loss function measures the difference between model predictions and ground-truth (or target) values. For neural network models, visualizing how this loss changes as model parameters are varied can provide insights into the local structure of the so-called loss landscape (e.g., smoothness) as well as global properties of the underlying model (e.g., generalization performance). While various methods for visualizing the loss landscape have been proposed, many approaches limit sampling to just one or two directions, ignoring potentially relevant information in this extremely high-dimensional space. This paper introduces a new representation based on topological data analysis that enables the visualization of higher-dimensional loss landscapes. After describing this new topological landscape profile representation, we show how the shape of loss landscapes can reveal new details about model performance and learning dynamics, highlighting several use cases, including image segmentation (e.g., UNet) and scientific machine learning (e.g., physics-informed neural networks). Through these examples, we provide new insights into how loss landscapes vary across distinct hyperparameter spaces: we find that the topology of the loss landscape is simpler for better-performing models; and we observe greater variation in the shape of loss landscapes near transitions from low to high model performance.

Figures (6)

• Figure 1: Our topological landscape profiles enable the visualization of higher-dimensional loss landscapes by capturing their underlying shape (or topology). Here we show loss landscapes based on the top $n$ Hessian eigenvectors. See Section \ref{['sec:methods']} for details.
• Figure 2: Representing the merge tree as a topological landscape profile. In (A) we show a single basin corresponding to a merge tree with a single branch, and in (B) we show multiple basins corresponding to multiple branches. In (C) we color the basins based on their average loss.
• Figure 3: Analyzing the loss function of a physics-informed neural network (PINN) trained to solve simple physical convection problems. See Section \ref{['sec:experiments_results_pinn']} for details.
• Figure 4: Loss landscapes over training for UNet models with a CRF layer trained on the Oxford-IIIT Pet dataset. See Section \ref{['sec:experiments_results_unet']} for details.
• Figure 5: Comparing topological landscape profiles based on (A) three-dimensional and (B) four-dimensional loss landscapes. See Section \ref{['sec:experiments_results_pinn']} for details.