Fast Gauss-Newton for Multiclass Cross-Entropy

Abstract: In multiclass softmax cross-entropy, the full generalized Gauss-Newton (GGN) curvature couples all output logits through the softmax covariance, making curvature-vector products harder to scale as the number of classes grows. We show that the standard multiclass GGN can be decomposed exactly into a true-vs-rest term and a positive semidefinite within-competitor covariance term. Fast Gauss-Newton (FGN) retains the first term and drops the second, yielding a positive semidefinite under-approximation of the multiclass GGN that is exact for binary classification. The derivation uses an exact true-vs-rest scalar-margin representation of softmax cross-entropy: the loss and gradient are unchanged, and the approximation enters only at the curvature level. Exploiting the FGN curvature structure, the damped update can be written as an equivalent whitened row-space system with one row per mini-batch example. We solve this system matrix-free by conjugate gradient using Jacobian-vector and vector-Jacobian products of the scalar margin map. Targeted mechanism experiments and an evaluation on a fixed-feature multiclass head support the predictions from the decomposition: FGN stays closest to the full softmax GGN when competitor mass is concentrated or damping is large, and deviates as the dropped within-competitor covariance grows.

• The paper introduces an efficient curvature approximation, Fast Gauss-Newton (FGN), that simplifies large-class softmax optimization by decomposing the output covariance.
• It employs a matrix-free conjugate-gradient solver that enables computation nearly independent of class count, thereby reducing wall time and enhancing accuracy.
• Empirical results on benchmarks like Cars196 show that FGN matches full Gauss-Newton accuracy while significantly improving computational scalability.

Fast Gauss-Newton for Multiclass Cross-Entropy: An Expert Analysis

Introduction

The paper "Fast Gauss-Newton for Multiclass Cross-Entropy" (2605.06081) introduces Fast Gauss-Newton (FGN), a curvature approximation specifically designed for the multiclass softmax cross-entropy regime. FGN addresses the primary computational bottleneck in deploying exact generalized Gauss-Newton (GGN) methods for multiclass settings—the class-coupled structure in the softmax covariance—by presenting an algebraically motivated decomposition of the GGN and constructing an efficient, matrix-free solution pathway that sidesteps the scaling issues of the standard approach.

Decomposition of the Multiclass GGN

The key technical result is an exact blockwise decomposition of the multiclass GGN matrix: $$H^{\mathrm{GGN}} = H^{\mathrm{FGN}} + R^{\mathrm{comp}}$$ where $H^{\mathrm{FGN}}$ is the rank-one, "true-vs-rest" outer product term, and $R^{\mathrm{comp}}$ is the positive semidefinite residual corresponding to within-competitor covariance in the output-logit space. In the binary scenario, $R^{\mathrm{comp}}$ is identically zero and FGN yields the exact GGN. The construction leverages an exact reparameterization of the cross-entropy loss as a softplus of the logit true-vs-rest margin, allowing the gradient and objective to remain untouched while the curvature is approximated at the level of the output geometry.

Figure 1: Mechanism checks for FGN—(a) update time scaling with class count, (b) output-space trace decomposition with shaded dropped trace $p_\dagger\xi$, (c) fraction of full-GGN quadratic decrease preserved by FGN as competitor dispersion varies.

This decomposition is not heuristic: the output-space factorization is exact, and the parameter-space gap is precisely the competitor-Jacobian-weighted covariance. FGN maintains the contribution from the mean of the competitor logit block while dropping covariance interactions within this block. The approximation is tight when probability mass on negative classes is concentrated and deteriorates as that mass becomes more diffuse within the competitor set.

Efficient Matrix-Free Solver

The retained curvature structure of FGN yields a batched row-space system entirely distinct from the standard multiclass GGN. For a minibatch of size $b$, the FGN curvature is expressed as $H^{\mathrm{FGN}} = \frac{1}{b} J^\top Q J$, where $J$ stacks the per-example true-vs-rest margin Jacobians and $Q$ is a strictly positive diagonal matrix. This structure admits an efficient matrix-free conjugate-gradient (CG) solution, which involves only Jacobian-vector and vector-Jacobian products with the scalar margin map. Notably, FGN bypasses computation over the $C \times C$ logit covariance, yielding wall time essentially independent of the class count for the curvature-product pathway.

