Block-Diagonal Hessian-Free Optimization

• Block-diagonal Hessian-free optimization is a second-order method that simplifies curvature matrices by splitting network parameters into independent blocks.
• It leverages local conjugate gradient solves with damping and warm-start techniques, enabling parallel, efficient subproblem resolution with reduced memory usage.
• Empirical studies show it achieves faster convergence and improved generalization on deep architectures compared to full-matrix and first-order methods.

Block-diagonal Hessian-free optimization is a second-order optimization approach for training deep neural networks that leverages a structured block-diagonal approximation to curvature matrices, notably the generalized Gauss–Newton (GGN) matrix or the parameter Hessian. By partitioning parameters into disjoint blocks and performing independent local quadratic solves within each block, this method achieves a significant reduction in computational and memory overhead compared to classical (full-matrix) second-order methods, while improving scalability, parallelism, and convergence properties in large-scale deep learning settings.

1. Theoretical Foundations: Curvature Structures in Neural Optimization

Second-order optimization methods commonly formulate a local quadratic objective around the current parameter vector $w\in \mathbb{R}^N$ for a loss $\ell(f(x;w),y)$: $$q(w+\Delta w) = \ell(w) + \Delta w^\top \nabla\ell(w) + \tfrac12 \Delta w^\top G(w) \Delta w$$ where $G(w)$ is a positive semi-definite curvature matrix, typically chosen as the GGN instead of the exact Hessian to ensure tractability and positive semi-definiteness. The GGN is defined over a curvature mini-batch $S_c$ as: $$G = \frac{1}{|S_c|} \sum_{(x, y) \in S_c} J^\top H_{\ell} J$$ with $J = \partial f / \partial w$ and $H_\ell = \partial^2 \ell / \partial f^2$. In Hessian-free (HF) optimization, $G$ is never represented explicitly; instead, products $Gv$ are computed efficiently via a sequence of forward- and reverse-mode automatic differentiation passes.

Block-diagonal structure enters naturally by partitioning $w$ into $B$ disjoint blocks $w_{(1)},\ldots,w_{(B)}$, leading to a curvature matrix $G$ with $B \times B$ blocks. The block-diagonal approximation $\tilde{G}$ retains only the intra-block curvature $G_{(b,b)}$ and discards cross-block entries, yielding: $$\tilde{G} = \mathrm{blockdiag}(G_{(1,1)}, \ldots, G_{(B,B)})$$ This reduces the global quadratic problem to $B$ independent smaller subproblems, one for each block (Zhang et al., 2017, Dangel et al., 2019).

2. Methodology: Block-Diagonal Hessian-Free Algorithm and Implementation

Block-diagonal HF solves, in each iteration, the block-separable quadratic subproblems: $$\min_{\Delta w_{(b)}}\; g_{(b)}^\top \Delta w_{(b)} + \tfrac12 \Delta w_{(b)}^\top G_{(b)} \Delta w_{(b)}$$ with $g_{(b)}$ the gradient and $G_{(b)}$ the GGN block for parameters in block $b$. These conjugate gradient (CG) inner loops are run in parallel for all blocks. Each CG subproblem includes:

• Damping: $G_{(b)}$ is replaced by $G_{(b)} + dI$ (for $d>0$) to ensure strict positive definiteness.
• Warm-start Momentum: CG initialization at $\rho s_{(b)}$ with $\rho \approx 0.95$.
• Truncation: Each CG solve halts after at most $\mathrm{max\_cg}$ iterations or when the relative residual norm is below a threshold $\epsilon$.

A typical outer iteration involves:

1. Sampling a large gradient mini-batch $S_g$ to compute block gradients.
2. Selecting a smaller curvature mini-batch $S_c$.
3. Solving each block's CG subproblem in parallel.
4. Aggregating the block updates and stepping the full $w$.

Block assignment often aligns with network modularity (e.g., per layer, encoder/decoder split), both to balance workloads and to exploit locality of parameter interactions (Zhang et al., 2017, Dangel et al., 2019).

3. Modular Block-Diagonal Approximations and Automatic Differentiation

A modular extension of backpropagation can propagate block-diagonal curvature matrices—Hessian, GGN, or Positive-Curvature Hessian (PCH)—through each network module. For a single module, the curvature backpropagation equation is: $$H_x = J_x^\top H_z J_x + \sum_k H_{z_k}(x)\,\delta z_k$$ GGN backpropagation omits the second term, valid for modules linear in parameters. This enables assembly of exact block-diagonal curvature matrices via modular, local routines, which is especially effective for feedforward and convolutional networks (Dangel et al., 2019).

Local blocks, particularly for linear layers or channel-wise parameters, can be further structured:

• For weights $W$ of a linear layer $z=Wx$, $H_W = (x x^\top) \otimes H_z$.
• For convolutional layers, im2col-based reshaping renders the block structure tractable.

Kronecker-factored approximations (e.g., KFAC) and batch-averaged curvature variants emerge as special cases within this modular block-diagonal scheme.

4. Computational Complexity, Scalability, and Parallelism

