SahaBose-KFAC: Making KFAC Stable via Spectral Annealing and Curvature Condensation

Illustrative example: optimizer loss landscape trajectories for SahaBose-KFAC, Vanilla KFAC, AdamW, and SGD.

Abstract

Kronecker-factored Approximate Curvature (KFAC) offers a powerful second-order optimization framework, yet it often suffers from instability and high computational overhead in large-scale deep learning. We present SahaBose-KFAC, a novel approach that decouples stability from efficiency through two core operators: Saha Spectral Annealing and Bose Curvature Condensation.

Saha Spectral Annealing introduces a time-varying spectral transformation \(\tilde{\sigma}_i(t) = (\sigma_i + \epsilon)^{\alpha_t}\) that flattens the curvature spectrum early in training, suppressing spurious spikes that lead to divergence. Bose Curvature Condensation identifies the smallest robust "head" subspace capturing the majority of curvature mass and caps the tail inverse-gain, significantly reducing the cost of inversion without sacrificing performance.

Our methodology provides a theoretical bridge between first-order smoothness and second-order convergence, ensuring stable descent even in highly non-convex landscapes. We evaluate SahaBose-KFAC across diverse benchmarks spanning LLMs, VLMs, and VLAs, demonstrating competitive accuracy with significantly improved stability metrics — e.g. Mean Kappa reduced from 3129.89 (Vanilla KFAC) to 1075.24 (SahaBose-KFAC) on AG_NEWS.

Methodology: SahaBose-KFAC Operators

SahaBose-KFAC does not change the KFAC factorization, the objective, or the outer optimizer. It modifies only the spectral behavior of the factor inverses at refresh time. The layer preconditioner retains the standard Kronecker form \[\mathbf{P}^\text{SB}_{\ell,t} = \hat{A}^{-1}_{\ell,t} \otimes \hat{G}^{-1}_{\ell,t}, \quad \Delta W_\ell = -\hat{G}^{-1}_{\ell,t}(\nabla_{W_\ell}\mathcal{L})\hat{A}^{-1}_{\ell,t},\] but the inverses \(\hat{A}^{-1}\) and \(\hat{G}^{-1}\) are produced by the two operators described below rather than by naïve Tikhonov damping.

Figure 1: SahaBose composition flowchart. Saha reshapes the factor spectrum; Bose caps the inverse tail. Ordering matters: Bose must act on the annealed spectrum, not the raw one.

1 · The Failure Mode: Inverse-Spectrum Tail Amplification

For a KFAC factor \(C \in \{A_\ell, G_\ell\}\) with eigendecomposition \(C = U\,\mathrm{diag}(\sigma)\,U^\top\), standard Tikhonov-damped inversion gives \[(C + \lambda I)^{-1} = U\,\mathrm{diag}\!\left((\sigma_i + \lambda)^{-1}\right)U^\top.\] Each eigendirection is scaled by inverse gain \(g_i(\lambda) = (\sigma_i + \lambda)^{-1}\).

The problem sharpens under noisy factor estimates. Writing \(\hat{\sigma}_i = \sigma_i + \epsilon_i\) and expanding to first order: \[f(\hat{\sigma}_i) - f(\sigma_i) \approx -\frac{\epsilon_i}{(\sigma_i + \lambda)^2}.\] Inverse error scales quadratically with tail gain. This motivates two distinct repairs: (i) inverse gain should depend on mode reliability, not a single global scalar; (ii) full-rank inversion should not be wasted uniformly on long noisy tails. These are exactly the roles of Saha and Bose.

2 · Saha Spectral Annealing Saha

The Saha operator is a time-dependent annealing map applied to factor eigenvalues before inversion. Given \(C = U\,\mathrm{diag}(\sigma)\,U^\top\) with \(\sigma_1 \ge \cdots \ge \sigma_d \ge 0\), define annealed eigenvalues \[\tilde{\sigma}_i(t) = (\sigma_i + \varepsilon)^{\alpha_t}, \quad 0 < \alpha_t \le 1,\] with \(\alpha_t\) increasing toward 1 over the course of training.

Condition-number compression. The key property follows immediately: \[\kappa(\tilde{C}_t) = \frac{\tilde{\sigma}_1(t)}{\tilde{\sigma}_d(t)} \;\le\; \left(\frac{\sigma_1 + \varepsilon}{\sigma_d + \varepsilon}\right)^{\alpha_t}.\] When \(\alpha_t < 1\), Saha compresses the condition number — sometimes dramatically. A raw spectrum with \(\kappa \approx 4.3 \times 10^4\) can be reduced to \(\kappa \approx 25\) with \(\alpha_t = 0.30\). As training proceeds and factor estimates stabilize, \(\alpha_t \to 1\) recovers the native KFAC geometry.

3 · Bose Curvature Condensation Bose

Even after early instability subsides, KFAC factors remain structurally heavy-tailed: a small head captures most curvature mass while a long tail contains weak or noisy modes.

Condensation rank. Define normalized spectral mass and its cumulative sum: \[p_i = \frac{\tilde{\sigma}_i(t)}{\sum_j \tilde{\sigma}_j(t)}, \qquad C_k = \sum_{i=1}^k p_i.\] Given a target mass fraction \(m \in (0,1]\), the condensation rank is \[k^* = \min\left\{ k : C_k \ge m \right\}.\]

Head-and-tail inverse weights. \[w_i(t) = \begin{cases}\bigl(\tilde{\sigma}_i(t) + \lambda_t\bigr)^{-1} & i \le k^*, \\[4pt]\lambda_\text{tail}^{-1} & i > k^*,\end{cases}\] yielding \(\hat{C}^{-1}_t = U\,\mathrm{diag}(w(t))\,U^\top\) with the hard guarantee \(w_i(t) \le \lambda_\text{tail}^{-1}\) for all tail modes.

4 · Full SahaBose-KFAC: Composition, Ordering, and Algorithm

For each refreshed factor \(C \in \{A_{\ell,t},\, G_{\ell,t}\}\), the full method applies the ordered composition \[C \xrightarrow{\;\text{Saha}\;} \tilde{C}_t \xrightarrow{\;\text{Bose}\;} \hat{C}^{-1}_t.\]

Results

Table 1: LLM results and ablations

We summarize text classification, natural-language inference, and reading comprehension under a matched optimizer protocol. The upper block compares first-order, matrix-adaptive, and KFAC-family baselines; the lower block isolates the proposed Saha and Bose operators. Full SahaBose-KFAC gives the strongest stability-performance tradeoff: it reduces ill-conditioning, lowers condensation rank, and cuts divergence while staying task-competitive.

Notes. N/A marks diagnostics not defined for non-curvature baselines. Curvature diagnostics are layer/time averages. Div. is the fraction of runs with spikes, numerical failure, or early termination.

Table 2: VLM results and ablations

We summarize image captioning, text-image retrieval, and visual question answering under a matched optimizer protocol. The upper block compares first-order, matrix-adaptive, and KFAC-family baselines; the lower block isolates the proposed Saha and Bose operators. Full SahaBose-KFAC preserves competitive task performance while reducing unstable curvature use through lower condition number, smaller condensation rank, and fewer divergent runs.

Notes. N/A marks diagnostics not defined for non-curvature baselines. Captioning: Flickr30k BLEU-4/BLEU-1; retrieval: T2I R@1/R@10; VQA: accuracy/eval loss. Curvature diagnostics are layer/time averages. Div. is the fraction of runs with spikes, numerical failure, or early termination.