Welcome to DS.SE ! This is a good question and gives me the opportunity to brush up on my own understanding in this area.

Addressing Second-Order Derivatives and the Hessian in Deep Neural Networks

Theoretical Background and Hessian Challenges

In neural networks with $w$ trainable parameters (weights and biases, collectively denoted by $\theta$ in what follows), directly computing the Hessian, which requires $w^2$ entries, is infeasible as $w$ grows large. For deep networks, $w$ can reach millions, as it scales with the number of neurons and connections across layers. For instance, a fully connected layer with 1,000 neurons connected to another layer of 1,000 neurons has $1,000 \times 1,000 = 1,000,000$ weights alone.

While backpropagation traditionally relies on first-order gradients, second-order information through the Hessian captures interlayer dependencies, offering a more detailed mapping of the loss landscape. Despite the memory demands, approximating the Hessian in a structured way can enhance convergence by making use of this curvature data.

The second-order derivatives are non-zero in neural networks primarily due to the presence of non-linear activation functions. While the linear transformations alone yield zero curvature (second derivatives), non-linear activations (eg., ReLU, sigmoid) following these transformations introduce non-zero curvature in the loss function. This is because the network’s overall function becomes non-linear, leading to a Hessian matrix with non-zero entries due to the chain rule applied through layers.

Kronecker-Factored BFGS Approach for Practical Approximation

The Kronecker-factored block-diagonal BFGS (K-BFGS) method, developed by Goldfarb et al. (2020), approximates the Hessian $H$ (with respect to all parameters $\theta$) using a block-diagonal structure, with each block $H^{(l)}$ corresponding to a network layer. The Hessian for a layer $l$ is approximated by a block $H^{(l)}$, which is represented as a Kronecker product:

$$ H^{(l)} = A^{(l)} \otimes G^{(l)} $$

where $A^{(l)}$ and $G^{(l)}$ are smaller matrices capturing layer-specific curvature. This decomposition retains essential curvature information while significantly reducing memory demands. By focusing on layer-wise blocks, this approach scales efficiently with network depth, making updates feasible even in high-dimensional settings. This structure uses intra-layer dependencies, allowing it to capture second-order interactions across layers without requiring a full Hessian calculation. This layer-wise approximation is particularly useful for large networks.

Damping for Stability

To address high variability and the potential non-convex nature of the Hessian in deep networks, Goldfarb et al. implement a "double damping" technique. This approach stabilises drastic shifts in Hessian eigenvalues, which is important for maintaining training stability when curvature varies widely. Inspired by Powell’s damping, this strategy ensures the Hessian remains positive definite by adjusting the scaling factor $\rho$ as follows:

$$ \rho = \frac{1}{y^T s + \text{damping}} $$ This adjustment helps stabilise the Hessian’s behaviour by limiting large eigenvalues and thus enhances training robustness without requiring frequent recalculations.

Alternative Second-Order Methods

Several alternative second-order methods, including Gauss-Newton and Hessian-Free, offer viable approaches but with trade-offs, as discussed in Tan and Lim (2019):. Gauss-Newton approximates the Hessian using the squared Jacobian, which stabilises updates but can struggle with non-positive-definite matrices in deep networks. Levenberg-Marquardt extends Gauss-Newton by adding a regularisation parameter to maintain definiteness, though this parameter may require careful tuning for optimal performance.

Hessian-Free optimisation bypasses explicit Hessian calculation by using finite differences and conjugate gradient methods to estimate curvature. While efficient in high dimensions, it can be sensitive to gradient noise, which is common in deep networks. Each method has specific strengths and limitations.

Implementation and Practical Benefits

Goldfarb et al demonstrate that K-BFGS can outperform or match first-order methods, such as SGD and Adam, on datasets like MNIST, proving its utility when memory headroom allows for second-order computations. Compared to KFAC, another second-order method that involves more computationally intensive matrix inversions, K-BFGS uses layer-wise updates, making it more efficient with large mini-batches.

Expanded Pseudocode for K-BFGS

Below is an expanded implementation of K-BFGS for a feedforward neural network, where each layer-wise update maintains positive definiteness and controls step size through damping.

import numpy as np

# Initialise Hessian approximation as an identity matrix for each layer
def initialise_hessian_approx(weights: dict) -> dict:
    """Initialises the Hessian approximation as identity matrices for each layer."""
    return {layer: (np.eye(weights[layer].shape[0]), np.eye(weights[layer].shape[1])) for layer in weights}