Block-diagonal HF significantly reduces both per-iteration computational and memory requirements versus full-matrix HF. In non-parallel settings, the total cost matches that of standard HF, since $\sum_b k_b \approx k$ (with $k$ the total and $k_b$ the per-block number of CG iterations). However, in parallel, each block can be dispatched independently, allowing for a theoretical $B$-fold speedup and a reduction of per-device memory to $O(N/B)$—$N$ the total number of parameters.

Table: Computational Characteristics

Method	Per-iteration Cost	Memory per Worker	Parallelism Potential
Full-matrix HF	$k$ network passes	$O(N)$	Low (centralized CG)
Block-diag HF	$\sum_b k_b$ network passes	$O(N/B)$	High ($B$-fold possible)

Block-diagonal methods are compatible with distributed and data-parallel settings and are robust to large mini-batch sizes, which is crucial for high-throughput many-core or distributed hardware (Zhang et al., 2017).

5. Practical Implementation, Hyper-parameters, and Adaptations

Effective block-diagonal HF requires careful hyper-parameter selection:

• Large gradient mini-batches ($|S_g|$ in $8$k–$32$k), small curvature mini-batches ($|S_c|$ in $512$–$2$k).
• Damping parameter $d \in \{0.01, 0.1, 1.0\}$.
• CG truncation between $30$–$100$ iterations, tolerance $\epsilon\approx10^{-3}$, warm-start momentum $\rho=0.95$.
• Learning rate $\alpha=0.1$ for HF, $\alpha=0.001$ for Adam (for baseline comparisons).

Partitioning into blocks with roughly equal parameter counts per block facilitates balanced computational loads and well-conditioned subproblems. Modular automatic differentiation and vector-Jacobian products underpin efficient curvature extraction for both block-diagonal GGN and Hessian (Zhang et al., 2017, Dangel et al., 2019).

Channel-wise or groupwise blocks are of special interest, as many normalization and convolutional modules yield exactly diagonal or small block-diagonal Hessians. Modern variants (e.g., SGD with Partial Hessian, or SGD-PH) exploit this fine structure to combine first- and second-order updates, updating channel-wise blocks via a Newton-like step extracted efficiently through Hessian-vector products (Sun et al., 2024).

6. Empirical Results and Comparative Performance

Experiments benchmarked block-diagonal HF on three canonical scenarios:

• Deep Autoencoder (MNIST): Encoder/decoder split (2 blocks). Block-diag HF with large batches reached a target reconstruction error with an order of magnitude fewer updates than Adam and fewer than full-matrix HF, showing robustness to curvature batch size.
• 3-layer LSTM (pooled MNIST): Per-layer block partition. Block-diag HF achieved target test accuracy much faster and with higher final accuracy than Adam or HF, especially in the large-batch regime.
• Simplified ResNet (CIFAR-10): 3 blocks. Block-diag HF attained faster and more stable convergence to better test accuracy than both Adam and full HF across a range of curvature batch sizes.

Empirical results consistently show that block-diagonal HF offers:

• Significantly fewer parameter updates than first-order (Adam) methods to reach similar or better minima.
• Improved generalization compared to full-matrix HF, especially under large-batch training (Zhang et al., 2017).

SGD-PH and related diagonal/block-diagonal methods on standard vision benchmarks (CIFAR-10/100, Mini-ImageNet, ImageNet) report $+0.8$–$2.9$\% accuracy gains over SGD with momentum and outperform established Hessian-free optimizers (AdaHessian, Apollo) in both convergence speed and generalization. The runtime overhead remains modest (about $2\times$ over SGDM)—significantly less than naive Hessian-free Newton with full CG solves (Sun et al., 2024).

7. Applications, Strengths, Limitations, and Recommended Use Cases

Key advantages:

• Fewer total parameter updates owing to second-order curvature adaptation.
• High stability and convergence in large-batch training, making efficient use of modern distributed hardware.
• Natural parallelism: block local CG solves can be performed independently with minimal inter-block communication.
• Improved optimization of very deep or recurrent architectures, where first-order methods often stall or require excessive updates.

Limitations:

• Per-iteration compute is higher than first-order methods (typically $5$–$10\times$ more than a simple gradient step).
• Complexity in implementation due to needs for forward-mode automatic differentiation and correct CG/damping routines.
• Tuning sensitivity: damping and CG truncation hyper-parameters may need adjustment but are found to be robust across problems.

Recommended use cases:

• Deep or recurrent network training scenarios where gradient-based methods plateau.
• Large-scale, distributed settings where few but heavy-weight parameter updates reduce communication costs.
• Regimes prioritizing fewer, more expensive steps to reach better minima, especially with large mini-batch training (Zhang et al., 2017, Sun et al., 2024, Dangel et al., 2019).

Block-diagonal Hessian-free optimization thus stands as a scalable, modular, and efficient second-order technique for neural network training, unifying and extending curvature-based optimization approaches in both feedforward and convolutional settings. Modular block-diagonal curvature extraction routines further facilitate integration into automatic differentiation-based machine learning libraries, with direct relevance to both research and large-scale production training pipelines.