PyTorch Loss Landscape Probes

A loss landscape probe is a walk along the straight line between two trained solutions. The scores along the line show whether it stays in one valley or climbs a ridge.

The blend, weight by weight

Alpha blends every tensor in the state_dict with the same weight. Alpha 0 is the first solution and alpha 1 is the second.

half = {name: 0.5 * first[name] + 0.5 * second[name] for name in first}
print(f"First solution, one weight: {float(first['0.weight'][0, 0]):.4f}")
print(f"Halfway blend, same weight: {float(half['0.weight'][0, 0]):.4f}")
print(f"Second solution, one weight: {float(second['0.weight'][0, 0]):.4f}")

Both endpoints are trained solutions

The first two losses on the line belong to the two solutions. Both are small, because both models finished training.

print(f"Loss at alpha 0.0: {losses[0]:.4f}")
print(f"Loss at alpha 1.0: {losses[-1]:.4f}")

A sparkline from raw numbers

The sparkline maps each loss to one shade character. The smallest loss takes the first character and the largest takes the last.

values = [0.1, 0.5, 0.9]
low, high = min(values), max(values)
print("".join(shades[min(8, int((v - low) / (high - low) * 8.99))] for v in values))

Example

import torch
from torch import nn


def make_blobs(count, seed):
    g = torch.Generator().manual_seed(seed)
    first = torch.randn(count // 2, 2, generator=g) * 0.7 + torch.tensor([-1.0, -1.0])
    second = torch.randn(count // 2, 2, generator=g) * 0.7 + torch.tensor([1.0, 1.0])
    x = torch.cat([first, second])
    y = torch.cat([torch.zeros(count // 2), torch.ones(count // 2)]).long()
    perm = torch.randperm(count, generator=g)
    return x[perm], y[perm]


train_x, train_y = make_blobs(160, seed=3)
loss_fn = nn.CrossEntropyLoss()


def train_solution(seed):
    torch.manual_seed(seed)
    model = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 2))
    optimizer = torch.optim.Adam(model.parameters(), lr=0.05)
    for step in range(80):
        optimizer.zero_grad()
        loss = loss_fn(model(train_x), train_y)
        loss.backward()
        optimizer.step()
    return {name: value.detach().clone()
            for name, value in model.state_dict().items()}


first = train_solution(0)
second = train_solution(1)

probe = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 2))
losses = []
for i in range(11):
    alpha = i / 10
    probe.load_state_dict({name: (1 - alpha) * first[name] + alpha * second[name]
                           for name in first})
    probe.eval()
    with torch.no_grad():
        losses.append(float(loss_fn(probe(train_x), train_y)))

print("Alpha  Loss")
for i, value in enumerate(losses):
    print(f"{i / 10:.1f}    {value:.4f}")

low, high = min(losses), max(losses)
shades = " .:-=+*#@"
line = "".join(shades[min(8, int((value - low) / (high - low) * 8.99))]
               for value in losses)
print("Loss along the line:", line)

Key points

Solution: one full set of trained weights, saved as a state_dict
Alpha: the blend weight along the line, running from 0 at one solution to 1 at the other
Each probe point loads blended weights into a fresh model and scores the training loss
A flat or dipping line means the two solutions share one valley
A hump in the middle means a ridge separates the two solutions
The spark-line maps the eleven losses to shade characters, smallest to largest

In this example, the same small MLP trains twice on the same blob data from two seeds. Both runs end at low loss, 0.0048 for the first solution and 0.0064 for the second. The probe then blends the two state_dicts at eleven alpha values. The loss climbs steadily from either end toward the middle. At alpha 0.5 it peaks at 0.0309, more than six times the first endpoint. The spark-line draws that hill in one row of characters. A ridge stands between these two solutions, so they sit in different valleys. Averaging these two weight sets would land on the ridge, not in a valley.

Note

The probe walks one straight line only, and a curved path between the solutions could still dip lower.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import torch
from torch import nn


def make_blobs(count, seed):
    g = torch.Generator().manual_seed(seed)
    first = torch.randn(count // 2, 2, generator=g) * 0.7 + torch.tensor([-1.0, -1.0])
    second = torch.randn(count // 2, 2, generator=g) * 0.7 + torch.tensor([1.0, 1.0])
    x = torch.cat([first, second])
    y = torch.cat([torch.zeros(count // 2), torch.ones(count // 2)]).long()
    perm = torch.randperm(count, generator=g)
    return x[perm], y[perm]


train_x, train_y = make_blobs(160, seed=3)
loss_fn = nn.CrossEntropyLoss()


def train_solution(seed):
    torch.manual_seed(seed)
    model = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 2))
    optimizer = torch.optim.Adam(model.parameters(), lr=0.05)
    for step in range(80):
        optimizer.zero_grad()
        loss = loss_fn(model(train_x), train_y)
        loss.backward()
        optimizer.step()
    return {name: value.detach().clone()
            for name, value in model.state_dict().items()}


first = train_solution(0)
second = train_solution(1)

probe = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 2))
losses = []
for i in range(11):
    alpha = i / 10
    probe.load_state_dict({name: (1 - alpha) * first[name] + alpha * second[name]
                           for name in first})
    probe.eval()
    with torch.no_grad():
        losses.append(float(loss_fn(probe(train_x), train_y)))

print("Alpha  Loss")
for i, value in enumerate(losses):
    print(f"{i / 10:.1f}    {value:.4f}")

low, high = min(losses), max(losses)
shades = " .:-=+*#@"
line = "".join(shades[min(8, int((value - low) / (high - low) * 8.99))]
               for value in losses)
print("Loss along the line:", line)

Exercise

What does the hump in the middle of the line prove?

A ridge separates the two solutions The two solutions are identical The second solution never trained The probe blended the wrong tensors