# BFGS update rule for inverse Hessian approximation
def update_bfgs(A: np.ndarray, G: np.ndarray, s: np.ndarray, y: np.ndarray, damping: float) -> np.ndarray:
    """
    Applies the BFGS formula with damping to update the Hessian approximation.

    Parameters:
        A (np.ndarray): Kronecker factor matrix for layer weights (2D array).
        G (np.ndarray): Kronecker factor matrix for layer gradient updates (2D array).
        s (np.ndarray): Update direction based on factor A (1D array).
        y (np.ndarray): Update direction based on factor G (1D array).
        damping (float): Damping term to stabilise updates.

    Returns:
        np.ndarray: Updated Hessian approximation for the layer (2D array).
        
    Note:
        Includes a small constant (1e-8) in the denominator to avoid division by zero.
    """
    # Compute scaling factor to adjust the step size based on vectors s and y
    rho = 1.0 / (y.T @ s + damping + 1e-8)
    I = np.eye(A.shape[0])
    V = I - rho * np.outer(y, s)
    H_new = V @ A @ V.T + rho * np.outer(s, s)
    return H_new

# Expanded K-BFGS update function with damping and initialisation
def kbfgs_update(weights: dict, grads: dict, hessian_approx: dict, learning_rate: float, damping: float = 1e-4):
    """Performs a K-BFGS update step on network weights using layer-wise Hessian approximations."""
    for layer in weights:
        A, G = hessian_approx[layer]  # Kronecker factors representing layer-wise Hessian approximation
        
        # Compute update directions based on Kronecker factors
        s = A @ grads[layer]
        y = G @ grads[layer]
        
        # Update the inverse Hessian approximation using the BFGS formula
        hessian_approx[layer] = update_bfgs(A, G, s, y, damping)
        
        # Compute the update direction and apply it to the weights
        update_direction = hessian_approx[layer] @ grads[layer]
        weights[layer] -= learning_rate * update_direction

# Example usage
weights = {'layer1': np.random.rand(3, 3), 'layer2': np.random.rand(4, 4)}
grads = {'layer1': np.random.rand(3, 3), 'layer2': np.random.rand(4, 4)}
hessian_approx = initialise_hessian_approx(weights)
learning_rate = 0.01
damping = 1e-4

# Perform a single update step
kbfgs_update(weights, grads, hessian_approx, learning_rate, damping)

Summing up

Quasi-Newton methods like K-BFGS offer a balanced alternative to full second-order approaches by using Kronecker-factored approximations to capture curvature without excessive memory demands. Although somewhat memory-intensive, these methods effectively handle layer-wise curvature and enable faster convergence in deep networks. For scenarios with some memory headroom, K-BFGS is a valuable middle ground between computational efficiency and the enhanced convergence benefits of second-order methods.

$\dagger$ In this context, the Hessian is a square matrix of second-order partial derivatives of the neural network's loss function with respect to all pairs of trainable parameters (weights and biases). Each entry in the Hessian matrix, $H_{ij}$, is defined as:

$$ H_{ij} = \frac{\partial^2 L}{\partial \theta_i \partial \theta_j} $$

where $L$ is the loss function, and $\theta_i$ and $\theta_j$ are individual parameters. This matrix captures how the gradient of one parameter changes with respect to another, effectively mapping the curvature of the loss landscape. Such curvature information helps optimise the training process by adapting update steps according to the local geometry of the loss function.

Refererences:

Goldfarb, D., Ren, Y., & Bahamou, A. (2020). Practical quasi-newton methods for training deep neural networks. Advances in Neural Information Processing Systems, 33, 2386-2396.
https://proceedings.neurips.cc/paper_files/paper/2020/file/192fc044e74dffea144f9ac5dc9f3395-Paper.pdf

Powell, M. J. D. (1978). Algorithms for nonlinear constraints that use Lagrange functions. Mathematical Programming, 14(1), 224-248.

Tan, H. H., & Lim, K. H. (2019, April). Review of second-order optimization techniques in artificial neural networks backpropagation. In IOP conference series: materials science and engineering (Vol. 495, No. 1, p. 012003). IOP Publishing: https://iopscience.iop.org/article/10.1088/1757-899X/495/1/012003/pdf