Package: loss-landscapes
loss-landscapesloss-landscapesis a PyTorch library for approximating neural network loss functions, and other related metrics,
in low-dimensional subspaces of the model's parameter space. The library makes the production of visualizations
such as those seen inVisualizing the Loss Landscape of Neural Netsmuch
easier, aiding the analysis of the geometry of neural network loss landscapes.This library does not provide plotting facilities, letting the user define how the data should be plotted. Other
deep learning frameworks are not supported, though a TensorFlow version,loss-landscapes-tf, is planned for
a future release.NOTE: this library is in early development. Bugs are virtually a certainty, and the API is volatile. Do not use
this library in production code. For prototyping and research, always use the newest version of the library.1. What is a Loss Landscape?LetL : Parameters -> Real Numbersbe a loss function, which maps a point in the model parameter space to a
real number. For a neural network withnparameters, the loss functionLtakes ann-dimensional input. We
can define the loss landscape as the set of alln+1-dimensional points(param, L(param)), for all pointsparamin the parameter space. For example, the image below, reproduced from the paper by Li et al (2018), link
above, provides a visual representation of what a loss function over a two-dimensional parameter space might look
like:Of course, real machine learning models have a number of parameters much greater than 2, so the parameter space of
the model is virtually never two-dimensional. Because we can't print visualizations in more than two dimensions,
we cannot hope to visualize the "true" shape of the loss landscape. Instead, a number of techniques
exist for reducing the parameter space to one or two dimensions, ranging from dimensionality reduction techniques
like PCA, to restricting ourselves to a particular subspace of the overall parameter space. For more details,
read Li et al's paper.2. Base Example: Supervised Loss in Parameter SubspacesThe simplest use case forloss-landscapesis to estimate the value of a supervised loss function in a subspace
of a neural network's parameter space. The subspace in question may be a point, a line, or a plane (these subspaces
can be meaningfully visualized). Suppose the user has trained a supervised learning model, of typetorch.nn.Module,
on a dataset consisting of samplesXand labelsy, by minimizing some loss function. The user now wishes to
produce a surface plot alike to the one in section 1.This is accomplished as follows:metric=Loss(loss_function,X,y)landscape=random_plane(model,metric,normalize='filter')As seen in the example above, the two core concepts inloss-landscapesaremetricsandparameter subspaces. The
latter define the section of parameter space to be considered, while the former define what quantity is evaluated at
each considered point in parameter space, and how it is computed. In the example above, we define aLossmetric
over dataXand labelsy, and instructloss_landscapeto evaluate it in a randomly generated planar subspace.This would return a 2-dimensional array of loss values, which the user can plot in any desirable way. Example
visualizations the user might use this type of data for are shown below.Check theexamplesdirectory forjupyternotebooks with more in-depth examples of what is possible.3. Metrics and Custom MetricsTheloss-landscapeslibrary can compute any quantity of interest at a collection of points in a parameter subspace,
not just loss. This is accomplished using aMetric: a callable object which applies a pre-determined function,
such as a cross entropy loss with a specific set of inputs and outputs, to the model. Theloss_landscapes.model_metricspackage contains a number of metrics that cover common use cases, such asLoss(evaluates a loss
function),LossGradient(evaluates the gradient of the loss w.r.t. the model parameters),PrincipalCurvatureEvaluator(evaluates the principal curvatures of the loss function), and more.Furthermore, the user can add custom metrics by subclassingMetric. As an example, consider the library
implementation ofLoss, fortorchmodels:classMetric(abc.ABC):""" A quantity that can be computed given a model or an agent. """def__init__(self):super().__init__()@abc.abstractmethoddef__call__(self,model_wrapper:ModelWrapper):passclassLoss(Metric):""" Computes a specified loss function over specified input-output pairs. """def__init__(self,loss_fn,inputs:torch.Tensor,target:torch.Tensor):super().__init__()self.loss_fn=loss_fnself.inputs=inputsself.target=targetdef__call__(self,model_wrapper:ModelWrapper)->float:returnself.loss_fn(model_wrapper.forward(self.inputs),self.target).item()The user may create customMetrics in a similar manner. One complication is that theMetricclass'__call__method is designed to take as input aModelWrapperrather than a model. This class is internal
to the library and exists to facilitate the handling of the myriad of different models a user may pass as
inputs to a function such asloss_landscapes.planar_interpolation(). It is sufficient for the user to know
that aModelWrapperis a callable object that can be used to call the model on a given input (see thecall_fnargument of theModelInterfaceclass in the next section). The class also provides aget_model()method
that exposes a reference to the underlying model, should the user wish to carry out more complicated operations
on it.In summary, theMetricabstraction adds a great degree of flexibility. An metric defines what quantity
dependent on model parameters the user is interested in evaluating, and how to evaluate it. The user could define,
for example, a metric that computes an estimate of the expected return of a reinforcement learning agent.4. More Complex ModelsIn the general case of a simple supervised learning model, as in the sections above, client code calls functions
such asloss_landscapes.linear_interpolationand passes as argument a PyTorch module of typetorch.nn.Module.For more complex cases, such as when the user wants to evaluate the loss landscape as a function of a subset of
the model parameters, or the expected return landscape for a RL agent, the user must specify to theloss-landscapeslibrary how to interface with the model (or the agent, on a more general level). This is accomplished using aModelWrapperobject, which hides the implementation details of the model or agent. For general use, the library
supplies theGeneralModelWrapperin theloss_landscapes.model_interface.model_wrappermodule.Assume the user wishes to estimate the expected return of some RL agent which provides anagent.act(observation)method for action selection. Then, the example from section 2 becomes as follows:metric=ExpectedReturnMetric(env,n_samples)agent_wrapper=GeneralModelWrapper(agent,[agent.q_function,agent.policy],lambdaagent,x:agent.act(x))landscape=random_plane(agent_wrapper,metric,normalize='filter')5. WIP: Connecting Paths, Saddle Points, and Trajectory TrackingA number of features are currently under development, but as of yet incomplete.A number of papers in recent years have shown that loss landscapes of neural networks are dominated by a
proliferation of saddle points, that good solutions are better described as large low-loss plateaus than as
"well-bottom" points, and that for sufficiently high-dimensional networks, a low-loss path in parameter space can
be found between almost any arbitrary pair of minima. In the future, theloss-landscapeslibrary will feature
implementations of algorithms for finding such low-loss connecting paths in the loss landscape, as well as tools to
facilitate the study of saddle points.Some sort of trajectory tracking features are also under consideration, though at the time it's unclear what this
should actually mean, as the optimization trajectory is implicitly tracked by the user's training loop. Any metric
along the optimization trajectory can be tracked with libraries such asignitefor PyTorch.6. Support for Other DL LibrariesTheloss-landscapeslibrary was initially designed to be agnostic to the DL framework in use. However, with the
increasing number of use cases to cover it became obvious that maintaining the original library-agnostic design
was adding too much complexity to the code.A TensorFlow version,loss-landscapes-tf, is planned for the future.7. Installation and UseThe package is available on PyPI. Install usingpip install loss-landscapes. To use the library, import as follows:importloss_landscapesimportloss_landscapes.metrics