Visualizing Loss Functions as Topological Landscape Profiles

TL;DR

This paper introduces a topological data analysis method to visualize high-dimensional loss landscapes in neural networks, revealing insights into model performance and learning dynamics across different tasks and hyperparameters.

Contribution

It presents a novel topological landscape profile representation for visualizing complex loss landscapes beyond traditional low-dimensional slices.

Findings

Simpler loss landscape topology correlates with better model performance.

Loss landscape shape varies significantly near performance transition points.

The method provides new insights into model generalization and training dynamics.

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.