PyTorch中LBFGS优化Rule110元胞自动机高维度不收敛问题排查

Let’s dive into why your PyTorch LBFGS setup fails to converge for 15x15 grids (while working fine for 10x10) and explore actionable fixes to match the performance of dlib’s BFGS implementation.

Key Convergence Barriers in PyTorch’s Setup #

1. Loss Function Numerical Instability #

Your eval() function returns torch.sqrt(s), which introduces a critical numerical issue: the gradient of sqrt(s) blows up as s approaches 0 (since the derivative is 1/(2*sqrt(s))). For high-dimensional grids, this gradient explosion can derail the LBFGS optimizer, causing it to oscillate or get stuck in local minima instead of converging to zero loss.

2. LBFGS Hyperparameter Misalignment #

PyTorch’s LBFGS defaults to a history_size=10, which is too small for capturing the curvature of the higher-dimensional loss surface. dlib’s BFGS implementation likely uses a larger history window to approximate the Hessian inverse more accurately, which is crucial for navigating complex non-convex landscapes like your Rule 110 polynomial fit.

Additionally, PyTorch’s default line search strategy may be less aggressive than dlib’s, making it harder to escape local minima in larger grids.

3. Suboptimal Initialization #

Starting with a fully zero tensor for 15x15 grids places the optimizer far from the true Rule 110 pattern. While 10x10 is small enough to navigate this gap, the larger parameter space makes it exponentially harder for PyTorch’s LBFGS to find the global minimum without a better starting point.

Fixes to Get PyTorch’s LBFGS Converging #

1. Stabilize the Loss Function #

Remove the torch.sqrt(s) call from eval(). Since sqrt() is a monotonically increasing function, optimizing s (the raw sum of squared errors) will lead to the same optimal solution, but with stable gradients even as s approaches zero:

def eval():
    s = 0.
    for i in range(W - 1):
        for j in range(1, W + 1):
            xx = x[i, (j - 1) % W]
            yy = x[i, j % W]
            zz = x[i, (j + 1) % W]
            r = x[i + 1, j % W]
            s += triangle(xx, yy, zz, r)
    # Enforce first row constraints
    for j in range(W - 1):
        s += x[0, j] ** 2
    s += (1 - x[0, W - 1]) ** 2
    return s  # No sqrt here

2. Tune LBFGS Hyperparameters #

Adjust the optimizer to use a larger history window and a more robust line search:

opt = torch.optim.LBFGS(
    [x],
    lr=0.1,
    history_size=100,  # Larger window for better Hessian approximation
    line_search_fn="strong_wolfe",  # More aggressive line search
    max_iter=20  # Increase per-step line search iterations
)

3. Improve Initialization #

Start with a tensor that already satisfies the first-row constraint and adds small noise to lower rows to nudge the optimizer toward the correct pattern:

# Initialize first row correctly, add tiny noise to other rows
x = torch.DoubleTensor(W, W).zero_()
x[0, W-1] = 1.0  # Enforce first row constraint upfront
x[1:] = torch.randn(W-1, W) * 0.01  # Small random noise to avoid local minima
x.requires_grad_(True)

4. Verify Automatic Differentiation Accuracy #

While unlikely, automatic differentiation could introduce subtle precision errors in your polynomial triangle() function. To rule this out, compare the auto-computed gradients with manually derived gradients for a small subset of parameters. If they don’t match, implement manual gradient calculations for the triangle() term.

5. Mimic dlib’s Convergence Criterion #

dlib uses an extremely strict stopping threshold (1e-87). Add a custom check to stop early once loss falls below a tiny threshold, and print progress to monitor convergence:

for i in range(15500):
    def closure():
        opt.zero_grad()
        s = eval()
        s.backward()
        return s
    opt.step(closure)
    # Early stopping and progress tracking
    current_loss = closure().item()
    if i % 100 == 0:
        print(f"Iteration {i}, Loss: {current_loss:.10f}")
    if current_loss < 1e-10:
        print(f"Converged at iteration {i}")
        break

Why dlib’s BFGS Works Better #

dlib’s find_min_using_approximate_derivatives uses numerical differentiation (instead of automatic differentiation) and a more sophisticated BFGS implementation with aggressive line search and larger history windows. These design choices make it more resilient to the non-convex, high-dimensional loss surface of your 15x15 grid problem.

Content的提问来源于stack exchange，提问作者Stepan Yakovenko