Machine learning / Fundamentals and overview
Exploring Loss Landscapes in Machine Learning
Contents
- Understanding Loss Landscapes
- Mathematical Foundations of Gradient Descent
- Implementing Basic Gradient Descent
- Stochastic Gradient Descent Implementation
- Momentum-based Optimization
- Adaptive Learning Rate Methods
- Implementing Adam Optimizer
- Loss Landscape Visualization Tool
- Escaping Local Minima with Simulated Annealing
- Implementing Grid Search for Global Minimum
- Loss Landscape Analysis Tools
- Real-World Example - Neural Network Training Analysis
- Results Analysis for Optimization Strategies
- Advanced Loss Surface Analysis
- Additional Resources
Understanding Loss Landscapes
The loss landscape represents the error surface over which optimization occurs during neural network training. It's a high-dimensional space where each dimension corresponds to a model parameter, making visualization and interpretation challenging for deep neural networks.
import numpy as np
import matplotlib.pyplot as plt
def loss_landscape_2d(x, y):
# Simple 2D loss landscape example with multiple local minima
return np.sin(4*x) * np.cos(4*y) + 2*(x**2 + y**2)
# Create meshgrid for visualization
x = np.linspace(-2, 2, 100)
y = np.linspace(-2, 2, 100)
X, Y = np.meshgrid(x, y)
Z = loss_landscape_2d(X, Y)
plt.figure(figsize=(10, 8))
plt.contour(X, Y, Z, levels=20)
plt.colorbar(label='Loss')
plt.xlabel('Parameter 1')
plt.ylabel('Parameter 2')
plt.title('2D Loss Landscape with Multiple Local Minima')
plt.show()
Mathematical Foundations of Gradient Descent
Gradient descent forms the backbone of neural network optimization, utilizing partial derivatives to determine the direction of steepest descent in the loss landscape. This mathematical foundation guides parameter updates during training.
def gradient_descent_example():
"""
Mathematical representation of gradient descent:
"""
code = '''
# Gradient Descent Update Rule
$$θ_{t+1} = θ_t - η∇L(θ_t)$$
# Where:
# θ_t: Current parameters
# η: Learning rate
# ∇L(θ_t): Gradient of loss with respect to parameters
# For a simple quadratic loss:
$$L(θ) = (θ - 2)^2$$
# The gradient would be:
$$∇L(θ) = 2(θ - 2)$$
'''
return code
print(gradient_descent_example())
Implementing Basic Gradient Descent
A practical implementation of gradient descent shows how parameters are updated iteratively to minimize the loss function. This example demonstrates the core mechanism using a simple quadratic function as the optimization target.
import numpy as np
def quadratic_loss(theta):
"""Simple quadratic loss function: (θ - 2)²"""
return (theta - 2) ** 2
def gradient(theta):
"""Gradient of the quadratic loss"""
return 2 * (theta - 2)
def gradient_descent(learning_rate=0.1, iterations=50):
theta = 10.0 # Starting point
history = []
for i in range(iterations):
grad = gradient(theta)
theta = theta - learning_rate * grad
loss = quadratic_loss(theta)
history.append((theta, loss))
print(f"Iteration {i}: θ = {theta:.4f}, Loss = {loss:.4f}")
return theta, history
final_theta, history = gradient_descent()
Stochastic Gradient Descent Implementation
import numpy as np
class SGD:
def __init__(self, learning_rate=0.01):
self.learning_rate = learning_rate
def minimize(self, X, y, model, epochs=100, batch_size=32):
n_samples = len(X)
losses = []
for epoch in range(epochs):
# Shuffle data
indices = np.random.permutation(n_samples)
# Mini-batch training
for i in range(0, n_samples, batch_size):
batch_indices = indices[i:i + batch_size]
X_batch = X[batch_indices]
y_batch = y[batch_indices]
# Forward pass
y_pred = model.forward(X_batch)
loss = model.compute_loss(y_pred, y_batch)
# Backward pass
gradients = model.backward()
# Update parameters
for param, grad in gradients.items():
model.parameters[param] -= self.learning_rate * grad
losses.append(loss)
if epoch % 10 == 0:
print(f"Epoch {epoch}, Loss: {loss:.4f}")
return losses
Momentum-based Optimization
Momentum helps accelerate gradient descent by accumulating a velocity vector in directions of persistent reduction in the objective. This implementation shows how momentum can help overcome local minima and speed up convergence.
class MomentumOptimizer:
def __init__(self, learning_rate=0.01, momentum=0.9):
self.learning_rate = learning_rate
self.momentum = momentum
self.velocity = {}
def minimize(self, gradients, parameters):
if not self.velocity:
# Initialize velocity for each parameter
self.velocity = {k: np.zeros_like(v) for k, v in parameters.items()}
for param_name in parameters:
# Update velocity
self.velocity[param_name] = (self.momentum * self.velocity[param_name] +
self.learning_rate * gradients[param_name])
# Update parameters
parameters[param_name] -= self.velocity[param_name]
return parameters
Adaptive Learning Rate Methods
Adaptive learning rate methods dynamically adjust the learning rate for each parameter during training. This implementation demonstrates how RMSprop works by maintaining a moving average of squared gradients to scale parameter updates.
class RMSprop:
def __init__(self, learning_rate=0.001, decay_rate=0.9, epsilon=1e-8):
self.learning_rate = learning_rate
self.decay_rate = decay_rate
self.epsilon = epsilon
self.cache = {}
def minimize(self, gradients, parameters):
if not self.cache:
self.cache = {k: np.zeros_like(v) for k, v in parameters.items()}
for param_name in parameters:
# Update moving average of squared gradients
self.cache[param_name] = (self.decay_rate * self.cache[param_name] +
(1 - self.decay_rate) * np.square(gradients[param_name]))
# Compute parameter update
update = (self.learning_rate * gradients[param_name] /
(np.sqrt(self.cache[param_name]) + self.epsilon))
# Update parameters
parameters[param_name] -= update
return parameters
Implementing Adam Optimizer
Adam combines the benefits of both momentum and RMSprop, utilizing first and second moments of gradients for more efficient optimization. This implementation shows the complete Adam algorithm.
class Adam:
def __init__(self, learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8):
self.learning_rate = learning_rate
self.beta1 = beta1
self.beta2 = beta2
self.epsilon = epsilon
self.m = {} # First moment
self.v = {} # Second moment
self.t = 0 # Timestep
def minimize(self, gradients, parameters):
if not self.m:
self.m = {k: np.zeros_like(v) for k, v in parameters.items()}
self.v = {k: np.zeros_like(v) for k, v in parameters.items()}
self.t += 1
for param_name in parameters:
# Update biased first moment estimate
self.m[param_name] = (self.beta1 * self.m[param_name] +
(1 - self.beta1) * gradients[param_name])
# Update biased second raw moment estimate
self.v[param_name] = (self.beta2 * self.v[param_name] +
(1 - self.beta2) * np.square(gradients[param_name]))
# Compute bias-corrected first moment estimate
m_hat = self.m[param_name] / (1 - self.beta1**self.t)
# Compute bias-corrected second raw moment estimate
v_hat = self.v[param_name] / (1 - self.beta2**self.t)
# Update parameters
parameters[param_name] -= (self.learning_rate * m_hat /
(np.sqrt(v_hat) + self.epsilon))
return parameters
Loss Landscape Visualization Tool
Creating a comprehensive tool for visualizing the loss landscape helps understand optimization dynamics. This implementation allows for 2D and 3D visualization of loss surfaces.
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
class LossLandscapeVisualizer:
def __init__(self, model, loss_fn):
self.model = model
self.loss_fn = loss_fn
def compute_loss_grid(self, param_range, resolution=50):
x = np.linspace(-param_range, param_range, resolution)
y = np.linspace(-param_range, param_range, resolution)
X, Y = np.meshgrid(x, y)
Z = np.zeros_like(X)
for i in range(resolution):
for j in range(resolution):
# Set model parameters
self.model.set_params([X[i,j], Y[i,j]])
# Compute loss
Z[i,j] = self.loss_fn(self.model)
return X, Y, Z
def plot_3d_surface(self, param_range=5.0):
X, Y, Z = self.compute_loss_grid(param_range)
fig = plt.figure(figsize=(12, 8))
ax = fig.add_subplot(111, projection='3d')
surf = ax.plot_surface(X, Y, Z, cmap='viridis')
ax.set_xlabel('Parameter 1')
ax.set_ylabel('Parameter 2')
ax.set_zlabel('Loss')
plt.colorbar(surf)
plt.title('3D Loss Landscape')
plt.show()
Escaping Local Minima with Simulated Annealing
Simulated annealing introduces controlled randomness to escape local minima by occasionally accepting worse solutions. This implementation demonstrates the technique with temperature scheduling.
import numpy as np
class SimulatedAnnealing:
def __init__(self, initial_temp=1000, cooling_rate=0.95, min_temp=1e-10):
self.temp = initial_temp
self.cooling_rate = cooling_rate
self.min_temp = min_temp
def optimize(self, loss_fn, initial_params, n_iterations=1000):
current_params = initial_params
current_loss = loss_fn(current_params)
best_params = current_params
best_loss = current_loss
history = []
for i in range(n_iterations):
# Generate neighbor solution with noise proportional to temperature
neighbor_params = current_params + np.random.normal(0, self.temp, size=current_params.shape)
neighbor_loss = loss_fn(neighbor_params)
# Calculate acceptance probability
delta_loss = neighbor_loss - current_loss
acceptance_prob = np.exp(-delta_loss / self.temp)
# Accept or reject new solution
if delta_loss < 0 or np.random.random() < acceptance_prob:
current_params = neighbor_params
current_loss = neighbor_loss
# Update best solution if necessary
if current_loss < best_loss:
best_params = current_params
best_loss = current_loss
history.append((current_loss, self.temp))
self.temp *= self.cooling_rate
if self.temp < self.min_temp:
break
return best_params, best_loss, history
Implementing Grid Search for Global Minimum
Grid search systematically explores the parameter space to identify potential global minima. This implementation includes parallel processing for efficiency.
import numpy as np
from multiprocessing import Pool
from functools import partial
class GridSearch:
def __init__(self, param_ranges, n_points=50, n_processes=4):
self.param_ranges = param_ranges
self.n_points = n_points
self.n_processes = n_processes
def _evaluate_point(self, point, loss_fn):
return (point, loss_fn(point))
def search(self, loss_fn):
# Generate grid points
grid_points = []
for param_range in self.param_ranges:
points = np.linspace(param_range[0], param_range[1], self.n_points)
grid_points.append(points)
# Create all combinations of parameters
param_combinations = np.array(np.meshgrid(*grid_points)).T.reshape(-1, len(self.param_ranges))
# Parallel evaluation of loss function
with Pool(self.n_processes) as pool:
results = pool.map(partial(self._evaluate_point, loss_fn=loss_fn),
param_combinations)
# Find best parameters
best_params, best_loss = min(results, key=lambda x: x[1])
return {
'best_params': best_params,
'best_loss': best_loss,
'all_results': results
}
Loss Landscape Analysis Tools
Advanced tools for analyzing loss landscape characteristics help identify problematic regions and optimize training strategies. This implementation provides metrics for landscape smoothness and connectivity.
import numpy as np
from scipy.ndimage import gaussian_filter
class LossLandscapeAnalyzer:
def __init__(self, model):
self.model = model
def compute_hessian(self, params, epsilon=1e-5):
n_params = len(params)
hessian = np.zeros((n_params, n_params))
for i in range(n_params):
for j in range(n_params):
# Compute second partial derivatives
h = np.zeros_like(params)
h[i] = epsilon
h[j] = epsilon
f_xy = self.model.loss(params + h)
f_x = self.model.loss(params + np.array([h[i] if k == i else 0 for k in range(n_params)]))
f_y = self.model.loss(params + np.array([h[j] if k == j else 0 for k in range(n_params)]))
f_0 = self.model.loss(params)
hessian[i,j] = (f_xy - f_x - f_y + f_0) / (epsilon * epsilon)
return hessian
def compute_smoothness(self, region, resolution=50):
"""Compute loss landscape smoothness using Gaussian filtering"""
smoothed = gaussian_filter(region, sigma=1.0)
roughness = np.mean(np.abs(region - smoothed))
return roughness
def analyze_critical_points(self, hessian):
"""Analyze critical points using eigenvalues"""
eigenvalues = np.linalg.eigvals(hessian)
is_minimum = np.all(eigenvalues > 0)
is_maximum = np.all(eigenvalues < 0)
is_saddle = not (is_minimum or is_maximum)
return {
'eigenvalues': eigenvalues,
'is_minimum': is_minimum,
'is_maximum': is_maximum,
'is_saddle': is_saddle,
'condition_number': np.max(np.abs(eigenvalues)) / np.min(np.abs(eigenvalues))
}
Real-World Example - Neural Network Training Analysis
This comprehensive example demonstrates how to analyze and optimize neural network training, incorporating multiple techniques to escape local minima and track optimization progress.
import numpy as np
import torch
import torch.nn as nn
class OptimizationAnalyzer:
def __init__(self, model, criterion, optimizer):
self.model = model
self.criterion = criterion
self.optimizer = optimizer
self.history = {'loss': [], 'gradients': [], 'weights': []}
def train_epoch(self, dataloader, analyze=True):
epoch_loss = 0.0
gradient_norms = []
weight_norms = []
for inputs, targets in dataloader:
# Forward pass
self.optimizer.zero_grad()
outputs = self.model(inputs)
loss = self.criterion(outputs, targets)
# Backward pass
loss.backward()
if analyze:
# Store gradient information
grad_norm = torch.norm(
torch.stack([p.grad.norm() for p in self.model.parameters()])
).item()
gradient_norms.append(grad_norm)
# Store weight information
weight_norm = torch.norm(
torch.stack([p.data.norm() for p in self.model.parameters()])
).item()
weight_norms.append(weight_norm)
self.optimizer.step()
epoch_loss += loss.item()
# Update history
self.history['loss'].append(epoch_loss / len(dataloader))
if analyze:
self.history['gradients'].append(np.mean(gradient_norms))
self.history['weights'].append(np.mean(weight_norms))
return epoch_loss / len(dataloader)
def detect_local_minimum(self, window_size=5, threshold=1e-5):
"""Detect if training is stuck in a local minimum"""
if len(self.history['loss']) < window_size:
return False
recent_losses = self.history['loss'][-window_size:]
loss_variance = np.var(recent_losses)
return loss_variance < threshold
def plot_optimization_landscape(self):
"""Visualize optimization progress"""
import matplotlib.pyplot as plt
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 12))
# Plot loss history
ax1.plot(self.history['loss'], label='Training Loss')
ax1.set_xlabel('Epoch')
ax1.set_ylabel('Loss')
ax1.set_title('Training Progress')
ax1.grid(True)
# Plot gradient and weight norms
ax2.plot(self.history['gradients'], label='Gradient Norm')
ax2.plot(self.history['weights'], label='Weight Norm')
ax2.set_xlabel('Epoch')
ax2.set_ylabel('Norm')
ax2.set_title('Gradient and Weight Evolution')
ax2.legend()
ax2.grid(True)
plt.tight_layout()
plt.show()
Results Analysis for Optimization Strategies
class OptimizationComparator:
def __init__(self, model_class, optimizers, dataset):
self.model_class = model_class
self.optimizers = optimizers
self.dataset = dataset
self.results = {}
def run_comparison(self, epochs=100):
for opt_name, opt_config in self.optimizers.items():
model = self.model_class()
optimizer = opt_config['class'](
model.parameters(),
**opt_config['params']
)
analyzer = OptimizationAnalyzer(model, nn.MSELoss(), optimizer)
print(f"\nTraining with {opt_name}")
for epoch in range(epochs):
loss = analyzer.train_epoch(self.dataset)
if epoch % 10 == 0:
print(f"Epoch {epoch}: Loss = {loss:.6f}")
self.results[opt_name] = {
'final_loss': loss,
'history': analyzer.history,
'convergence_epoch': self._find_convergence(analyzer.history['loss'])
}
def _find_convergence(self, losses, threshold=1e-5):
"""Find epoch where training converged"""
for i in range(1, len(losses)):
if abs(losses[i] - losses[i-1]) < threshold:
return i
return len(losses)
def plot_comparison(self):
"""Plot comparison of different optimizers"""
import matplotlib.pyplot as plt
plt.figure(figsize=(12, 6))
for opt_name, result in self.results.items():
plt.plot(result['history']['loss'], label=opt_name)
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Optimization Methods Comparison')
plt.legend()
plt.grid(True)
plt.show()
Advanced Loss Surface Analysis
A deep dive into analyzing loss surface characteristics using eigenvalue decomposition and curvature analysis helps identify problematic regions during training and optimize hyperparameters.
import numpy as np
from scipy import linalg
class LossSurfaceAnalyzer:
def __init__(self, model, criterion):
self.model = model
self.criterion = criterion
def compute_hessian_eigenspectrum(self, data, labels):
params = np.concatenate([p.data.numpy().flatten()
for p in self.model.parameters()])
n_params = len(params)
hessian = np.zeros((n_params, n_params))
def loss_fn(params):
self._set_params(params)
output = self.model(data)
return self.criterion(output, labels).item()
# Compute Hessian using finite differences
epsilon = 1e-6
for i in range(n_params):
for j in range(i, n_params):
params_ij = params.copy()
params_ij[i] += epsilon
params_ij[j] += epsilon
fpp = loss_fn(params_ij)
params_i = params.copy()
params_i[i] += epsilon
fp = loss_fn(params_i)
params_j = params.copy()
params_j[j] += epsilon
fp_ = loss_fn(params_j)
f = loss_fn(params)
hessian[i,j] = (fpp - fp - fp_ + f) / (epsilon * epsilon)
hessian[j,i] = hessian[i,j]
# Compute eigenvalues and eigenvectors
eigenvals, eigenvecs = linalg.eigh(hessian)
return {
'eigenvalues': eigenvals,
'eigenvectors': eigenvecs,
'condition_number': np.abs(eigenvals[-1] / eigenvals[0]),
'positive_curvature_ratio': np.sum(eigenvals > 0) / len(eigenvals),
'negative_curvature_ratio': np.sum(eigenvals < 0) / len(eigenvals)
}
def analyze_loss_surface_geometry(self, point, directions, resolution=50):
"""Analyze loss surface along specific directions"""
alphas = np.linspace(-1, 1, resolution)
surface = np.zeros((len(directions), resolution))
for i, direction in enumerate(directions):
for j, alpha in enumerate(alphas):
params = point + alpha * direction
self._set_params(params)
surface[i,j] = self._compute_loss()
return {
'surface': surface,
'smoothness': np.mean(np.abs(np.diff(surface, axis=1))),
'convexity': np.mean(np.diff(surface, n=2, axis=1) > 0)
}
def _set_params(self, params):
"""Helper to set model parameters"""
offset = 0
for p in self.model.parameters():
numel = p.numel()
p.data = torch.from_numpy(
params[offset:offset + numel].reshape(p.shape)
)
offset += numel
def _compute_loss(self):
"""Helper to compute current loss"""
with torch.no_grad():
return self.criterion(self.model(self.data), self.labels).item()
Additional Resources
- "Visualizing the Loss Landscape of Neural Nets" https://arxiv.org/abs/1712.09913
- "The Mechanics of n-Player Differentiable Games" https://arxiv.org/abs/1802.05642
- "ADAM: A Method for Stochastic Optimization" https://arxiv.org/abs/1412.6980
- "On the Difficulties of Training Deep Neural Networks" https://arxiv.org/abs/1212.0975
- "Random Matrix Theory and the Evolution of Deep Neural Network Eigenvalues" https://arxiv.org/abs/2008.00724
- "Gradient Descent Converges to Minimizers" https://arxiv.org/abs/1602.04915