Visualizing the Invisible: Mastering Loss Landscapes with Filter-Wise Normalization
Visualizing the Invisible: Mastering Loss Landscapes with Filter-Wise Normalization
Ever wondered why two neural networks with the same training loss can perform drastically differently on a test set? The answer often lies in the geometry of the loss landscape.
In deep learning, we often talk about "sharp" vs. "flat" minima. A flat minimum—where the loss remains low even if you nudge the weights slightly—is generally associated with better generalization. But there is a catch: visualizing this landscape is deceptively difficult.
In this post, we dive into a powerful diagnostic framework that solves the "scale-invariance" problem in loss visualization using Filter-Wise Normalization.
The Problem: The Illusion of Flatness
To visualize a loss landscape, we typically take a trained model $\theta^$ and perturb it in a random direction $\delta$: $$L(\alpha) = L(\theta^ + \alpha\delta)$$ By varying $\alpha$, we can plot a 1D slice of the loss. To get a 3D surface, we use two directions, $\delta$ and $\eta$: $$f(\alpha, \beta) = L(\theta^* + \alpha\delta + \beta\eta)$$
The Trap: Neural networks (especially those using ReLU and Batch Normalization) are scale-invariant. If you multiply the weights of a layer by 10 and divide the next layer by 10, the output remains the same.
If you use a standard random Gaussian vector for $\delta$, a model with very large weights will naturally appear "flatter" simply because the perturbation is tiny relative to the weight magnitude. You aren't seeing the intrinsic geometry; you're seeing an artifact of weight scaling.
The Solution: Filter-Wise Normalization
The core contribution of this framework is Filter-Wise Normalization. Instead of using a raw random vector, we rescale the perturbation to match the norm of the actual filters in the trained model.
The Intuition
If a specific filter in your network has a large L2 norm, the perturbation applied to that filter should be proportionally large. This ensures that the "step" we take in the parameter space is meaningful relative to the learned weights.
The Math
For every filter $d_{i,j}$ (the perturbation for layer $i$, filter $j$), we apply the following transformation: $$d_{i,j} \leftarrow \frac{d_{i,j}}{|d_{i,j}|} \cdot |\theta_{i,j}|$$ Where $|\theta_{i,j}|$ is the Frobenius norm of the trained weights for that specific filter.
System Architecture
The following workflow describes how the visualization pipeline transforms a trained model into a geometric map.
Production Implementation (PyTorch)
Below is a complete implementation. The LossLandscapeVisualizer class encapsulates the normalization logic, allowing you to pass any PyTorch model and dataset.
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import make_classification
from torch.utils.data import DataLoader, TensorDataset
class SimpleNet(nn.Module):
def __init__(self, input_dim=20, hidden_dim=64, output_dim=2):
super(SimpleNet, self).__init__()
self.net = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, output_dim)
)
def forward(self, x):
return self.net(x)
class LossLandscapeVisualizer:
def __init__(self, model, criterion, data_loader):
self.model = model
self.criterion = criterion
self.data_loader = data_loader
def _get_filter_normalized_direction(self, direction):
"""
Implements Filter-Wise Normalization:
delta_normalized = delta * (|
|weight|
| / |
|delta|
|)
"""
normalized_dir = {}
with torch.no_grad():
for name, param in self.model.named_parameters():
if 'weight' in name:
weight_norm = torch.norm(param.data)
dir_norm = torch.norm(direction[name])
normalized_dir[name] = direction[name] * (weight_norm / (dir_norm + 1e-8))
else:
normalized_dir[name] = direction[name]
return normalized_dir
def compute_loss(self, params):
self.model.load_state_dict(params)
self.model.eval()
total_loss = 0
with torch.no_grad():
for inputs, targets in self.data_loader:
outputs = self.model(inputs)
loss = self.criterion(outputs, targets)
total_loss += loss.item()
return total_loss / len(self.data_loader)
def plot_2d_landscape(self, resolution=20, range_val=1.0):
theta_star = {n: p.clone() for n, p in self.model.named_parameters()}
delta = {n: torch.randn_like(p) for n, p in self.model.named_parameters()}
eta = {n: torch.randn_like(p) for n, p in self.model.named_parameters()}
delta_norm = self._get_filter_normalized_direction(delta)
eta_norm = self._get_filter_normalized_direction(eta)
x = np.linspace(-range_val, range_val, resolution)
y = np.linspace(-range_val, range_val, resolution)
X, Y = np.meshgrid(x, y)
Z = np.zeros((resolution, resolution))
for i in range(resolution):
for j in range(resolution):
current_params = {}
for name, param in theta_star.items():
current_params[name] = param + X[i, j] * delta_norm[name] + Y[i, j] * eta_norm[name]
Z[i, j] = self.compute_loss(current_params)
self.model.load_state_dict(theta_star)
return X, Y, Z
# --- Execution Block ---
if __name__ == '__main__':
# Data Setup
X_np, y_np = make_classification(n_samples=1000, n_features=20, random_state=42)
loader = DataLoader(TensorDataset(torch.FloatTensor(X_np), torch.LongTensor(y_np)), batch_size=32)
# Model Training
model = SimpleNet()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)
for epoch in range(20):
for inputs, targets in loader:
optimizer.zero_grad()
criterion(model(inputs), targets).backward()
optimizer.step()
# Visualization
visualizer = LossLandscapeVisualizer(model, criterion, loader)
X, Y, Z = visualizer.plot_2d_landscape(resolution=25, range_val=0.5)
fig = plt.figure(figsize=(12, 5))
ax1 = fig.add_subplot(1, 2, 1, projection='3d')
ax1.plot_surface(X, Y, Z, cmap='viridis')
ax1.set_title("Filter-Normalized Landscape")
ax2 = fig.add_subplot(1, 2, 2)
ax2.contourf(X, Y, Z, levels=20, cmap='viridis')
ax2.set_title("Contour Map")
plt.show()Key Takeaways for Practitioners
- Don't Trust Raw Perturbations: If you are comparing two different architectures (e.g., ResNet vs. DenseNet), standard random directions will lie to you. Always use normalization.
- Generalization Proxy: When analyzing your model, look for "wide valleys" in the contour plot. If your minimum is a "sharp needle," your model is likely overfitting and will be sensitive to distribution shift.
- Computational Cost: Note that generating a $25 \times 25$ grid requires 625 full forward passes over your dataset. For very large models, consider using a representative subset of your validation data to speed up the sampling process.
By applying Filter-Wise Normalization, we move from guessing the shape of our model's intelligence to actually seeing the terrain it has learned.