Abstract
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.
Introduction: The Large-Batch Training Paradox
Modern hardware, such as the latest H100 or MI300X GPUs, possesses massive memory capacities that facilitate "large-batch training" with mini-batches exceeding 10,000 samples. While this maximizes hardware efficiency, it presents a significant challenge for optimization.
First-order optimizers, including SGD and AdamW, rely on the stochastic "noise" inherent in smaller batches to help the model navigate the loss landscape and escape sharp, suboptimal minima. When batch sizes become excessively large, the gradient becomes nearly deterministic. This lack of noise causes models to get stuck or generalize poorly. Second-order methods, specifically Natural Gradient Descent (NGD) via Kronecker-Factored Approximate Curvature (KFAC), attempt to use the landscape's curvature to navigate more effectively. However, at large scales, these methods often become numerically unstable. To maintain stability, practitioners must increase "damping," which unfortunately "washes out" the curvature information, causing the expensive second-order method to perform no better than simple SGD.
Methods: Fisher-Orthogonal Projection (FOP)
The authors introduce Fisher-Orthogonal Projection (FOP) to resolve this dilemma. Rather than simply calculating the average gradient of a massive batch, FOP treats intra-batch variance as a signal rather than noise. The process begins by splitting the large batch into two sub-batches, $g_1$ and $g_2$, to extract both the mean ($g_{avg}$) and the difference ($g_{diff}$).
The core innovation is Fisher-Orthogonalization. FOP projects the difference component so that it is orthogonal to the mean gradient specifically under the Fisher Information metric. This ensures that the added signal provides novel, curvature-sensitive information that is not already captured by the mean. To ensure stability and efficiency, FOP incorporates:
- Layer-wise Adaptation: A coefficient ($\beta$) balances the components, while an adaptive step size ($\eta_\ell^*$) automatically damps updates in high-curvature regions.
- Distributed Implementation: A "Dual-Gradient" reduction strategy allows different GPU groups to compute sub-batches in parallel, minimizing the communication overhead typically associated with second-order methods.
Empirical Evaluation and Key Results
The authors evaluated FOP across diverse computer vision benchmarks, including CIFAR-10 (ResNet-18), ImageNet-100 (T2T-ViT), and ImageNet-1K (ResNet-50). The results demonstrate that FOP is uniquely capable of maintaining performance at extreme scales.
Key Findings:
- Scaling Where Others Fail: On CIFAR-10 with batch sizes of 32,768 and 50,000, FOP was the only method capable of reaching the 91% accuracy threshold. Baselines like SGD, AdamW, and KFAC stalled significantly below the target.
- Massive Speedups: FOP achieved a 1.5x–1.7x speedup over SGD/AdamW at moderate batch sizes. At extreme scales (e.g., ImageNet-1K at BS=8192), it reached target accuracy 7.5x faster than standard methods.
- Superior Generalization: On long-tailed benchmarks (CIFAR-LT), FOP reduced Top-1 error by 2.3%–3.3% compared to strong baselines, proving more robust to imbalanced data due to its geometry-aware updates.
- The Damping Advantage: Mathematically, FOP’s correction terms decay much slower than standard KFAC terms as damping increases. This allows FOP to remain "curvature-aware" even in regimes where KFAC is forced to revert to first-order behavior for stability.
Conclusion and Future Work
Fisher-Orthogonal Projection represents a significant step forward for large-scale optimization, proving that intra-batch variance contains critical geometric information that should be harvested rather than suppressed. While FOP currently requires a higher memory footprint in single-GPU settings and carries the computational cost of Fisher block inversion, these are offset by significantly faster convergence in distributed environments.
Future work includes extending this geometry-aware variance approach to Large Language Models (LLMs), where the massive scale of training could benefit immensely from the improved stability and generalization offered by FOP.