Mechanism and Empirical Properties

A detailed mechanistic analysis confirms several theoretical properties:

• Class Scaling: FGN exhibits significantly better scaling with the number of classes than full multiclass GGN curvature solves. Update-times scale favorably (Figure 2a).
• Curvature Gap Visibility: The algebraic gap—the trace of the dropped output-space residual—grows as the probability mass on wrong classes is spread more evenly (Figure 2b).
• Step-Level Robustness: The proportion of the full-GGN quadratic decrease retained by FGN remains near unity when competitor dispersion is low or damping is large, with graceful degradation as dispersion and class entropy increase (Figure 2c).

These properties are confirmed experimentally, with ablations directly inspecting the sensitivity of the quadratic model decrease to the structure of the competitor class block. The trace metric $H^{\mathrm{FGN}}$0, with $H^{\mathrm{FGN}}$1 for the conditional competitor distribution $H^{\mathrm{FGN}}$2, effectively diagnoses when the FGN approximation is nearly exact.

Optimization in Practice: Frozen-Feature Affine Head

In a fixed-feature scenario on Cars196 with a frozen ResNet-50 backbone, FGN is benchmarked against the standard first-order Adam optimizer and a truncated-CG, full-GGN reference ("SGN", for stochastic Gauss-Newton) for linear classifier training. FGN matches and slightly outpaces SGN, with improved time-to-accuracy thresholds and indistinguishable final prediction performance under the same objective—despite never explicitly constructing or applying the full softmax GGN.

Figure 2: On Cars196, test accuracy as a function of head-optimization wall time for Adam, SGN, and FGN. Bands represent $H^{\mathrm{FGN}}$3 std. dev. over 10 runs.

Absolute performance metrics reinforce the qualitative advantages:

• FGN achieves the full-GGN reference accuracy ($H^{\mathrm{FGN}}$4 at 3.65s vs. SGN's $H^{\mathrm{FGN}}$5), but requires less wall time to reach intermediate thresholds (e.g., $H^{\mathrm{FGN}}$6s to $H^{\mathrm{FGN}}$7, versus $H^{\mathrm{FGN}}$8s for SGN).
• The first-order method (Adam) saturates at a lower ceiling ($H^{\mathrm{FGN}}$9 in the same regime).

Theoretical Implications

FGN represents a principled alteration of the output-space curvature for softmax models. It stands distinct from diagonal preconditioners (e.g., AdaHessian, Sophia), layerwise approximations (e.g., K-FAC, Shampoo), and output sampling (e.g., sampled softmax) by modifying the algebraic structure of the output covariance directly. The approach is loss-specific and exact at the level of the objective and the gradient, with approximation confined to the curvature.

The provided analysis is tight: the PSD residual term is explicitly identified and quantified. The practical implication is that the main scalability obstacle in Gauss-Newton methods for large-class softmax is not inherent to the architecture or first-order computation, but is a consequence of the output covariance structure—once this is decoupled, both computational and statistical bottlenecks are alleviated.

Future Directions

Potential future research includes:

• Extension of FGN with low-rank or structured approximations to the within-competitor covariance, balancing the tradeoff between fidelity and computational overhead.
• Adoption of the FGN structure in end-to-end training of deep networks, not just for linear heads but also in scenarios with nontrivial encoder component curvature.
• Integration with adaptive damping, warm-started CG, and preconditioner learning for further efficiency.
• Theoretical study of FGN under nonstationary and non-i.i.d. competitor distributions, and its interaction with label noise and adversarial targets.

Conclusion

FGN is a structured, algebraically grounded method for efficient second-order optimization in multiclass deep learning, directly addressing the core scaling barrier in the softmax output geometry via a true-vs-rest margin decomposition. It retains fidelity to the objective and gradient, offers precise theoretical characterization of its approximation error, and delivers superior empirical scaling and robust convergence in practice. Its distinct separation from sampled output approximations and classical curvature preconditioners positions it as a strong candidate for scalable, second-order methods in large-output-class domains.