Beyond the Mean: Fisher-Orthogonal Projection for Natural Gradient Descent in Large Batch Training

Modern GPUs are equipped with large amounts of high-bandwidth memory, enabling them to support mini-batch sizes of up to tens of thousands of training samples. However, most existing optimizers struggle to perform effectively at such a large batch size. As batch size increases, gradient noise decreases due to averaging over many samples, limiting the ability of first-order methods to escape sharp or suboptimal minima and reach the global minimum. Meanwhile, second-order methods like the natural gradient with Kronecker-Factored Approximate Curvature (KFAC) often require excessively high damping to remain stable at large batch sizes. This high damping effectively washes out the curvature information that gives these methods their advantage, reducing their performance to that of simple gradient descent. In this paper, we introduce Fisher-Orthogonal Projection (FOP), a novel technique that restores the effectiveness of the second-order method at very large batch sizes, enabling scalable training with improved generalization and faster convergence. FOP constructs a variance-aware update direction by leveraging gradients from two sub-batches, enhancing the average gradient with a component of the gradient difference that is orthogonal to the average under the Fisher-metric.

💡 Research Summary

The paper addresses a pressing challenge in modern deep‑learning: training with extremely large mini‑batches (tens of thousands of samples) on high‑bandwidth GPUs. As batch size grows, stochastic gradient noise vanishes, which hampers first‑order optimizers (SGD, Adam, AdamW) that rely on noise to escape sharp minima. Natural‑gradient methods, particularly Kronecker‑Factored Approximate Curvature (K‑FAC), promise curvature‑aware updates but become numerically unstable at large batches because the Fisher information matrix (FIM) becomes ill‑conditioned. Practitioners therefore apply strong damping, which effectively nullifies the curvature information and reduces K‑FAC to plain gradient descent.

The authors propose Fisher‑Orthogonal Projection (FOP), a simple yet powerful augmentation to natural‑gradient descent that restores the usefulness of second‑order information without expensive sketching or stale statistics. At each iteration the model computes two independent sub‑batch losses L₁ and L₂, yielding gradients g₁ and g₂. The average gradient g_avg = (g₁+g₂)/2 captures the common signal, while the difference g_diff = g₁−g₂ encodes intra‑batch variability. To avoid redundancy, g_diff is orthogonalized with respect to g_avg under the Fisher inner product:
 s_proj = (g_diffᵀ F g_avg) / (g_avgᵀ F g_avg + ε)
 g_⊥diff = g_diff − s_proj·g_avg.
By construction ⟨g_avg, g_⊥diff⟩_F = 0, so the orthogonal component adds only novel curvature‑sensitive information that would be lost by naïve averaging.

The combined update direction is
 g_combined = g_avg + β g_⊥diff,
where β is a layer‑wise mixing coefficient. The paper derives a closed‑form optimal β* by minimizing a second‑order Taylor surrogate of the total loss L₁+L₂, assuming the Hessian ≈ F. The result is β* = D/E with D = g_avgᵀ F⁻¹ g_⊥diff and E = g_⊥diffᵀ F⁻¹ g_⊥diff. This adaptive β automatically shrinks to zero when the orthogonal signal is noisy or unhelpful, reverting FOP to standard K‑FAC, and grows when the variance carries useful descent information.

A further refinement is a layer‑wise adaptive step size η*_ℓ, obtained by locally minimizing a quadratic model of the loss along the natural‑gradient direction:
 η*_ℓ = (g_tot,ℓᵀ F⁻¹_ℓ g_combined,ℓ) / (g_combined,ℓᵀ F⁻¹_ℓ g_combined,ℓ).
The final parameter update for layer ℓ is d_ℓ = η₀ η*_ℓ F⁻¹_ℓ g_combined,ℓ, where η₀ is a global base learning rate.

The authors analyze the KL‑norm of the update, showing that the FOP step decomposes into a base term (∝ λ⁻²) and two correction terms (∝ λ⁻¹) when damping λ is large. Consequently, even under heavy damping the orthogonal correction decays more slowly than the base term, preserving useful variance information. In early training, when ‖g_avg‖ is large and the orthogonal component reflects high‑frequency noise, β can become negative, causing a cross‑term that partially cancels the base KL contribution and effectively reduces the required damping.

From a systems perspective, FOP is implemented with a “dual‑gradient” strategy: two disjoint GPU groups compute g₁ and g₂ in parallel via AllReduce. Each GPU that specializes in a subset of layers stores and updates its local Fisher block, inverting it as needed. After the Fisher inverse is available, the specialist applies the orthogonal projection locally and broadcasts the resulting preconditioned gradient to all workers. This design avoids the massive communication overhead typical of full second‑order methods while still exploiting intra‑batch variance.

Empirical evaluation spans ResNet‑18, Vision Transformers, and other architectures on CIFAR‑10/100 and ImageNet‑like tasks. Across batch sizes from 2¹¹ (2048) to 2¹⁶ (65536), FOP consistently outperforms K‑FAC (1.2–1.3× faster convergence) and first‑order baselines (1.5–1.7× faster). At the extreme batch sizes (≥ 32768), wall‑clock speedups reach up to 7.5×. Moreover, on long‑tailed CIFAR‑LT benchmarks, FOP reduces Top‑1 error by 2.3–3.3 % without any additional tricks, while preserving small‑batch accuracy when scaling down. The method requires no extra hyper‑parameters beyond those already present in K‑FAC, and the authors release a pip‑installable implementation that can be added to existing training scripts with a single line.

In summary, Fisher‑Orthogonal Projection revitalizes natural‑gradient descent for the large‑batch regime by extracting and orthogonalizing intra‑batch gradient variance under the Fisher metric, adaptively mixing it with the mean gradient, and scaling updates per‑layer. This approach mitigates the need for excessive damping, retains curvature information, and delivers substantial speed and generalization gains on modern hardware, making second‑order optimization practical for today’s massive deep‑learning workloads.