marcellodebernardi/loss-landscapes

What is loss‑landscapes?

loss‑landscapes is a small Python library built on top of PyTorch that helps researchers and engineers sample and evaluate any scalar quantity (loss, gradient norm, curvature, expected return, etc.) on low‑dimensional slices of a neural network’s huge parameter space. By providing the raw numbers on a line, plane or other sub‑space, it makes it easy to draw the classic loss‑landscape visualisations that show how a model’s loss changes as you move its weights.

---

Core ideas

Concept	What it does
Metric	A callable that, given a wrapped model/agent, returns a single float (e.g., loss, gradient norm, curvature, expected return). The library ships with common metrics like Loss, LossGradient, PrincipalCurvatureEvaluator and lets you subclass Metric for anything else.
Parameter sub‑space	A low‑dimensional region of the full weight space (a point, a line, a random plane, etc.). Functions such as random_plane or linear_interpolation generate these sub‑spaces and evaluate a metric at a grid of points inside them.
ModelWrapper	An internal adapter that hides how to call a model or an RL agent. It presents a uniform forward‑like interface so the same metric code works for plain torch.nn.Modules and for more exotic agents.

---

Typical workflow (Python‑style)

import torch, loss_landscapes as ll

# 1️⃣  Train a model (any torch.nn.Module) …
model = MyNet()
# … training code …

# 2️⃣  Choose a metric – here the standard supervised loss
metric = ll.metrics.Loss(loss_fn=torch.nn.CrossEntropyLoss(),
                         inputs=X_train, target=y_train)

# 3️⃣  Pick a sub‑space – a random 2‑D plane around the current weights
plane = ll.random_plane(model, metric, normalize="filter")

# 4️⃣  `plane` is a 2‑D NumPy / torch array of loss values.
#    Plot it however you like (matplotlib, plotly, etc.)

The library does not provide plotting utilities; you feed the returned array into your favourite visualisation tool.

---

Why use it?

• Re‑useable metric abstraction – evaluate loss, gradients, curvature, or custom quantities without rewriting boiler‑plate code.
• Supports complex agents – via ModelWrapper you can explore expected‑return landscapes for reinforcement‑learning agents, not just supervised losses.
• Normalization options – e.g., normalize="filter" rescales directions to avoid misleading scale differences.
• Open source & extensible – add your own metrics by subclassing Metric.

---

Current status

• The original 2019 release (3.0.6) is broken (missing explicit torch dependency and contains mathematically incorrect normalisation).
• After a long hiatus the repository has been revived; a modern rewrite targeting current Python and PyTorch versions is in progress. The API may change, and the authors advise not to use the published 3.0.6 package until the new version is released.

---

Getting started

pip install loss-landscapes   # pulls the last published wheel (still old)
# For development with the latest code:
uv sync --all-groups          # install deps via uv
uv run pytest                 # run the test suite

Check the examples/ folder for Jupyter notebooks that demonstrate loss‑contour, 3‑D surface, curvature, and RL‑return landscapes.

---

Who might find it useful?

• ML researchers investigating optimisation geometry, flat vs. sharp minima, or the effect of regularisation on the loss surface.
• Practitioners who want to sanity‑check training dynamics by visualising how a model’s loss behaves around the converged weights.
• RL developers curious about how policy parameters affect expected return.
• Educators looking for a lightweight way to generate concrete loss‑landscape plots for lectures.

---

In short: loss‑landscapes is a focused, PyTorch‑native toolkit that turns the abstract notion of a neural‑network loss surface into concrete numeric grids, leaving the actual plotting to you. It is actively being modernised, so keep an eye on the repository for the upcoming stable release.