Scalable Kronecker-Factored Fisher Approximation for Neural Network Parameter Sensitivity
Abstract
The Fisher Information Matrix (FIM) provides a principled geometric framework for parameter sensitivity in neural networks, but directly computing and using the full FIM is infeasible in high-dimensional models. As a result, most existing methods rely on diagonal approximations that discard important correlation structure. We introduce Matrix-free Fisher Factorization (MFF), a GPU-tractable algorithm that captures both diagonal and off-diagonal dependencies without materializing the full matrix. For post-training neural network layer compression, we prove that under Matrix-Variate Normal assumptions, MFF yields GFWSVD, a unique closed-form linear layer decomposition that optimally minimizes the expected second-order loss increase. Experiments on controlled numerical benchmarks with large neural networks show that GFWSVD achieves up to 50\% compression while matching or exceeding state-of-the-art diagonal and activation-based baselines across most tasks, and it reliably avoids collapse in dense architectures such as Llama 3. Moreover, when used to initialize existing optimization pipelines (e.g., Dobi-SVD), GFWSVD better preserves accuracy at 40\% parameter reduction in regimes where standard methods substantially degrade. Together, these results position MFF and GFWSVD as foundational algorithmic primitives for scalable, second-order-aware neural network approximation and parameter sensitivity.
Lay Summary
The Hessian matrix measures a neural network's sensitivity to parameter changes and captures how these parameters interact. While valuable for optimization and model compression, the Hessian is computationally intractable; its size scales quadratically with the number of parameters, making it impossible to store or process directly. To bypass this, standard approaches resort to diagonal approximations, completely discarding crucial parameter correlations. We introduce a fast, matrix-free algorithm that factorizes the Hessian without materializing it, fully preserving these vital interactions. Building on this approach, we propose a novel low-rank compression method. This method is provably optimal, generalizes existing frameworks in this scope, and achieves state-of-the-art efficiency and performance.